load "~/loadFiles.ncl"

begin
;************************************************
; Calculate PC correlation and make plots by YM Cheng
;************************************************
;;=============================================
;;=========== beginning of settings ===========
;;=============================================
fin="EOF.nc"
fout="correlation.nc"
diri="."
diro="."
varname="vwnd"; EOF varname
yrs=1984; year start, set the time period and season of interest for EOF
yre=2013; year end
mns=7; month start
mne=9; month end
jump_extract=153	; 153 days from June 1 to Oct 31
jump_store	=92		; 92  days from July 1 to Sep 30
;;=============================================
;;============== end of settings ==============
;;========== can leave the rest unchanged======
;;=============================================

print("====================================")
print("=====Calculating PC correlation=====")
print("====================================")
print(" fin: "+diri+"/"+fin)
print("fout: "+diro+"/"+fout)

print("reading PCs")
if(varname .eq. "brtmp")
	incre=2
else
	incre=1
end if

inf=addfile(diri+"/"+fin,"r")
tmp=inf->eof_ts(:,::incre)
cd=cd_calendar(tmp&time,0)

ind_jjaso=ind(cd(:,0) .ge. yrs .and. cd(:,0) .le. yre .and. \
cd(:,1) .ge. mns-1 .and. cd(:,1) .le. mne+1)
var=tmp(:,ind_jjaso)

ind_jas  =ind(cd(:,0) .ge. yrs .and. cd(:,0) .le. yre .and. \
cd(:,1) .ge. mns .and. cd(:,1) .le. mne)
pc=tmp(:,ind_jas)
delete(tmp)
ndim=dimsizes(pc)
maxlag=48
tres=4
lags=ispan(-maxlag,maxlag,1)
nlags=dimsizes(lags)
var_lag=new(ndim(1),typeof(var)); to hold the pc data for correlation
cor=new((/ndim(0),ndim(0),nlags/),typeof(pc))

cor!0="eof"
cor!1="eof2"
cor!2="lag"

cor&lag=fspan(-12,12,2*maxlag+1)

print("doing correlation!")
cd_cal=cd_calendar(var&time,0)
ind_jul_1_1984=ind(cd_cal(:,0) .eq. yrs .and. cd_cal(:,1) .eq. mns .and. cd_cal(:,2) .eq. 1 .and. cd_cal(:,3) .eq. 0)
; print(ind_jul_1_1984)

do x=0, ndim(0)-1
  ;print("====== x= "+x+" ======")
  do y=0, ndim(0)-1
    ;	print("====== y= "+y+" ======")
    do i=0,nlags-1,1
    ;		print("=====Lag "+ lags(i) +" now=====")
    var_lag=0.0;; var_lag will differ for every lag, set to zero every time
    ;;======================================
    ;;;; taking out the required time series
    ind_extract_start=ind_jul_1_1984+lags(i)
    ind_extract_end  =ind_extract_start+92*tres-1
    ;		pvs(var_lag)
    do yr=yrs, yre,1
      ;					print("=== Year "+yr+" ===")
      j=yr-yrs
      ind_store_start:=j*92*tres
      ind_store_end:=ind_store_start+92*tres-1
      var_lag(ind_store_start:ind_store_end)=var(y,ind_extract_start:ind_extract_end)
      ind_extract_start:=ind_extract_start+jump_extract*tres
      ind_extract_end:=ind_extract_start+jump_store*tres-1

      end do
      ;			print("Initial time="+cd_calendar(var_lag&time(0),-3))
      ;			print("  Last time ="+cd_calendar(var_lag&time(ndim(1)-1),-3))
      ; Done with time series extraction
      cor(x,y,i)=(/escorc(pc(x,:),var_lag)/)
    end do
  end do
end do

system("/bin/rm -f " + diro +"/"+ fout)
outf=addfile(diro+"/"+fout,"c")
cor@info_tag="autocorrelation matrix"
outf->cor=	cor

print("================================")
print("=======plotting correlation======")
print("================================")

wksName=diro+"/pc.correlation"
wksType="x11"
wks   = gsn_open_wks (wksType,wksName)                ; open workstation

do i=0 ,7
  res       = True
  res@gsnDraw=False
  res@gsnFrame=False
  res@gsnLeftString	= ""
  res@gsnRightString	= "Correlation with PC "+(i+1)
  res@gsnStringFontHeightF=0.012
  res@tiXAxisFontHeightF=0.014
  res@tiYAxisFontHeightF=0.014
  res@tiXAxisString     = "Day"
  res@tiYAxisString     = "Correlation"
  res@xyLineColors	=(/"red","red3","blue","royalblue1","gold","gold4","darkseagreen1","darkseagreen4"/)
  res@xyLineThicknesses=(/5,5,5,5,3,3,3,3/)
  res@xyExplicitLegendLabels = "  PC "+ispan(1,8,1)         ; create explicit labels
  res@trYMinF	= -1.0
  res@trYMaxF	= 1.0
  res@vpWidthF =.7
  res@vpHeightF =.45
  res@tmYLMode="Explicit"
  res@tmYLValues=sprintf("%4.1f",fspan(-1,1,11))
  res@tmYLLabels=sprintf("%4.1f",fspan(-1,1,11))
  res@pmLegendDisplayMode    = "Always"            ; turn on legend
  res@pmLegendSide           = "Bottom"               ; Change location of
  res@pmLegendParallelPosF   = .125                  ; move units right
  res@pmLegendOrthogonalPosF = -0.48                ; move units down
  res@pmLegendWidthF         = 0.10                ; Change width and
  res@pmLegendHeightF        = 0.13                ; height of legend.
  res@lgPerimOn              = False               ; turn off box around
  res@lgLabelFontHeightF     = .01                 ; label font height

  plot = gsn_csm_xy (wks,cor&lag,cor(i,0:7,:),res) ; create plot
  draw(plot)
  frame(wks)
end do
end