#!/usr/bin/env ruby

require "numru/ggraph"      # ggraph library
require "numru/dcl"
include NumRu

######################################################################

# 2nd order centered derivation
class GPhys
def dz(gp)  

   dzvar = NArray.sfloat(gp.data.val.shape[0], gp.data.val.shape[1]).fill!(0.0)
   dzvar[0..-1, 1..-2] = ( gp.data.val[0..-1, 0..-3] - gp.data.val[0..-1, 2..-1] ) / ( 2.0 * 100.0 ) 

   dzvar = VArray.new(dzvar,'long_name'=>'temparature gradient', 'units'=>'K/m')

   dzgp = GPhys.new( Grid.new(gp.axis(0), gp.axis(1)) , dzvar)

   return dzgp

end

def dx(gp)

   dxvar = NArray.sfloat(gp.data.val.shape[0], gp.data.val.shape[1]).fill!(0.0)
   dxvar[1..-2, 0..-1] = ( gp.data.val[0..-3, 0..-1] - gp.data.val[2..-1, 0..-1] ) / ( 2.0 * 100.0 ) 

   dxvar = VArray.new(dxvar,'long_name'=>'temparature gradient', 'units'=>'K/m')

   dxgp = GPhys.new( Grid.new(gp.axis(0), gp.axis(1)) , dxvar)

   return dxgp

end
end

def dz(var)  

   dzvar = NArray.sfloat(var.shape[0], var.shape[1]).fill!(0.0)
   dzvar[0..-1, 1..-2] = ( var[0..-1, 0..-3] - var[0..-1, 2..-1] ) / ( 2.0 * 200.0 ) 
   return dzvar


end

def dx(var)

   dxvar = NArray.sfloat(var.shape[0], var.shape[1]).fill!(0.0)
   dxvar[1..-2, 0..-1] = ( var[0..-3, 0..-1] - var[2..-1, 0..-1] ) / ( 2.0 * 200.0 ) 
   return dxvar

end

######################################################################

# file

file_d = "/work01/kitamo/2005-08-02_kitamo/nc/arare-mars-20050801.nc"  # output netCDF file (disturb field value)
file_b = "/work01/kitamo/2005-08-02_kitamo/nc/arare-mars-20050801-restart.nc"  # output netCDF file (basic field value)

# definition of constant value

Bulk = 0.01  # bulk coefficient
Delz = 200.0   # vartical grid interval [m]
Sfctemp = 270.0 # ground tempareture [K]
Tstep = 300.0    # time step 
Tbegin = 39600.0 
Tend = 43200.0

# GPhys object initialization

gp_velx    = GPhys::IO.open(file_d, "VelX")
gp_velz    = GPhys::IO.open(file_d, "VelZ")
gp_pottemp_d = GPhys::IO.open(file_d, "PotTemp")
gp_pottemp_b = GPhys::IO.open(file_b, "PotTempBasicZ")
gp_exner_d = GPhys::IO.open(file_d, "Exner")
gp_exner_b = GPhys::IO.open(file_b, "ExnerBasicZ")
gp_kh = GPhys::IO.open(file_d, "Kh")

ax_x = gp_velz.axis(0)
ax_z = gp_velz.axis(1)
ax_t = gp_velz.axis(2)

# uniform array size

gp_velz = gp_velz.cut('x'=>0..25600, 'z'=>0..10000)
gp_pottemp_d = gp_pottemp_d.cut('x'=>0..25600, 'z'=>0..10000)
gp_pottemp_b = gp_pottemp_b.cut('x'=>0..25600, 'z'=>0..10000)
gp_exner_d = gp_exner_d.cut('x'=>0..25600, 'z'=>0..10000)
gp_exner_b = gp_exner_b.cut('x'=>0..25600, 'z'=>0..10000)
gp_kh = gp_kh.cut('x'=>0..25600, 'z'=>0..10000)

# get fundamental value

xnum = gp_velz.coord(0).val.shape[0]
znum = gp_velz.coord(1).val.shape[0]

# initialize GPhys object 

data_adv = NArray.sfloat(xnum, znum).fill!(0.0)
data_adv = VArray.new(data_adv, 'long_name'=>'advection', 'units'=>'K/day')
gp_adv = GPhys.new( Grid.new(ax_x, ax_z), data_adv)

data_rad = NArray.sfloat(xnum, znum).fill!(0.0)
data_rad = VArray.new(data_rad, 'long_name'=>'cooling rate', 'units'=>'K/day')
gp_rad = GPhys.new( Grid.new(ax_x, ax_z), data_rad)

data_diff = NArray.sfloat(xnum, znum).fill!(0.0)
data_diff = VArray.new(data_diff, 'long_name'=>'diffusion', 'units'=>'K/day')
gp_diff = GPhys.new( Grid.new(ax_x, ax_z), data_diff)


# calculate tendency

time = Tbegin
loopcount = 0

while time <= Tend do

  gp_pottemp = gp_pottemp_b + gp_pottemp_d.cut('t'=>time)
  gp_exner = gp_exner_b + gp_exner_d.cut('t'=>time)

  # derivation

  dptdx = dx(gp_pottemp_d.cut('t'=>time).data.val)
  dptdz = dz(gp_pottemp_d.cut('t'=>time).data.val)
  dptbdz = dz(gp_pottemp_b.data.val)

  # calculate advection

  gp_adv.data.val = gp_adv.data.val - ( gp_velz.cut('t'=>time).data.val * (dptdz + dptbdz) + gp_velx.cut('t'=>time).data.val * dptdx ) * 86400.0


  # calculate diffusion

  gp_diff.data.val = gp_diff.data.val + dx(gp_kh.cut('t'=>time).data.val * dx(gp_pottemp.data.val) ) + dz(gp_kh.cut('t'=>time).data.val * dz(gp_pottemp.data.val) ) * 86400.0

  # calculate radiation cooling

  gp_rad.cut('z'=>0..5000).data.val = gp_rad.cut('z'=>0..5000).data.val - 50.0 / gp_exner.cut('z'=>0..5000).data.val 

  # handling for next loop

  time = time + Tstep
  loopcount = loopcount + 1

end

gp_adv.data.val = gp_adv.data.val / loopcount
gp_rad.data.val = gp_rad.data.val / loopcount
gp_diff.data.val = gp_diff.data.val / loopcount

gp_total = gp_adv + gp_rad + gp_diff

# calculate radiation

gp_rad = gp_rad.mean('x')
gp_adv = gp_adv.mean('x')
gp_diff = gp_diff.mean('x')
gp_total = gp_total.mean('x')

# draw
#
DCL.gropn(1)
DCL.sgpset('lcntl', false)
DCL.uzfact(0.7)
GGraph.line( gp_total, true, 'exchange'=>true, 'title'=>'THETA TENDENCY' )
GGraph.line( gp_adv, false, 'exchange'=>true )
GGraph.line( gp_rad, false, 'exchange'=>true )
GGraph.line( gp_diff, false, 'exchange'=>true )
DCL.grcls