library(s2dverification)
library(plyr)
library(gridExtra)
library(grid)
library(abind)
source("https://earth.bsc.es/gitlab/es/s2dverification/raw/develop-teleconnections/R/REOF.R")
source("https://earth.bsc.es/gitlab/es/s2dverification/raw/develop-teleconnections/R/EOF.R")
source("https://earth.bsc.es/gitlab/es/s2dverification/raw/develop-teleconnections/R/EuroAtlanticTC.R")

#---------------------
# Choose the closest CPC mode with a space or time correlation
#---------------------
order_by_spatial_corr=T

#---------------------
# Read CPC indices to check signs and order of modes
#---------------------
TCI.CPC <- read.table("TCI_from_CPC.tab",skip=17,header=T,na.strings=c("-99.90"))
TCI.CPC <- TCI.CPC[TCI.CPC$yyyy %in% (1981:2016),c("yyyy","mm","NAO","EA","EAWR","SCA")]

#---------------------
# Read 500hPa Geopotential
#---------------------
data <- Load('g500', exp = NULL,
             obs = 'erainterim',
             sdates = paste0(1981:2016, '0101'),
             storefreq = 'monthly',
             leadtimemin = 1, leadtimemax = 12,
             latmin = 20, latmax = 80,
             lonmin = 270, lonmax = 60, output = 'lonlat')

g500 <- drop(data$obs)
names(dim(g500)) <- c("year","month","lat","lon")

#---------------------
# Compute anomalies
#---------------------
clim.g500 <- Mean1Dim(g500,1)
anom.g500 <- sweep(g500,c(2,3,4), clim.g500,"-")

#---------------------
# An auxiliar function for generating permutations
#---------------------
genperm <- function (n, r, v = 1:n) {
        if (r == 1) { matrix(v, n, 1) }
        else if (n == 1) { matrix(v, 1, r) }
        else {
                X <- NULL
                for (i in 1:n) X <- rbind(X, cbind(v[i], genperm(n - 1, r - 1, v[-i])))
                X
        }
}
pp <- genperm(4,4)

#---------------------
# Loop all months.
#---------------------
CPC_corr <- array(NA,dim=c(month=12,mode=4,lat=85,lon=214))
EATC_corr <- array(NA,dim=c(month=12,mode=4,lat=85,lon=214))
EATC_patterns <- array(NA,dim=c(month=12,mode=4,lat=85,lon=214))
EATC_indices <- array(NA,dim=c(month=12,year=36,mode=4))
sc.ls <- list()
tc.ls <- list()
for (m in 1:12) {
	#---------------------
	# Do a REOF for each month 
	# including also data from previous
	# and next month.
	#---------------------
	months <- (m-1 + (-1:1))%%12 + 1
	ano <- aperm(anom.g500[,months,,],c(2,1,3,4))
	dim(ano) <- c(dim(ano)[1]*dim(ano)[2],dim(ano)[3:4])
	telec <- EuroAtlanticTC(ano = ano, lon = data$lon, lat = data$lat, ntrunc = 4)

	EATC_patterns[m,,,] <- telec$patterns
	EATC_indices[m,,] <- telec$indices[rep(c(F,T,F),36),]
	# Get anomalies for the central month only
	anoc <- ano[rep(c(F,T,F),36),,]

	#---------------------
	# Compute CPC correlation patterns
	#---------------------
	CPC_corr[m,,,] <- laply(1:4,function(i) apply(anoc,c(2,3),function(x) cor(x,TCI.CPC[TCI.CPC$mm==m,c("NAO","EA","EAWR","SCA")[i]])))
	EATC_corr[m,,,] <- laply(1:4,function(i) apply(anoc,c(2,3),function(x) cor(x,EATC_indices[m,,i])))

	#---------------------
	# Compute spatial correlations
	#---------------------
	# TODO: latitudinal weights
	tmp.a <- t(array(CPC_corr[m,,,],dim=c(4,85*214)))
	tmp.b <- t(array(EATC_corr[m,,,],dim=c(4,85*214)))
	colnames(tmp.a) <- c("NAO","EA","EAWR","SCA")
	colnames(tmp.b) <- c("m1","m2","m3","m4")
	sc_mat <- cor(tmp.b,tmp.a)

	#---------------------
	# Compute time correlations
	#---------------------
	tc_mat <- cor(EATC_indices[m,,],TCI.CPC[TCI.CPC$mm==m,c("NAO","EA","EAWR","SCA")])
	rownames(tc_mat) <- c("m1","m2","m3","m4")

	#---------------------
	# Reorder modes and change sign
	# based on spatial or temporal correlations
	#---------------------
	if(order_by_spatial_corr) { cmat <- sc_mat }
	else { cmat <- tc_mat }
	bestperm <- pp[which.max(lapply(1:24,function(i) sum(diag(abs(cmat[pp[i,],]))))),]
	# manual tweaks!!!
	if(m==2) { bestperm <- c(1,2,4,3) } 
	#if(m==4) { bestperm <- c(1,3,4,2) }
	sign <- diag(sign(cmat[bestperm,]))

	EATC_indices[m,,] <- EATC_indices[m,,bestperm]*rep(sign,each=36)
	EATC_patterns[m,,,] <- EATC_patterns[m,bestperm,,]*sign
	EATC_corr[m,,,] <- EATC_corr[m,bestperm,,]*sign
	sc.ls[[m]] <- sc_mat[bestperm,]*sign
	tc.ls[[m]] <- tc_mat[bestperm,]*sign
}

PlotLayout(PlotEquiMap,c(3,4),aperm(EATC_patterns,c(2,1,3,4)),lon=data$lon,lat=data$lat,filled.continents=F,fileout="EATC_patterns.pdf",height=23,width=11,ncol=4,row_titles=month.abb[1:12],col_titles=c("NAO","EA","EA/WR","SCA"),bar_scale=0.3)
PlotLayout(PlotEquiMap,c(3,4),aperm(EATC_corr,c(2,1,3,4)),lon=data$lon,lat=data$lat,filled.continents=F,fileout="EATC_corr.pdf",height=23,width=11,ncol=4,row_titles=month.abb[1:12],col_titles=c("NAO","EA","EA/WR","SCA"),bar_scale=0.3,brks=seq(-1,1,0.1))
PlotLayout(PlotEquiMap,c(3,4),aperm(CPC_corr,c(2,1,3,4)),lon=data$lon,lat=data$lat,filled.continents=F,fileout="CPC_corr.pdf",height=23,width=11,ncol=4,row_titles=month.abb[1:12],col_titles=c("NAO","EA","EA/WR","SCA"),bar_scale=0.3,brks=seq(-1,1,0.1))

pdf("by_month.pdf",width=8,height=6)
for (m in 1:12) {
	PlotLayout(PlotEquiMap,c(3,4),abind(CPC_corr[m,,,,drop=F],EATC_corr[m,,,,drop=F],along=1),data$lon,data$lat,filled.continents=F,ncol=2,brks=seq(-1,1,0.1),row_titles=c("NAO","EA","EA/WR","SCA"),col_titles=c("CPC patterns","S2S4E patterns"),bar_scale=0.8,fileout=paste0(month.abb[[m]],".png"),width=7,height=7)
	
	pl <- grid::rasterGrob(png::readPNG(paste0(month.abb[[m]],".png")))
	t1 <- tableGrob(round(tc.ls[[m]],2))
	t2 <- tableGrob(round(sc.ls[[m]],2))
	title <- textGrob(month.name[[m]],gp=gpar(fontsize=30))
	title1 <- textGrob("Time correlations",gp=gpar(fontsize=20))
	title2 <- textGrob("Space correlations",gp=gpar(fontsize=20))
        grid.arrange(title,pl,title1,t1,title2,t2,layout_matrix=rbind(c(1,1),c(2,3),c(2,4),c(2,5),c(2,6)),widths=c(2,1),heights=c(1,0.5,3,0.5,3))
}
dev.off()
