# This is a script to plot correlations between FD PC1 and SPEI02
#===================================================================================================================
library(zoo)
library(reshape2)
library(colorRamps)
source('./my.functions.R')
library(fields)
library(maps)
#===================================================================================================================

years=c(2000:2013)
psig=0.05
col.pal <- colorRampPalette(c("darkblue","steelblue2","white","indianred2","red4"))(101) 
spei=c('SPEI01','SPEI02','SPEI03')
ind=c('01','02','03')
seas=c('DJF','MAM','JJA','SON')

pdf(file=paste('./plots/EOF/corr.PC1-SPEI02.2000-2013.pdf', sep=''), width=11, height=8.5) # landscape orientation

layout(matrix(c(1:4), 2, 2, byrow = TRUE)) # layout.show(nf)
par(oma=c(0.5,0,3,5), mar=c(1,3.5,3,1.5), mgp = c(0,0.8,0))

#for (n in 1:length(spei)){
n=2
print(spei[n])
  for (k in 1:length(seas)){
  print(seas[k])

  #--- Load PCs 
  load(paste('./data/EOF/EOFs_2000-2013-',seas[k],'_FD.RData',sep=''))
  PCs = results$PCs[,1] 
  sites=results$sites
  
  #---Load SPEI
  load(paste('./data/LOCA.obs/dtds/dtds.mm-spei',ind[n],'.2000-2013.RData',sep=''))
  df = ss$iav  
  temp=array(df, dim=c(dim(df)[1], dim(df)[2], 12, dim(df)[3]/12))
  df2=temp[,,(3*k-2):(3*k),]
  df3 = array(df2, dim=c(dim(df)[1], dim(df)[2], dim(df2)[3]*dim(df2)[4]))
  
  #--- Calculate corr map
  corr.map = array(NA, dim=c(length(ss$lon), length(ss$lat))) # lon x lat
  for(i in 1:length(ss$lon)){ # lon
    for (j in 1:length(ss$lat)){ # lat
      ts = df3[i,j,]
      ts.cmpl = ts[complete.cases(ts)]
      if (length(ts.cmpl) >= 3){ 
        pearson.corr = cor.test(PCs, ts, na.action='na.exclude')
        if (is.na(pearson.corr$p.value)==FALSE & pearson.corr$p.value <= psig){ # only include statistically significant correlations
          corr.map[i,j] = pearson.corr$estimate
        } else{ corr.map[i,j] = NA }
      }
    }}
  
  res = list('hcorr'=corr.map, 'psig'=psig, 'lat'=ss$lat, 'lon'=ss$lon)
  save(res, file=paste('./data/EOF/corr/FD/corr_',seas[k],'-PC1_',spei[n],'.RData',sep=''))
  
  #--- Plot
  image.plot(x=res$lon, y=res$lat, z=res$hcorr, col=col.pal, zlim=c(-1,1), xlab='', ylab='', xaxt='n', yaxt='n')
  map("world", add=TRUE, col='black') 
  map("state", add=TRUE, col='black')
  mtext(paste(seas[k]), cex=1.5, side=3, line=0.5)
 
  lwd=2.5
  if(k==1){rect(-117,30,-105,38, lwd=lwd)}
  if(k==2){rect(-115,28,-103,38, lwd=lwd)}
  if(k==3){rect(-115,29,-103,40, lwd=lwd)}
  if(k==4){rect(-113,30,-103,39, lwd=lwd)}

  }
mtext(paste('2000-2013 corr(FD PC1 & SPEI',ind[n],'), p<',psig,sep=''), cex=1.5, side=3, line=0.5, outer=TRUE)
#}

dev.off()

#===================================================================================================================
#===================================================================================================================