require "numru/gphys"
require "numru/dcl"
include NumRu

ARGV.length<3 && raise("Usage: ruby #$0 Omega G U")

omega = ARGV.shift
g = ARGV.shift
u0 = ARGV.shift

fname1 = "../data/spshallow-zd_rn4cn_freedecay_Omega#{omega}_G#{g}_U#{u0}.nc"
file1 = NetCDF.open(fname1)
fname2 = "../data/spshallow-zd_rn4cn_freedecay_Omega#{omega}_G#{g}_U#{u0}-2.nc"
if File.exist?(fname2)
  file2 = NetCDF.open(fname2)
  file = [file1,file2]
else
  file = file1
end


gu = GPhys::IO.open(file,"vellon")
gv = GPhys::IO.open(file,"vellat")
gvor = GPhys::IO.open(file,"vor")
ghsfc = GPhys::IO.open(file,"hsfc")
nlon,nlat,ntim = gu.shape
lat = gu.coord("lat").val
tim = gu.coord("t").val
nm = file1.dim("n").length
lonw = file1.var("lon_weight").get
latw = file1.var("coslat_lat_weight").get
omega = file1.att("Omega").get[0]
g = file1.att("Grav").get[0]
h = file1.att("Hbar").get[0]
a = file1.att("Radius").get[0]


f = NMath::sin(lat)*2*omega
f.reshape!(1,nlat)

lonw /= lonw.sum
latw /= latw.sum
kenergy = NArray.sfloat(ntim)
penergy = NArray.sfloat(ntim)
enstrophy = NArray.sfloat(ntim)
for n in 0...ntim
  u = gu[true,true,n].val
  v = gv[true,true,n].val
  hsfc = ghsfc[true,true,n].val
  vor = gvor[true,true,n].val
  kenergy[n] = 0.5*((u**2+v**2)*(hsfc+h)).mul_add(lonw,0).mul_add(latw,0)
  penergy[n] = 0.5*g*((hsfc+h)**2).mul_add(lonw,0).mul_add(latw,0)
  enstrophy[n] = ((vor+f)**2*(hsfc+h)).mul_add(lonw,0).mul_add(latw,0)
end
penergy -= 0.5*g*h**2
kenergy -= 0.5

iws = ARGV.delete("-ps") ? 2 : 4
if ARGV.delete("-png")
  DCL::swlset("lwnd",false)
  DCL::swlset("ldump",true)
end
DCL::sglset("lfull",true)
DCL::gropn(iws)


2.times do |i|
  DCL::grfrm
  DCL::grsvpt(0.09,0.98,0.08,0.67)
  DCL::grswnd(0,tim[-1],-0.2,0.2)
  DCL::grstrn(1)
  DCL::grstrf
  DCL::uzfact(0.7)

  DCL::uxaxdv("b",0.2,1)
  DCL::uxaxdv("t",0.2,1)
  DCL::uyaxdv("l",0.1,0.2)
  DCL::uyaxdv("r",0.1,0.2)
  DCL::uxsttl("b","time [s]",0)

  case i
  when 0
    kei = 0.5
    pei = 0.5*g*h**2
    DCL::sgplzu(tim,kenergy+penergy,1,2)
    DCL::sgplzu(tim,kenergy,3,2)
    DCL::sgplzu(tim,penergy,2,2)
    x0 = 0.7
    x1 = 0.78
    x2 = 0.80
    y0 = 0.65
    DCL::sgplzr([x0,x1],[y0]*2,1,2)
    DCL::sgtxzr(x2,y0,"total (#{kei+pei} t=0)",0.015,0,-1,1)
    y0 = 0.63
    DCL::sgplzr([x0,x1],[y0]*2,3,2)
    DCL::sgtxzr(x2,y0,"kinetic (#{kei} t=0)",0.015,0,-1,1)
    y0 = 0.61
    DCL::sgplzr([x0,x1],[y0]*2,2,2)
    DCL::sgtxzr(x2,y0,"potential (#{pei} t=0)",0.015,0,-1,1)
    title = "energy"
  when 1
    DCL::sgplzu(tim,enstrophy,1,2)
    title = "enstrophy"
  end
  DCL::sgtxzr(0.02,0.70,"Total #{title} change",0.02,0,-1,2)
  DCL::sgtxzr(0.98,0.035,"Omega=#{omega.to_i}, g=#{g.to_i}, H=#{h.to_i}, U=1, a=#{a.to_i}",0.012,0,1,1)
end

DCL::grcls