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

begin
;************************************************
; Calcuate map regression with PCs by YM Cheng
; The design is for summer (JAS) from 1984 - 2013
; Using data from Geroge:
; 1. 2.5 degree ERA interim.4x
; (e.g., uwnd.erai.pl.2.5.gauss.nc)
; 2. 2.5 degree clausir.8x.30ns.olrgrid.nc
; 3. PCs from EOF analysis
; Note that errors may occrus if using different seasons, specifically
; when including Feburary (due to leap years) and winter (e.g., Dec-Feb due to
; the subscriptting of data)
;************************************************

;;=============================================
;;=========== beginning of settings ===========
;;=============================================
maxlag=48; total lags
tres=4 ; data are 4 times a day for Interim, will read brtmp 4 times a day

diri = "./raw" ; dir of raw data for brtmp
dir_eof="."; where the EOF.nc located
diro="./regression"; directory out, where to store the regression file
feof="EOF.nc"
varname="vwnd"; varname of the EOF analysis
re_varname="vwnd"; "vwnd"/"uwnd"/"sf"/"brtmp" the variable name of the field to regress with
reg_lev=700; specify the level of interest to read in file, null for brtmp
re_pc=1; conduct the regression on PC1

yrs=1984; year start, set the time period and season of interest for regression
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("=========Regression "+re_varname+"=========")
print("=================================")

if(re_varname .eq. "brtmp"	)
	reg_lev=0 ; reg_lev is null here
	fin:= "clausir.8x.30ns.olrgrid.nc"
	fout="brtmp.regres.pc"+re_pc+".nc"
else
	reg_lev=reg_lev
	fin:= re_varname+".erai.pl.2.5.gauss."+reg_lev+"hPa.nc"
	; fin="vwnd.700.4x.era2.nc";***
	fout= re_varname+".regres.pc"+re_pc+"."+reg_lev+"hPa.nc"
end if
print("=================================")
print("========reading raw data=========")
print("Could take a while for Interim data")
print("=================================")

inf			= addfile(diri+"/"+fin,"r")
time		= inf->time
cd      = cd_calendar(time,0)
indx    = ind(cd(:,0) .ge. yrs .and. cd(:,0) .le. yre .and. \
          cd(:,1) .ge. (mns-1) .and. cd(:,1) .le. (mne+1))
          ; get the index for the period and season of interst, adding extra
					; two months'worth of data for lag correlation
					; i.e., the field to be regressed against needs to extend
					; from Jun 1 to Oct 30 to do lag regression in JAS case

if(re_varname .eq. "brtmp"	)
	incre=2 ; brtmp is every 3 hours, take only every other time step
	vartmp	:= inf->$re_varname$
	vartmp  := vartmp(indx,:,:)
	vartmp 	:= vartmp(::incre,:,:)
else
	incre=1
	vartmp	:= inf->$re_varname$; for only level of data
	; vartmp	:= inf->v; for only level of data ***
	vartmp	:= vartmp(indx,:,:);
	vartmp 	:= vartmp(::incre,:,:)
end if
; test out if lonFlip is needed, i.e., applied when lon=[0,360)
if(any(getvardims(vartmp) .eq. "lon" .or. getvardims(vartmp) .eq. "LON"))
  if(vartmp&lon(dimsizes(vartmp&lon)-1) .gt. 180.)
    vartmp = lonFlip(vartmp)
  end if
end if
var	=	vartmp(lat|:,lon|:,time|:); rearrange the dimension order
delete(vartmp)
; printVarSummary(var)
ndim	=	dimsizes(var)



print("reading PC")
inf :=	addfile(dir_eof+"/"+feof,"r")
cd	:= cd_calendar(inf->time,0)
indx:=	ind(cd(:,0) .ge. yrs .and. cd(:,0) .le. yre .and. \
cd(:,1) .ge. mns .and. cd(:,1) .le. mne)
pc_tmp	=	inf->eof_ts(re_pc-1,indx)
if(varname .eq. "brtmp") ; the varname of the EOF
	incre:=2
else
	incre:=1
end if
pc	=	pc_tmp(::incre) ; pc_tmp(::2) when regress against clausir pc (8x a day)
delete(pc_tmp)
; printVarSummary(pc)
ntim=dimsizes(pc&time)
print("finished reading data")
;************************************************
; Extract data for lag correlation and calculate the regression coefficients
;************************************************
lags=ispan(-maxlag,maxlag,1)
nlags=dimsizes(lags)

; create a matrix for each grid point to hold data for regression at
; different lags, dim=(lat,lon,time)
var_lag=new((/ndim(0),ndim(1),ntim/),typeof(var))

; create a coefficient matrix to hold regression coefficient, dim=(lags,lat,lon)
regresCoeff=new((/nlags,ndim(0),ndim(1)/),typeof(var))
regresCoeff(0,:,:)=var(:,:,0); copy meta data from the original var
regresCoeff!0="lag"
regresCoeff&lag=lags*(1./tres) ; put lags in terms of days
regresCoeff&lag@long_name="Days w.r.t. to Jul 1 00Z"
regresCoeff=0.0
regresCoeff@long_name := "regression coefficient of "+re_varname+\
" and PC"+re_pc+" from "+varname+" EOF"

correlation=regresCoeff
correlation=0.0
correlation@long_name="Correlation coefficient"

; create a predictand matrix to hold predictands on a map, same dimension as regresCoeff
var_predict=regresCoeff
var_predict=0.0
var_predict@long_name="predictand using 1 standard deviation anomaly"

; extract data for different lags
; Here's the messy part, may think of a better way to extract data for different lags.
print("doing correlation/regression!")
cd:=cd_calendar(var&time,0)
ind_first_time =ind(cd(:,0) .eq. yrs .and. cd(:,1) .eq. mns .and. \
cd(:,2) .eq. 1 .and. cd(:,3) .eq. 0)
; the index for the first season of interest,
; In this case, 00Z, July 1, 1984, we need this index as a reference point
; to extract data for different lags
;print(ind_first_time)

std=stddev(pc)


do i=0,nlags-1,1
	print("======= Lag "+ lags(i)+" ======="); lags=-48~+48
	var_lag=0.0;;; var_lag will differ for every lag, set to zero every time
  ; extract the required time series at different lags
  ; ind_var_start and ind_var_end are the index in var,
	; use these two indices to extract data and store in var_lag
  ; ind_store_start and ind_store_end are indices to put the extracted var data
	; at proper postion in var_lag at different lags
	ind_extract_start=ind_first_time+lags(i)
	ind_extract_end  =ind_extract_start+jump_store*tres-1
	do yr=yrs, yre,1
;		print("=== Year "+yr+" ===")
		j=yr-yrs
		; print("j="+j)
		ind_store_start:=j*jump_store*tres
		ind_store_end:=ind_store_start+jump_store*tres-1

		var_lag(:,:,ind_store_start:ind_store_end)=var(:,:,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(ntim-1),-3))
;;======================================
; Done with time series extraction
; calculate regression coefficient and correlation coefficient
; scale the perturbation to one standard deviation of PC
	rc=regCoef(pc,var_lag)
	regresCoeff(i,:,:)= (/rc/)
	correlation(i,:,:)=(/escorc(pc,var_lag)/)
	var_predict(i,:,:)=(/std*regresCoeff(i,:,:)/)

end do

	system("rm -f "+diro+"/"+fout)
	outf=addfile(diro+"/"+fout,"c")
	outf->correlation=correlation
	outf->regrecoeff=regresCoeff
	outf->$re_varname$=var_predict
exit
end