; ==============================================================
; Calculate gmoc EOFs 
; ==============================================================
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"

begin
; ==============================================================
; User defined parameters that specify region of globe to consider
; ==============================================================
focn = "/home/project/MST108251/yhtseng00/model_output/cesm1/archive/b_20ct_ctrl_1/ocn/hist/b_20ct_ctrl_1.pop.h.1921-01.nc"
  
  datadir= "/work/yiwen89419/CCSM_ncl/"
; fin = datadir+"./Yearly/BSFyr_b_20ct_ctrl_1.pop.h.192101-200012.nc"
; fin = datadir+"f09.B-PI.tn15.cmip6.j01.re/yearly/BSFyr_taiesm_501_700.nc"
  fin = datadir+"G_JRA55_test1_vdc013l/5_6th_cycle/BSFyr_024501-036612.nc" 
 
  rEarth = 6.37122e8 		; Earth radius in m

  field = "BSF"
  eofvar = "BSF"

  my0 = 245.		;; model year start
  my1 = 367.		;; model year end
  detrend = True	;; 1 ==> remove linear trend (& mean)
			;; 0 ==> remove mean
  neof     = 4         ; number of EOFs
  optEOF   = False      
  optETS   = False
  netCDF   = True
  smooth   = False
  highpass = False

; region = "natl"               ; "glob","pac","atl","natl",or "box"
  region = "box"                ; "glob","pac","atl","natl",or "box"
;  latN = 70.                    ; coords for computing EOF by "box"
;  latS = 15.
;  lonW = -80.
;  lonE = 0.
  latN = 60.                    ; coords for computing EOF by "box"
  latS = 50.
  lonW = -55.
  lonE = -35.

  fout = datadir+eofvar+"."+flt2string(my0)+"-"+flt2string(my1-1)+".eof."+region+"."
  if (smooth) then
;    fout = fout+"sm5."
;    fout = fout+"lz5+."
     fout = fout+"lz30+."
  end if
  if (highpass) then
     fout = fout + "hp15-."
  end if
  if (.not.detrend) then
     fout = fout + "nodt."
  end if
  fout = fout + "nc"

; ==============================================================
  f      = addfile (focn, "r")
  kmt    = f->KMT
  rmask  = f->REGION_MASK
  tarea  = f->TAREA
  uarea  = f->UAREA
  tlon   = f->TLONG
  tlat   = f->TLAT
  ulon   = f->ULONG
  ulat   = f->ULAT
  delete(f)
  ulon = where(ulon.gt.180.,ulon-360.,ulon)

  f      = addfile (fin, "r")
  tmp   = f->$field$
  time   = f->time

  time = time/365.
  time@units = "model year"
  tmp&time = tmp&time/365.
  tmp&time@units = "model year"

; data = tmp({my0:my1},0,:,:) 
  data = tmp({my0:my1},:,:) 
  delete(tmp)

  dims   = dimsizes(data)
  ntim   = dims(0)
  ny   = dims(1)
  nx   = dims(2)

; -----------------------------------------------------------------
; do region masking
; -----------------------------------------------------------------
  do y=0,ntim-1
    tmp = data(y,:,:)
   if (region.eq."glob") then
      mymask = rmask.gt.0
      tmp = where(mymask,tmp,tmp@_FillValue)
    end if

    if (region.eq."box") then
      mymask = rmask.gt.0.and.ulat.le.latN.and.ulat.ge.latS.and.ulon.le.lonE.and.ulon.ge.lonW
      tmp = where(mymask,tmp,tmp@_FillValue)
    end if

   ; 1=SO,2=Pac,3=Ind,6=Atl,7=Med,8=Lab
    if (region.eq."atl") then
      rmaskval = (/ 6, 8 /)          ; Basin for computing EOF by "basin"
      dimr = dimsizes(rmaskval)
      nr = dimr(0)
      if (nr.gt.1) then
        mymask = rmask.eq.rmaskval(0)
        do ir=1,nr-1
          mymask = mymask.or.rmask.eq.rmaskval(ir)
        end do
      else
        mymask = rmask.eq.rmaskval
      end if
      tmp = where(mymask,tmp,tmp@_FillValue)
    end if

    if (region.eq."pac") then
      rmaskval = (/ 2 /)          ; Basin for computing EOF by "basin"
      dimr = dimsizes(rmaskval)
      nr = dimr(0)
      if (nr.gt.1) then
        mymask = rmask.eq.rmaskval(0)
        do ir=1,nr-1
          mymask = mymask.or.rmask.eq.rmaskval(ir)
        end do
      else
        mymask = rmask.eq.rmaskval
      end if
      tmp = where(mymask,tmp,tmp@_FillValue)
    end if

   if (region.eq."natl") then
      rmaskval = (/ 6, 8, 9 /)          ; Basin for computing EOF by "basin"
      dimr = dimsizes(rmaskval)
      nr = dimr(0)
      if (nr.gt.1) then
        mymask = rmask.eq.rmaskval(0).and.tlat.ge.latS.and.tlat.le.latN
        do ir=1,nr-1
          mymask = (mymask.or.rmask.eq.rmaskval(ir)).and.tlat.ge.latS
        end do
      else
        mymask = rmask.eq.rmaskval
      end if
      tmp = where(mymask,tmp,tmp@_FillValue)
    end if

    data(y,:,:) = tmp
  end do
  
; -----------------------------------------------------------------
; Remove the temporal mean or trend at each point: make anomalies
; -----------------------------------------------------------------
if (.not.detrend) then
  WORK = dim_rmvmean_Wrap( data(nlat|:,nlon|:,time|:) )
  delete(data)
  data = WORK(time|:,nlat|:,nlon|:) 
  delete(WORK)
else
  WORK = dtrend_leftdim( data, False )
  delete(data)
  data = WORK
  delete(WORK)
end if

; -----------------------------------------------------------------
; weight all observations 
; -----------------------------------------------------------------
  wgt    =  tarea
  wdat   = data*conform(data, wgt, (/1,2/))  ; same units as "data"
  copy_VarMeta(data,wdat)

; -----------------------------------------------------------------
; do temporal smoothing pass?
; -----------------------------------------------------------------
if (smooth) then
  ;;;;;;;;;;;;;;;;
  ;  sm5
  ;;;;;;;;;;;;;;;;
  ;  twgt  = new(5,float)
  ;  twgt(:) = 1.
  ;  twgt  = twgt/sum(twgt)
  ; WORK = wgt_runave_Wrap(wdat(nlat|:,nlon|:,time|:),twgt,0)

  ;;;;;;;;;;;;;;;;
  ;  lz5+
; ;;;;;;;;;;;;;;;;
;  nwt = 15
;  pda = 5            ; longest period
;  pdb = 1            ; shortest period
;  fca = 1./pda       ;  ==> lowest allowed frequency
;  fcb = 1./pdb       ;  ==> highest allowed frequency
;  ihp = 0            ;  0 ==> low pass filter, fcb ignored
;  nsigma = 1.
;  twgt = filwgts_lanczos (nwt, ihp, fca, fcb, nsigma)
;  WORK = wgt_runave_Wrap(wdat(nlat|:,nlon|:,time|:),twgt,0)

  ;;;;;;;;;;;;;;;;
  ;  lz30+
  ;;;;;;;;;;;;;;;;
   nwt = 51
   fca = 1./30.         ; 20-yr period ==> highest allowed frequency
   ihp = 0              ;  1 ==> low pass filter
   nsigma = 1.
   twgt = filwgts_lanczos (nwt, ihp, fca, -999., nsigma)
   WORK = wgt_runave_Wrap(wdat(nlat|:,nlon|:,time|:),twgt,-1)

   wdat = WORK(time|:,nlat|:,nlon|:)
   delete(WORK)

end if

; -----------------------------------------------------------------
; reorder (z,lat,time) the *weighted* data: input to EOFs
; -----------------------------------------------------------------
  x      = wdat(nlat|:,nlon|:,time|:)
  x@long_name = "area weighted: "+data@long_name

  workeof    = eofunc_Wrap(x, neof, optEOF)     ; [evn | neof] x [nlat | 384] x [nlon | 320]
  workeof_ts = eofunc_ts_Wrap (x, workeof, optETS)  ; [evn | neof] x [time | nt]

; -----------------------------------------------------------------
; reweight eof and eof_ts by one std dev of PC time series 
; -----------------------------------------------------------------
  sd =  dim_stddev_Wrap(workeof_ts)
  eof_ts = workeof_ts/conform(workeof_ts,sd,(/0/))
  eof = workeof*conform(workeof,sd,(/0/))/conform(workeof,wgt,(/1,2/))
  copy_VarMeta(workeof,eof)
  copy_VarMeta(workeof_ts,eof_ts)

; -----------------------------------------------------------------
; reconstruct *weighted* anomaly data, as a check 
; -----------------------------------------------------------------
  x_recon = x                               ; (nlat,nlon,time)  
  x_recon = 0.0                    ; set to missing        

  do nt=0,ntim-1
    do ne=0,neof-1
     x_recon(:,:,nt) = x_recon(:,:,nt) + eof(ne,:,:)* eof_ts(ne,nt) 
    end do  
  end do  
  x_recon@long_name = "x_recon: "+x@long_name

; -----------------------------------------------------------------
; reconstruct original *unweighted* *anomaly* data, as a check 
; -----------------------------------------------------------------
; X_RECON   = x_recon/conform(x_recon, wgt, (/0,1/))
  X_RECON   = x_recon
  copy_VarMeta(x_recon, X_RECON)
  X_RECON@long_name = "X_RECON: x_recon/wgt"

;------------------------------------------------------------
; save EOFs to netcdf
;------------------------------------------------------------
  if (netCDF) then
      system("/bin/rm -f "+fout)
      outf = addfile(fout,"c")
      outf->time=(/eof_ts&time/)     ; yyyy
      outf->data=data
      outf->data_recon=X_RECON         
      outf->eof=eof
      outf->eofts=eof_ts
  end if
 
end