; ==============================================================
; NCL script to compute POP MOC field offline from POP netcdf
; history files.  This routine is designed for the CCSM4 ocean 
; component.
; ==============================================================
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 "/work/yiwen89419/CCSM_ncl/ncl/yeager_util.ncl"

begin
; ==============================================================
; Input/Output
; ==============================================================
; workdir = getenv("WORKDIR")
; infile  = getenv("INFILE")
; sfx = get_file_suffix(infile,0)
; outfile  = sfx@fBase+".rho.nc"
; fin   = workdir+"/"+infile
; fout  = workdir+"/"+outfile

; fin = "./Monthly/TSR_b_20ct_ctrl_1.pop.h.192101-200012.nc"  
; fmean = "meanTSR_b_20ct_ctrl_1.pop.h.192101-200012.nc"
; fin ="./G_JRA55_test1_vdc013l/yearly/TSRyr_G_JRA55_test1_vdc013l.pop.030601-036612.nc"
; fmean = "./G_JRA55_test1_vdc013l/yearly/meanTSRyr_G_JRA55_test1_vdc013l.pop.030601-036612.nc"
fin ="./f09.B-PI.tn15.cmip6.j01.re/yearly/TSRyr_taiesm_501_700.nc"
fmean = "./f09.B-PI.tn15.cmip6.j01.re/yearly/meanTSRyr_taiesm_501_700.nc"

  sfx = get_file_suffix(fin,0)
  fout  = sfx@fBase+".rho_TS.nc"

 focn = "/home/project/MST108251/yhtseng00/model_output/cesm1/archive/b_20ct_ctrl_1/ocn/hist/b_20ct_ctrl_1.pop.h.1921-01.nc"

; ==============================================================
; Bounds
; ==============================================================
  x0 = -55.
  x1 = -35.
  y0 = 50.
  y1 = 60.
  z0 = 0.
  z1 = 200.
  if (x0.lt.0.) x0 = x0+360. end if
  if (x1.lt.0.) x1 = x1+360. end if
  z0 = z0*100.
  z1 = z1*100.
; ==============================================================
; Shift POP arrays since focus is on Atlantic where the "seam" is 
; ==============================================================
ishift = 200
nx = 320
ny = 384
nx0 = nx-ishift
 
; ==============================================================
; Get variables and reduce and shift
; ==============================================================
  f  = addfile (focn, "r")
  tlat = f->TLAT
  tlon = f->TLONG
  z_t  = f->z_t
  dz  = f->dz
  tarea = f->TAREA
  delete(f)
  z_t = z_t/100.
  dz = dz/100.
  tarea = tarea/(1.e4)

  tmp1 = tlat(:,0:ishift-1)
  tmp2 = tlat(:,ishift:)
  tlat(:,0:nx0-1) = tmp2
  tlat(:,nx0:) = tmp1
  delete(tmp1)
  delete(tmp2)

  tmp1 = tlon(:,0:ishift-1)
  tmp2 = tlon(:,ishift:)
  tlon(:,0:nx0-1) = tmp2
  tlon(:,nx0:) = tmp1
  delete(tmp1)
  delete(tmp2)

  tmp1 = tarea(:,0:ishift-1)
  tmp2 = tarea(:,ishift:)
  tarea(:,0:nx0-1) = tmp2
  tarea(:,nx0:) = tmp1
  delete(tmp1)
  delete(tmp2)

;; GET i-, j-indices which contain box defined by x0,y0,x1,y1
 ji = region_ind (tlat,tlon, y0, y1, x0, x1)
 
 f  = addfile (fmean, "r")
 tmn  = f->TEMP({z0:z1},:,:);
 smn  = f->SALT({z0:z1},:,:);
 rmn  = f->RHO({z0:z1},:,:);

;tmn=rm_single_dims(t);
;smn=rm_single_dims(s);
;rmn=rm_single_dims(r);
;delete(t)
;delete(s)
;delete(r)
 delete(f)

 f  = addfile (fin, "r")
 t0  = f->TEMP(:,{z0:z1},:,:)
 s0  = f->SALT(:,{z0:z1},:,:)
 r0  = f->RHO(:,{z0:z1},:,:)

 tmnfull = conform(t0,tmn,(/1,2,3/))
 smnfull = conform(s0,smn,(/1,2,3/))
 rmnfull = conform(r0,rmn,(/1,2,3/))
 tanom = t0 - tmnfull
 sanom = s0 - smnfull
 ranom = r0 - rmnfull

 tmp1 = t0(:,:,:,0:ishift-1)
 tmp2 = t0(:,:,:,ishift:)
 t0(:,:,:,0:nx0-1) = tmp2
 t0(:,:,:,nx0:) = tmp1
 delete(tmp1)
 delete(tmp2)

 tmp1 = s0(:,:,:,0:ishift-1)
 tmp2 = s0(:,:,:,ishift:)
 s0(:,:,:,0:nx0-1) = tmp2
 s0(:,:,:,nx0:) = tmp1
 delete(tmp1)
 delete(tmp2)

 tmp1 = r0(:,:,:,0:ishift-1)
 tmp2 = r0(:,:,:,ishift:)
 r0(:,:,:,0:nx0-1) = tmp2
 r0(:,:,:,nx0:) = tmp1
 delete(tmp1)
 delete(tmp2)

 tmp1 = tanom(:,:,:,0:ishift-1)
 tmp2 = tanom(:,:,:,ishift:)
 tanom(:,:,:,0:nx0-1) = tmp2
 tanom(:,:,:,nx0:) = tmp1
 delete(tmp1)
 delete(tmp2)

 tmp1 = sanom(:,:,:,0:ishift-1)
 tmp2 = sanom(:,:,:,ishift:)
 sanom(:,:,:,0:nx0-1) = tmp2
 sanom(:,:,:,nx0:) = tmp1
 delete(tmp1)
 delete(tmp2)

 tmp1 = ranom(:,:,:,0:ishift-1)
 tmp2 = ranom(:,:,:,ishift:)
 ranom(:,:,:,0:nx0-1) = tmp2
 ranom(:,:,:,nx0:) = tmp1
 delete(tmp1)
 delete(tmp2)


t = t0(:,:,ji(0):ji(1),ji(2):ji(3))
s = s0(:,:,ji(0):ji(1),ji(2):ji(3))
r = r0(:,:,ji(0):ji(1),ji(2):ji(3))
delt = tanom(:,:,ji(0):ji(1),ji(2):ji(3))
dels = sanom(:,:,ji(0):ji(1),ji(2):ji(3))
delr_on = ranom(:,:,ji(0):ji(1),ji(2):ji(3))

tmp = tarea(ji(0):ji(1),ji(2):ji(3))
delete(tarea)
tarea = tmp
delete(tmp)

r = r*1000.
r@units="kg/m^3"
delr_on = delr_on*1000.
delr_on@units="kg/m^3"

;delt =  dtrend_leftdim( t, False )
;dels =  dtrend_leftdim( s, False )
;delr_on =  dtrend_leftdim( r, False )
;delt = dim_rmvmean_n(t,0)
;dels = dim_rmvmean_n(s,0)
;delr_on = dim_rmvmean_n(r,0)

dims   = dimsizes(r)
nt = dims(0)
nz = dims(1)
ny = dims(2)
nx = dims(3)

tvol = new((/nz,ny,nx/),double)
do iz=0,nz-1
  tvol(iz,:,:) = tarea*dz(iz)
end do
tvol!0 = "z_t"
tvol!1 = "nlat"
tvol!2 = "nlon"
tvol@long_name = "Volume of T-grid cells"
tvol@units = "m^3"
tvol@_FillValue = tarea@missing_value
tvol@missing_value = tarea@missing_value

; ==============================================================
; Compute upperocean rho from mean T and S
; ==============================================================
rho = r
rhot = r
rhos = r
do it=0,nt-1
do iz=0,nz-1
  tmp = rho_mwjf(t(it,iz,:,:),s(it,iz,:,:),z_t(iz))
  rho(it,iz,:,:) = tmp

  tmp = rho_alphabeta(t(it,iz,:,:),s(it,iz,:,:),z_t(iz),0)
  rhot(it,iz,:,:) = -tmp*delt(it,iz,:,:)*rho(it,iz,:,:)

  tmp = rho_alphabeta(t(it,iz,:,:),s(it,iz,:,:),z_t(iz),1)
  rhos(it,iz,:,:) = tmp*dels(it,iz,:,:)*rho(it,iz,:,:)
end do
end do
copy_VarMeta(r,rho)
copy_VarMeta(r,rhot)
copy_VarMeta(r,rhos)
rho = rho*1000.
rhot = rhot*1000.
rhos = rhos*1000.

;delr_off =  dtrend_leftdim( rho, False )
delr_off = dim_rmvmean_n(rho,0)
delr_tpluss =  rhot+rhos

copy_VarMeta(r,delr_on)
copy_VarMeta(r,delr_off)
copy_VarMeta(r,delr_tpluss)

delr_off@long_name = "Density anomaly from monthly: rho(T, S)"
delr_tpluss@long_name = "Density anomaly from monthly: -alpha*deltaT+beta*deltaS"
rhot@long_name = "Density anomaly from monthly: -alpha*deltaT"
rhos@long_name = "Density anomaly from monthly: beta*deltaS"

printVarSummary(rho)
printVarSummary(rhot)
printVarSummary(rhos)
; ==============================================================
; Save to netcdf
; ==============================================================
  system("/bin/rm -f "+fout)
  outf = addfile(fout,"c")
  filedimdef(outf, "time",-1,True)    ; make "time" UNLIMITED
  outf->tvolume=tvol
  outf->rho_on=r
  outf->rho_off=rho
  outf->delr_on=delr_on
  outf->delr_off=delr_off
  outf->delr_tpluss=delr_tpluss
  outf->rho_T=rhot
  outf->rho_S=rhos
end