# coding: utf-8
#!/usr/bin/ruby
# coding: utf-8


require "numru/ggraph"
include NumRu
include NMath

## NCEP / NCAR reanalysis data, which can be downloaded from
## ftp://ftp.cdc.noaa.gov/Datasets/ncep.reanalysis.derived/pressure/vwnd.mon.mean.nc

u  = GPhys::IO.open("u.nc", "u")[true,true,true,900..1500]
v  = GPhys::IO.open("v.nc", "v")[true,true,true,900..1500]
rho = GPhys::IO.open("rho.nc", "rho")[true,true,true,900..1500]
lat = GPhys::IO.open("u.nc", "lat")[true]
#[lon, lat, lev,time]

lat_1  = lat.copy

#stream function
u_var  = u.mean('time').mean('lon')
v_var  = (v.mean('time')).mean('lon')
#w_var  = (w.mean('time')).mean('lon')
rho_var  = (rho.mean('time')).mean('lon')

lat    = [-89.0, -87.0, -85.0, -83.0, -81.0, -79.0, -77.0, -75.0, -73.0, -71.0, -69.0, -67.0, -65.0, -63.0, -61.0, -59.0, -57.0, -55.0, -53.0, -51.0, -49.0, -47.0, -45.0, -43.0, -41.0, -39.0, -37.0, -35.0, -33.0, -31.0, -29.0, -27.0, -25.0, -23.0, -21.0, -19.0, -17.0, -15.0, -13.0, -11.0, -9.0, -7.0, -5.0, -3.0, -1.0, 1.0, 3.0, 5.0, 7.0, 9.0, 11.0, 13.0, 15.0, 17.0, 19.0, 21.0, 23.0, 25.0, 27.0, 29.0, 31.0, 33.0, 35.0, 37.0, 39.0, 41.0, 43.0, 45.0, 47.0, 49.0, 51.0, 53.0, 55.0, 57.0, 59.0, 61.0, 63.0, 65.0, 67.0, 69.0, 71.0, 73.0, 75.0, 77.0, 79.0, 81.0, 83.0, 85.0, 87.0, 89.0]
for i in 0..89
  lat_1[i] =  cos(lat[i]* Math::PI  / 180.0)
end

delta_z = 600
psi_0   = 0 
psi     = v_var.copy

for i in 0..39
  if i == 0 then
    psi[true,i] = 0
  else
    psi[true,i]  = psi_0 - v_var[true,i-1]*rho_var[true,i-1]*delta_z/lat_1[true]
    psi_0 = psi[true,i]
  end
end

p psi.shape
    
DCL.swpset("ifl",1)
DCL.gropn(2)
#DCL.uzfact(0.6) 
DCL.sgpset('lfull',true)
GGraph.set_fig('itr'=>1,'viewport' => [0.1,0.85,0.2,0.55])

GGraph.tone(psi[true,true], true, 'max'=>6000, 'min'=>-6000 ,'nlev'=>30)
#GGraph.tone( psi[true,true], true,'max'=>4000,'min'=>-4000 )
GGraph.color_bar
GGraph.contour( psi[true,true], false, 'max'=>6000, 'min'=>-6000 ,'nlev'=>30)


#GGraph.tone( u_var[true,true], true,'max'=>36,'min'=>-36 )
#GGraph.color_bar("landscape"=>true)

DCL.grcls