#!/usr/bin/env ruby

require "numru/gphys"
include NumRu

#< Define constants >

CpDry    = 820.0
GasRDry  = 287
Grav     = 9.806

#< Define filenames >

tempfilename     = "venus_Temp.nc"

bvfreqsqfilename = "venus_BVFreqSq.nc"

#< Opening GPhys variables >

temp   = GPhys::IO.open(tempfilename,'Temp')
sig    = GPhys::IO.open(tempfilename,'sig')
sigw   = GPhys::IO.open(tempfilename,'sig_weight')
sigm   = GPhys::IO.open(tempfilename,'sigm')

#< Creating working variables >

imax,jmax,kmax,tmax = temp.shape

xyr_Temp     = NArray.float(imax,jmax,kmax+1,tmax)
xyr_Height   = NArray.float(imax,jmax,kmax+1,tmax)
xyz_BVFreqSq = NArray.float(imax,jmax,kmax,tmax)

xyz_Temp = temp.val

z_Sig    = sig.val
r_Sig    = sigm.val
z_DelSig = sigw.val

#< Interpolate temperature at sigm levels >

xyr_Temp[true,true,0,true]    = xyz_Temp[true,true,0,true]
xyr_Temp[true,true,kmax,true] = xyz_Temp[true,true,kmax-1,true]

for k in 1..kmax-1
  x = (log(r_Sig[k])-log(z_Sig[k-1]))/(log(z_Sig[k])-log(z_Sig[k-1]))
  xyr_Temp[true,true,k,true] \
  = x * xyz_Temp[true,true,k,true] + (1-x) * xyz_Temp[true,true,k-1,true]
end

#< Calculate height field at sigm levels >

xyr_Height[true,true,0,true] = 0.0

for k in 0..kmax-1
  xyr_Height[true,true,k+1,true] = xyr_Height[true,true,k,true] \
     + GasRDry / Grav * xyz_Temp[true,true,k,true]  \
             * z_DelSig[k] / z_Sig[k]
end


#< Calculate Brunt Vaisala frequency squared  >

for k in 0..kmax-1
  for j in 0..jmax-1
    for i in 0..imax-1
      xyz_BVFreqSq[i,j,k,true] \
      = Grav/xyz_Temp[i,j,k,true]*(Grav/CpDry + (xyr_Temp[i,j,k+1,true]-xyr_Temp[i,j,k,true])/(xyr_Height[i,j,k+1,true]-xyr_Height[i,j,k,true]))
    end
  end
end

bvfreqsq = temp.copy
bvfreqsq.replace_val(xyz_BVFreqSq)
bvfreqsq.rename("BVFreqSq")
bvfreqsq.set_att("long_name","Brunt Vaisala frequency squared")
bvfreqsq.set_att("units","s-2")

#< Output to file >

bvfreqsqfile=NetCDF.create(bvfreqsqfilename)
GPhys::IO.write( bvfreqsqfile, bvfreqsq )
bvfreqsqfile.close