load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_code.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/gsn_csm.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/contributed.ncl"
load "$NCARG_ROOT/lib/ncarg/nclscripts/csm/shea_util.ncl"

begin
 filename = "potvorticity"
 dir_in  = "/home/netapp-clima-scratch/ggraffin/MOM_"
 dir_out = "/home/netapp-clima-users1/ggraffin/archive/script/"
 NAME = "CONTROL"
 name = str_lower(NAME)
 
;; Remove the file if it already exists
 system("rm "+dir_out+"potentialvorticity_"+name+"_new.nc")
  
;; Open files
 fi = addfile(dir_in+NAME+"/"+name+".ocean_month.nc","r")
 fm = addfile(dir_out+"regionmask_v6.nc","r")
 fg = addfile(dir_out+"dummygrid_one.nc","r")
 fc = addfile(dir_out+"coriolis_param.nc","r")
 fo = addfile(dir_out+"potentialvorticity_"+name+"_new.nc","c")
 
;; Read data
 pot_rho = fi->pot_rho_0
 f_cor = fc->F
 lev = fg->z_depth
 zlv = fi->st_ocean
 zrh = fi->neutral
 lat = fi->geolat_t(:,0)
 lon = fi->geolon_t(0,:)
 nzlv = dimsizes(zlv)
 nzrh = dimsizes(zrh)
 nlat = dimsizes(lat)
 nlon = dimsizes(lon)
 delete([/fi,fm,fg,fc/]) 

; printVarSummary(pot_rho)
; printVarSummary(f_cor) 
 
;; Create time average
 pot_rho_ave = new((/nzlv,nlat,nlon/),typeof(pot_rho),pot_rho@_FillValue)
 pot_rho_ave = dim_avg_n_Wrap(pot_rho,0)
 delete(pot_rho)
 
;; Compute PV
 rho0 = 1025.
 pot_rho_ddc = new((/nzlv,nlat,nlon/),typeof(pot_rho_ave),pot_rho_ave@_FillValue)
 pot_rho_ddc = center_finite_diff_n(pot_rho_ave,tofloat(zlv),False,0,0)
 pot_vor = new((/nzlv,nlat,nlon/),typeof(pot_rho_ddc),pot_rho_ddc@_FillValue)
do i = 0, nzlv-1
 pot_vor(i,:,:) = (tofloat(f_cor)/rho0)*pot_rho_ddc(i,:,:)
end do
 
;; PV on density
 interp = 1
 pot_vor_rho = new((/nzrh,nlat,nlon/),typeof(pot_vor),pot_vor@_FillValue)
 pot_vor_rho = int2p_n_Wrap(zlv,pot_vor,zrh,interp,0)
 
;; Density on z
 pot_rho_z = new((/nzrh,nlat,nlon/),typeof(lev),lev@_FillValue)
 pot_rho_z = int2p_n_Wrap(zlv,lev,zrh,interp,0)
 delete([/pot_rho_ave,pot_rho_ddc/])
 
;; Create global attributes of the file (optional)
 fAtt               = True
 fAtt@title         = "PV and isopycnal surfaces"
 fAtt@source_file   = name+".ocean_month.nc"
 fAtt@Conventions   = "None"
 fAtt@creation_date = systemfunc ("date")
 fileattdef(fo,fAtt)                     

;; Define dimensions
 dim_names = (/"zrh","zlv","lat","lon"/)
 dim_sizes = (/nzrh,nzlv,nlat,nlon/)
 dim_unlim = (/False,False,False,False/)
 filedimdef(fo,dim_names,dim_sizes,dim_unlim)

;; Define file variables
 filevardef(fo,"pv",typeof(pot_vor),(/"zlv","lat","lon"/))
 filevardef(fo,"pv_rho",typeof(pot_vor_rho),(/"zrh","lat","lon"/))
 filevardef(fo,"z_rho",typeof(pot_rho_z),(/"zrh","lat","lon"/))
 
;; Write output file 
 pot_vor!0 = "zlv"
 pot_vor!1 = "lat"
 pot_vor!2 = "lon"
 pot_vor&zlv = zlv       
 pot_vor&lat = lat
 pot_vor&lon = lon 
 pot_vor@long_name = "Potential Vorticity" 
 pot_vor@units = "m-1 s-1"
 pot_vor@_FillValue = pot_vor@_FillValue
 pot_vor_rho!0 = "zrh"
 pot_vor_rho!1 = "lat"
 pot_vor_rho!2 = "lon"
 pot_vor_rho&zrh = zrh       
 pot_vor_rho&lat = lat
 pot_vor_rho&lon = lon 
 pot_vor_rho@long_name = "Potential Vorticity on isopycnal surface" 
 pot_vor_rho@units = "m-1 s-1"
 pot_vor_rho@_FillValue = pot_vor_rho@_FillValue
 pot_rho_z!0 = "zrh"
 pot_rho_z!1 = "lat"
 pot_rho_z!2 = "lon"
 pot_rho_z&zrh = zrh       
 pot_rho_z&lat = lat                 
 pot_rho_z&lon = lon                 
 pot_rho_z@long_name = "Elevation of isopycnal surfaces" 
 pot_rho_z@units = "m"
 pot_rho_z@_FillValue = pot_rho_z@_FillValue
 fo->zlv = (/zlv/)
 fo->zrh = (/zrh/)
 fo->lat = (/lat/)
 fo->lon = (/lon/)
 fo->pv = (/pot_vor/)
 fo->pv_rho = (/pot_vor_rho/)
 fo->z_rho  = (/pot_rho_z/)
 
 delete([/fo,zlv,zrh,lat,lon/])
 delete([/pot_vor,pot_vor_rho,pot_rho_z/])
end