#!/usr/bin/env ruby

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

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

def dx(var, del)

   dxvar = NArray.sfloat(var.shape[0], var.shape[1]).fill!(0.0)

   # 2nd order
#   dxvar[1..-2, 0..-1] = ( var[2..-1, 0..-1] - var[0..-3, 0..-1] ) / ( 2.0 * del ) 

   # 4th order
   dxvar[2..-3, 0..-1] = 
      ( 8 * ( var[3..-2, 0..-1] - var[1..-4, 0..-1] ) - ( var[4..-1, 0..-1] - var[0..-5, 0..-1] ) ) / ( 12.0 * del )  

   return dxvar

end

def dz(var, del)  

   dzvar = NArray.sfloat(var.shape[0], var.shape[1]).fill!(0.0)

#  2nd order
#   dzvar[0..-1, 1..-2] = ( var[0..-1, 2..-1] - var[0..-1, 0..-3] ) / ( 2.0 * del ) 

#  4th order
   dzvar[0..-1, 2..-3] = 
     ( 8 * ( var[0..-1, 3..-2] - var[0..-1, 1..-4] ) - ( var[0..-1, 4..-1] - var[0..-1, 0..-5] ) ) / ( 12.0 * del )  

   return dzvar

end

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

# file

file_d = "./arare-20050903-1.nc"  # output netCDF file (disturb field value)
file_b = "./arare-odaka1998-dx200-init.nc"  # output netCDF file (basic field value)

# definition of constant value

Bulk = 0.01        # bulk coefficient
Delz = 100.0       # vartical grid interval [m]
Sfctemp = 270.0    # ground tempareture [K]
Tstep = 300.0      # time step [s]
Tbegin = 39600.0   # begin time [s]
Tend = 43200.0     # end time [s]
Dense = 0.016      # density [kg /m^3]

# 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_pottemp = 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, Delz)
  dptdz = dz(gp_pottemp_d.cut('t'=>time).data.val, Delz)
  dptbdz = dz(gp_pottemp_b.data.val, Delz)

  # 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 )

  # calculate diffusion

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


  #gp_diff.cut('z'=>Delz/2).data.val = gp_diff.cut('z'=>Delz).data.val + Bulk * Dense * gp_velx.cut('t'=>time, 'z'=>Delz/2).data.val.abs * ( 25.0 - gp_pottemp_d.cut('t'=>time, 'z'=>Delz/2).data.val )

  # 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

# calculate radiation

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

gp_diff.data.val = gp_diff.data.val * 86400
gp_adv.data.val = gp_adv.data.val * 86400

gp_total = gp_adv + gp_rad + gp_diff

# draw
#
DCL.gropn(4)
DCL.sgpset('lcntl', false)
DCL.uzfact(0.7)
GGraph.line( gp_total, true, 'exchange'=>true, 'title'=>'THETA TENDENCY',
             'max'=>150, 'min'=>-150, 'index'=>12 )
GGraph.line( gp_adv, false, 'exchange'=>true, 'index'=>22 )
GGraph.line( gp_rad, false, 'exchange'=>true, 'index'=>32 )
GGraph.line( gp_diff, false, 'exchange'=>true, 'index'=>42 )
DCL.grcls