
# Creation:  4/2018
# Author: Nicola Cortesi
# Aim: To compute and visualize the first X modes of variability for a given spatial domain and period from a reanalysis or from a forecast system
#
# Assumption:  # I/O: input must be in a format compatible with Load(); one output .RData file is created in the 'workdir' folder.
#
# Branch: teleconnections?


rm(list=ls())
library(s2dverification)  # for the function Load()
library(future)
library(Kendall)
source('/esarchive/scratch/scripts/Rfunctions.R')

## Coordinates of the box enclosing the study region (the Euro-Atlantic region in this case):
lat.max <- 80 
lat.min <- 20 
lon.max <- 60 
lon.min <- 270 

## years:
year.start <- 1981
year.end <- 2015

## choose how many teleconnections you want to project:
n.telec = 4 

## choose a variable you want to project to: "psl" or "g500" or "sfcWind"
psl <- "g500"

## choose a month from 1 to 12 of a season (13 = DJF, 14 = JJA, 15 = MAM, 16 = SON):
m <- 13

## chose a reanalysis for the variable you want to project to:
##rean  <- paste0('/esnas/recon/noaa/ncep-reanalysis/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc') ## for NCEP (but there is no g500 here!!!)
rean  <- paste0('/esnas/recon/ecmwf/erainterim/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc') ## for ERA-Interim

## working dir where all output files will be generated:
workdir <- "/esarchive/scratch/ncortesi/Regimes/86_TPs_NCEP_psl_1981-2015"

###################################################################################################################################################
##                                                             Composition                                                                       ##
###################################################################################################################################################

n.years <- year.end - year.start + 1

## load monthly time series of teleconnection indices from CPC: (only works if n.telec = 4)
NAO <- read.table(file="/esarchive/scratch/ncortesi/Regimes/series_teleconnexion_CPC/NAO.txt", header=FALSE)
EA <- read.table(file="/esarchive/scratch/ncortesi/Regimes/series_teleconnexion_CPC/EA.txt", header=FALSE)
EAWR <- read.table(file="/esarchive/scratch/ncortesi/Regimes/series_teleconnexion_CPC/EAWR.txt", header=FALSE)
SCAND <- read.table(file="/esarchive/scratch/ncortesi/Regimes/series_teleconnexion_CPC/SCAND.txt", header=FALSE)

names(NAO) <- names(EA) <- names(EAWR) <-  names(SCAND) <- c("year","month","index")

## create winter averages and select only years in the chosen period:
NAO.winter <- data.frame(year=1950:2016, month=13, dec=NAO[which(NAO[,2] == 12), 3], jan=NAO[which(NAO[,2] == 1), 3][-1], feb=NAO[which(NAO[,2] == 2), 3][-1])
NAO.winter$index <- (NAO.winter$dec + NAO.winter$jan + NAO.winter$feb) / 3
NAO.winter.y <- NAO.winter[(year.start - 1950 + 1):(year.end - 1950 + 1),c(1,2,6)]

EA.winter <- data.frame(year=1950:2016, month=13, dec=EA[which(EA[,2] == 12), 3], jan=EA[which(EA[,2] == 1), 3][-1], feb=EA[which(EA[,2] == 2), 3][-1])
EA.winter$index <- (EA.winter$dec + EA.winter$jan + EA.winter$feb) / 3
EA.winter.y <- EA.winter[(year.start - 1950 + 1):(year.end - 1950 + 1),c(1,2,6)]

EAWR.winter <- data.frame(year=1950:2016, month=13, dec=EAWR[which(EAWR[,2] == 12), 3], jan=EAWR[which(EAWR[,2] == 1), 3][-1], feb=EAWR[which(EAWR[,2] == 2), 3][-1])
EAWR.winter$index <- (EAWR.winter$dec + EAWR.winter$jan + EAWR.winter$feb) / 3
EAWR.winter.y <- EAWR.winter[(year.start - 1950 + 1):(year.end - 1950 + 1),c(1,2,6)]

SCAND.winter <- data.frame(year=1950:2016, month=13, dec=SCAND[which(SCAND[,2] == 12), 3], jan=SCAND[which(SCAND[,2] == 1), 3][-1], feb=SCAND[which(SCAND[,2] == 2), 3][-1])
SCAND.winter$index <- (SCAND.winter$dec + SCAND.winter$jan + SCAND.winter$feb) / 3
SCAND.winter.y <- SCAND.winter[(year.start - 1950 + 1):(year.end - 1950 + 1),c(1,2,6)]

## create summer averages and select only years in the chosen period:
NAO.summer <- data.frame(year=1950:2016, month=14, jun=NAO[which(NAO[,2] == 6), 3], jul=NAO[which(NAO[,2] == 7), 3], aug=NAO[which(NAO[,2] == 8), 3])
NAO.summer$index <- (NAO.summer$jun + NAO.summer$jul + NAO.summer$aug) / 3
NAO.summer.y <- NAO.summer[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

EA.summer <- data.frame(year=1950:2016, month=14, jun=EA[which(EA[,2] == 6), 3], jul=EA[which(EA[,2] == 7), 3], aug=EA[which(EA[,2] == 8), 3])
EA.summer$index <- (EA.summer$jun + EA.summer$jul + EA.summer$aug) / 3
EA.summer.y <- EA.summer[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

EAWR.summer <- data.frame(year=1950:2016, month=14, jun=EAWR[which(EAWR[,2] == 6), 3], jul=EAWR[which(EAWR[,2] == 7), 3], aug=EAWR[which(EAWR[,2] == 8), 3])
EAWR.summer$index <- (EAWR.summer$jun + EAWR.summer$jul + EAWR.summer$aug) / 3
EAWR.summer.y <- EAWR.summer[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

SCAND.summer <- data.frame(year=1950:2016, month=14, jun=SCAND[which(SCAND[,2] == 6), 3], jul=SCAND[which(SCAND[,2] == 7), 3], aug=SCAND[which(SCAND[,2] == 8), 3])
SCAND.summer$index <- (SCAND.summer$jun + SCAND.summer$jul + SCAND.summer$aug) / 3
SCAND.summer.y <- SCAND.summer[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

## create spring averages and select only years in the chosen period:
NAO.spring <- data.frame(year=1950:2016, month=15, mar=NAO[which(NAO[,2] == 3), 3], apr=NAO[which(NAO[,2] == 4), 3], may=NAO[which(NAO[,2] == 5), 3])
NAO.spring$index <- (NAO.spring$mar + NAO.spring$apr + NAO.spring$may) / 3
NAO.spring.y <- NAO.spring[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

EA.spring <- data.frame(year=1950:2016, month=15, mar=EA[which(EA[,2] == 3), 3], apr=EA[which(EA[,2] == 4), 3], may=EA[which(EA[,2] == 5), 3])
EA.spring$index <- (EA.spring$mar + EA.spring$apr + EA.spring$may) / 3
EA.spring.y <- EA.spring[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

EAWR.spring <- data.frame(year=1950:2016, month=15, mar=EAWR[which(EAWR[,2] == 3), 3], apr=EAWR[which(EAWR[,2] == 4), 3], may=EAWR[which(EAWR[,2] == 5), 3])
EAWR.spring$index <- (EAWR.spring$mar + EAWR.spring$apr + EAWR.spring$may) / 3
EAWR.spring.y <- EAWR.spring[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

SCAND.spring <- data.frame(year=1950:2016, month=15, mar=SCAND[which(SCAND[,2] == 3), 3], apr=SCAND[which(SCAND[,2] == 4), 3], may=SCAND[which(SCAND[,2] == 5), 3])
SCAND.spring$index <- (SCAND.spring$mar + SCAND.spring$apr + SCAND.spring$may) / 3
SCAND.spring.y <- SCAND.spring[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

## create autumn averages and select only years in the chosen period:
NAO.autumn <- data.frame(year=1950:2016, month=16, sep=NAO[which(NAO[,2] == 9), 3], oct=NAO[which(NAO[,2] == 10), 3], nov=NAO[which(NAO[,2] == 11), 3])
NAO.autumn$index <- (NAO.autumn$sep + NAO.autumn$oct + NAO.autumn$nov) / 3
NAO.autumn.y <- NAO.autumn[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

EA.autumn <- data.frame(year=1950:2016, month=16, sep=EA[which(EA[,2] == 9), 3], oct=EA[which(EA[,2] == 10), 3], nov=EA[which(EA[,2] == 11), 3])
EA.autumn$index <- (EA.autumn$sep + EA.autumn$oct + EA.autumn$nov) / 3
EA.autumn.y <- EA.autumn[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

EAWR.autumn <- data.frame(year=1950:2016, month=16, sep=EAWR[which(EAWR[,2] == 9), 3], oct=EAWR[which(EAWR[,2] == 10), 3], nov=EAWR[which(EAWR[,2] == 11), 3])
EAWR.autumn$index <- (EAWR.autumn$sep + EAWR.autumn$oct + EAWR.autumn$nov) / 3
EAWR.autumn.y <- EAWR.autumn[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

SCAND.autumn <- data.frame(year=1950:2016, month=16, sep=SCAND[which(SCAND[,2] == 9), 3], oct=SCAND[which(SCAND[,2] == 10), 3], nov=SCAND[which(SCAND[,2] == 11), 3])
SCAND.autumn$index <- (SCAND.autumn$sep + SCAND.autumn$oct + SCAND.autumn$nov) / 3
SCAND.autumn.y <- SCAND.autumn[(year.start - 1950 + 1):(year.end - 1950 + 1), c(1,2,6)]

## select only teleconnection in the chosen years:
row.start1 <- which(NAO[,1] == year.start)[1]
row.end1 <- which(EA[,1] == year.end)[l(which(EA[,1] == year.end))]
row.start2 <- which(EA[,1] == year.start)[1]
row.end2 <- which(EA[,1] == year.end)[l(which(EA[,1] == year.end))]
row.start3 <- which(EAWR[,1] == year.start)[1]
row.end3 <- which(EAWR[,1] == year.end)[l(which(EAWR[,1] == year.end))]
row.start4 <- which(SCAND[,1] == year.start)[1]
row.end4 <- which(SCAND[,1] == year.end)[l(which(SCAND[,1] == year.end))]

NAOy <- NAO[row.start1:row.end1,]
EAy <- EA[row.start2:row.end2,]
EAWRy <- EAWR[row.start3:row.end3,]
SCANDy <- SCAND[row.start4:row.end4,]

NAOy <- rbind(NAOy, NAO.winter.y, NAO.summer.y, NAO.spring.y, NAO.autumn.y) ## add winter and summer years after the 12 months
EAy <- rbind(EAy, EA.winter.y, EA.summer.y, EA.spring.y, EA.autumn.y) ## add winter years after the 12 months
EAWRy <- rbind(EAWRy, EAWR.winter.y, EAWR.summer.y, EAWR.spring.y, EAWR.autumn.y) ## add winter years after the 12 months
SCANDy <- rbind(SCANDy, SCAND.winter.y, SCAND.summer.y, SCAND.spring.y, SCAND.autumn.y) ## add winter years after the 12 months

test.CPC <- array(c(NAOy[which(NAOy[,2] == m), 3], EAy[which(EAy[,2] == m), 3], EAWRy[which(EAWRy[,2] == m), 3], SCANDy[which(SCANDy[,2] == m), 3]), c(n.years, 4))

## Load DJF data:
if(m == 13) user.period <- '1201'
if(m == 14) user.period <- '0601'
if(m == 15) user.period <- '0301'
if(m == 16) user.period <- '0901'

pslFull <-  Load(var = psl, exp = NULL, obs = list(list(path=rean)), 
                 paste0(year.start:year.end, user.period), storefreq = 'monthly', leadtimemin = 1, leadtimemax = 3, 
                 output = 'lonlat', latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max, nprocs=availableCores())

## average data over DJF of JJA:
seasonal <- Mean1Dim(pslFull$obs, 4)

## average data over all years:
clim <- Mean1Dim(seasonal, 3)

## measure DJF or JJA anomalies:
anom <- drop(seasonal - InsertDim(clim, 3, n.years))

lat <- pslFull$lat
lon <- pslFull$lon
n.lat <- length(lat)
n.lon <- length(lon)

## Project the PCs of the CPC on g500 or SLP anomalies for the chosen month or season:
CorrPatterns2(anom, lon, lat, test.CPC, nrow=4, ncol=1, show=T)

## select numb of truncated EOFs for the varimax rotation: (0 means no rotation)
my.rot <- 4 ##c(0, 4, 5, 6, 7)

n.rot <- length(my.rot)
proj.cor <- raw.REOFs <- array(NA, c(4*n.rot, n.lat, n.lon)) ## array with the pattern correlations for each teleconnection
spat.cor <- temp.cor <- var.exp <- array(NA, c(4, n.rot)) ## arrays with the the spatial and temporal correlations between the PCs of REOFs()and of the CPC
i <- 1
for(t in my.rot){
    
    if(t == 0){
        ## ## compute EOFs with EOF(): (EOFs already correct the anomalies for the latitude)
        test <- EOF(ano = anom, lon = lon, lat = lat, neofs = 4)
        test$RPCs <- test$PCs ## to correlate with CPC
        test$REOFs <- test$EOFs ## to be consistent with REOF()
    } else {
        test <- REOF(ano = anom, lon = lon, lat = lat, ntrunc = t)
    }
    
    ## modify the sign of the CPCs to provide the same sign of the patterns of the CPC:
    test2 <- test
    
    if(psl == "g500" && m == 13) { test2$RPCs[,2] <- -test$RPCs[,2]; test2$REOFs[2,,] <- -test$REOFs[2,,] }
    if(psl == "g500" && m == 13 && t <= 4) { test2$RPCs[,3] <- -test$RPCs[,3]; test2$REOFs[3,,] <- -test$REOFs[3,,] }
    if(psl == "psl" && m == 13 && t == 4) {  test2$RPCs[,1] <- -test$RPCs[,1]; test2$REOFs[1,,] <- -test$REOFs[1,,] }
    if(psl == "psl" && m == 13 && t == 4) {  test2$RPCs[,3] <- -test$RPCs[,3]; test2$REOFs[3,,] <- -test$REOFs[3,,] }

    if(psl == "g500" && m == 14) { test2$RPCs[,2] <- -test$RPCs[,2]; test2$REOFs[2,,] <- -test$REOFs[2,,] }
    if(psl == "g500" && m == 14) { test2$RPCs[,3] <- -test$RPCs[,3]; test2$REOFs[3,,] <- -test$REOFs[3,,] }
    ##if(psl == "psl" && m == 14) { test2$RPCs[,2] <- -test$RPCs[,2]; test2$REOFs[2,,] <- -test$REOFs[2,,] }

    if(psl == "g500" && m == 15) {
        test2$RPCs[,1] <- -test$RPCs[,1]; test2$REOFs[1,,] <- -test$REOFs[1,,]
        ##test2$RPCs[,2] <- -test$RPCs[,2]; test2$REOFs[2,,] <- -test$REOFs[2,,]
        test2$RPCs[,3] <- -test$RPCs[,3]; test2$REOFs[3,,] <- -test$REOFs[3,,]
        if(t > 0) { test2$RPCs[,4] <- -test$RPCs[,4]; test2$REOFs[4,,] <- -test$REOFs[4,,] }
    }

     if(psl == "g500" && m == 16) {
        test2$RPCs[,2] <- -test$RPCs[,2]; test2$REOFs[2,,] <- -test$REOFs[2,,]
        test2$RPCs[,3] <- -test$RPCs[,3]; test2$REOFs[3,,] <- -test$REOFs[3,,]
    }


    ## draw the patterns to check if their sign is the same of CPC's:
    ## CorrPatterns(anom, lon, lat, PCs=test2$RPCs, nrow=4, ncol=1, show=T)

    ## save temporal correlation values of the pattern map:
    REOF.proj.corr <- CorrPatterns(anom, lon, lat, PCs=test2$RPCs, nrow=4, ncol=1, show=F)

    proj.cor[i + 0*n.rot,,] <- REOF.proj.corr[1,,]
    proj.cor[i + 1*n.rot,,] <- REOF.proj.corr[2,,]
    proj.cor[i + 2*n.rot,,] <- REOF.proj.corr[3,,]
    proj.cor[i + 3*n.rot,,] <- REOF.proj.corr[4,,]
    
    ##CorrPatterns2(anom, lon, lat, test.CPC, nrow=4, ncol=1, show=T)
    REOF.proj.corr.CPC <- CorrPatterns(anom, lon, lat, test.CPC, nrow=4, ncol=1, show=F)

    ## spatial correlations:
    spat.cor[1,i] <- cor(as.vector(REOF.proj.corr.CPC[1,,]), as.vector(REOF.proj.corr[1,,]))  ## first REOFs
    spat.cor[2,i] <- cor(as.vector(REOF.proj.corr.CPC[2,,]), as.vector(REOF.proj.corr[2,,]))  ## second REOFs
    spat.cor[3,i] <- cor(as.vector(REOF.proj.corr.CPC[3,,]), as.vector(REOF.proj.corr[3,,]))  ## third REOFs
    spat.cor[4,i] <- cor(as.vector(REOF.proj.corr.CPC[4,,]), as.vector(REOF.proj.corr[4,,]))  ## fourth REOFs

    ## temporal correlations between PCs of REOF() and of CPC:
    temp.cor[1,i] <- round(cor(test$RPCs[,1], NAOy[which(NAOy[,2] == m), 3]), 2)
    temp.cor[2,i] <- round(cor(test$RPCs[,2], EAy[which(EAy[,2] == m), 3]), 2)
    temp.cor[3,i] <- round(cor(test$RPCs[,3], EAWRy[which(EAWRy[,2] == m), 3]), 2)
    temp.cor[4,i] <- round(cor(test$RPCs[,4], SCANDy[which(SCANDy[,2] == m), 3]), 2)

    ## explained variance:
    var.exp[,i] <- test$var[1:4]

    ## raw PCs:
    raw.REOFs[i + 0*n.rot,,] <- test2$REOFs[1,,]
    raw.REOFs[i + 1*n.rot,,] <- test2$REOFs[2,,]
    raw.REOFs[i + 2*n.rot,,] <- test2$REOFs[3,,]
    raw.REOFs[i + 3*n.rot,,] <- test2$REOFs[4,,]
    
    ## Normalize RPCs to plot trend:
    test3 <- test2
    for(d in 1:dim(test3$RPCs)[2]) test3$RPCs[,d] <- test2$RPCs[,d] / sd(test2$RPCs[,d])
       
    i <- i + 1
}

## Plot TPs:
fileoutput <- paste0(workdir,"/TPs_varimax_composite_proj_ERA-Interim.png")
png(filename=fileoutput, width=4500, height=1800) # width=15000, height=3700

plot.titles <- "" ##paste0("pattern #", c(rep(1,n.rot),rep(2,n.rot),rep(3,n.rot))) ##, "  (", round(var[1:4],1), "% tot.var)")
brks.corr <- seq(-1, 1, 0.1)
cols.corr <- colorRampPalette(c('#a50026','#d73027','#f46d43','#fdae61','#fee090','white','#e0f3f8','#abd9e9','#74add1','#4575b4','#313695'))(length(brks.corr)-1)

PlotLayout(fun = PlotEquiMap, plot_dims = c(2,3), var = proj.cor, lat = lat, lon = lon, filled.continents=FALSE, bar_scale=1, brks=brks.corr, cols=cols.corr, nrow=4, ncol=n.rot, titles=plot.titles, mar=c(2,2,2,2), coast_color="black", coast_width=2, axes_label_scale=2.5)

dev.off()

## Plot raw rotated EOFs:
fileoutput <- paste0(workdir,"/TPs_varimax_composite_EOFs_ERA-Interim.png")
png(filename=fileoutput, width=4500, height=1800) # width=15000, height=3700

plot.titles <- "" ##paste0("pattern #", c(rep(1,n.rot),rep(2,n.rot),rep(3,n.rot))) ##, "  (", round(var[1:4],1), "% tot.var)")
brks.corr <- seq(-1, 1, 0.1)
cols.corr <- colorRampPalette(c('#a50026','#d73027','#f46d43','#fdae61','#fee090','white','#e0f3f8','#abd9e9','#74add1','#4575b4','#313695'))(length(brks.corr)-1)

PlotLayout(fun = PlotEquiMap, plot_dims = c(2,3), var = raw.REOFs, lat = lat, lon = lon, filled.continents=F, bar_scale=1, cols=cols.corr, bar_label_scale=5)     ## Visualize the EOFs
dev.off()

## plot TPs tables:
library(gridExtra)
library(grid)
grid.table(as.data.frame(round(spat.cor,2)))
grid.table(as.data.frame(round(temp.cor,2)))
grid.table(as.data.frame(round(var.exp,1)))

## plot time series RPCs with trend:
for(PC in 1:4){
    fileoutput <- paste0(workdir,"/TPs_varimax_composite_freq_", PC,"_ERA-Interim.png")
    png(filename=fileoutput, width=400, height=300) # width=15000, height=3700

    plot(test3$RPCs[,PC], type="o", cex=.5, xaxt="n", ylim=c(-3.5, 3.5), xlab="Year", ylab="PC", cex.lab=1, cex.axis=1)
    axis(1, at=1:n.years, labels=year.start:year.end, cex.axis=1)
    lines(x=1:n.years, y=rep(0, n.years), col="gray")
    p.value <- MannKendall(test3$RPCs[,PC])
    x.data <- 1:n.years
    slope <- lm(test3$RPCs[,PC] ~ x.data)$coefficients[2]
    intercept <-  lm(test3$RPCs[,PC] ~ x.data)$coefficients[1]
    trend <- slope * n.years
    if(p.value[2] < 0.05) plot.lty <- 1 else plot.lty <- 2
    lines(x=1:n.years, y=slope*(0:(n.years-1)) + intercept, col="red", lty=plot.lty)

    dev.off()
}



