require "numru/ggraph"

module NumRu
  class GPhys
    def deriv(dim)
      dim = dim_index(dim)    # input dim can be a String or an Integer
      x = axis(dim).to_gphys
      a = [true]*dim + [1..-1,false]
      b = [true]*dim + [0..-2,false]
      dydx = ( self[*a] - self[*b] ) / ( x[1..-1] - x[0..-2] )
      xi = ( x[1..-1] + x[0..-2] ) / 2
      dydx.axis(dim).set_pos(xi.coord(0))
      dydx.long_name = "d #{name}/d #{x.name}"
      dydx
    end
  end
end


require "date"
include NumRu

filedir = "time_000000000-000080000"
filehead = "BS1998_"
cuttime = 200..80000
zcut = 'XYZmean'
kappa = 155.0
comax =5.0e-7
comin = -5.0e-7
ins = 2

dens_file = 'restart_long-DensBZ.nc'
exner_file = 'restart_long-ExnerBZ.nc'
ptemp_file = 'restart_long-PTempBZ.nc'

ptemp0_file = "../#{filedir}/#{filehead}PTemp_rank000000.nc"
ptemp1_file = "../#{filedir}/#{filehead}PTemp_rank000001.nc"
ptemp2_file = "../#{filedir}/#{filehead}PTemp_rank000002.nc"
ptemp3_file = "../#{filedir}/#{filehead}PTemp_rank000003.nc"
ptemp4_file = "../#{filedir}/#{filehead}PTemp_rank000004.nc"
ptemp5_file = "../#{filedir}/#{filehead}PTemp_rank000005.nc"

exner0_file = "../#{filedir}/#{filehead}Exner_rank000000.nc"
exner1_file = "../#{filedir}/#{filehead}Exner_rank000001.nc"
exner2_file = "../#{filedir}/#{filehead}Exner_rank000002.nc"
exner3_file = "../#{filedir}/#{filehead}Exner_rank000003.nc"
exner4_file = "../#{filedir}/#{filehead}Exner_rank000004.nc"
exner5_file = "../#{filedir}/#{filehead}Exner_rank000005.nc"

velx0_file = "../#{filedir}/#{filehead}VelX_rank000000.nc"
velx1_file = "../#{filedir}/#{filehead}VelX_rank000001.nc"
velx2_file = "../#{filedir}/#{filehead}VelX_rank000002.nc"
velx3_file = "../#{filedir}/#{filehead}VelX_rank000003.nc"
velx4_file = "../#{filedir}/#{filehead}VelX_rank000004.nc"
velx5_file = "../#{filedir}/#{filehead}VelX_rank000005.nc"

velz0_file = "../#{filedir}/#{filehead}VelZ_rank000000.nc"
velz1_file = "../#{filedir}/#{filehead}VelZ_rank000001.nc"
velz2_file = "../#{filedir}/#{filehead}VelZ_rank000002.nc"
velz3_file = "../#{filedir}/#{filehead}VelZ_rank000003.nc"
velz4_file = "../#{filedir}/#{filehead}VelZ_rank000004.nc"
velz5_file = "../#{filedir}/#{filehead}VelZ_rank000005.nc"


ptemp  = GPhys::IO.open([ptemp0_file, \
                         ptemp2_file, \
                         ptemp1_file, \
                         ptemp3_file, \
                         ptemp4_file, \
                         ptemp5_file ], \
                        'PTemp')

exner  = GPhys::IO.open([exner0_file, \
                         exner2_file, \
                         exner1_file, \
                         exner3_file, \
                         exner4_file, \
                         exner5_file ], \
                        'Exner')

velx   = GPhys::IO.open([velx0_file, \
                         velx2_file, \
                         velx1_file, \
                         velx3_file, \
                         velx4_file, \
                         velx5_file ], \
                        'VelX')
velz   = GPhys::IO.open([velz0_file, \
                         velz2_file, \
                         velz1_file, \
                         velz3_file, \
                         velz4_file, \
                         velz5_file ], \
                        'VelZ')


densBZ  = GPhys::IO.open(dens_file, 'DensBZ')
exnerBZ  = GPhys::IO.open(exner_file, 'ExnerBZ')
ptempBZ  = GPhys::IO.open(ptemp_file, 'PTempBZ')



# dt rho theta
  dens_prime = densBZ * ( exner / exnerBZ - ptemp /ptempBZ )
  rho_theta_prime = dens_prime * ptemp                      
  dt_rho_theta_prime_XYZmean = rho_theta_prime.mean('x').mean('y').mean('z').deriv('t')

# rho theta grad v
  dx_theta = ptemp.deriv('x')
  dz_theta = ptemp.deriv('z')

#  grad =  ( velx[1..1199,0..-1,1..99,0..-1] * dx_theta[0..-1,0..-1,1..99,0..-1] + velz[1..1199,0..-1,1..99,0..-1] * dz_theta[1..1199,0..-1,0..-1,0..-1] )
  u_theta = velx[1..1199,0..-1,1..99,0..-1] * dx_theta[0..-1,0..-1,1..99,0..-1]
  w_theta = velz[1..1199,0..-1,1..99,0..-1] * dz_theta[1..1199,0..-1,0..-1,0..-1]
  grad_prime = dens_prime[1..1199,0..-1,1..99,0..-1] * ( u_theta + w_theta )

  grad_prime_XYZmean = grad_prime.mean('x').mean('y').mean('z')

DCL.gropn(ins)
GGraph.line( dt_rho_theta_prime_XYZmean.cut('t'=>cuttime), 
             true, 'exchange'=>false ,
             'index'=>12, 'type'=>1, 'label'=>'dens',
             'title'=>"z=#{zcut}",
             'max'=>comax, 'min'=>comin)    
GGraph.line( grad_prime_XYZmean.cut('t'=>cuttime), 
             false, 'exchange'=>false ,
             'index'=>22, 'type'=>1, 'label'=>'grad',
             'max'=>comax, 'min'=>comin)    

DCL.grcls