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
scale_indices=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) }
	if(m==12) {bestperm <- c(1,3,4,2) }
	sign <- diag(sign(cmat[bestperm,]))
	if(m==4) { sign[2] <- sign[2]*-1 }

	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
}

#---------------------
# Plot all months together (sabana)
#---------------------
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))

#---------------------
# One page per month, detailed comparison
#---------------------
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()


#---------------------
# Save to netcdf
#---------------------
lats <- data$lat
lons <- data$lon
lons1<-c(lons[which(lons>180)]-360,lons[which(lons<180)])

for (m in 1:12) {
	Pattern1 <- t(EATC_patterns[m,1,,c(which(lons>180),which(lons<180))])
	Pattern2 <- t(EATC_patterns[m,2,,c(which(lons>180),which(lons<180))])
	Pattern3 <- t(EATC_patterns[m,3,,c(which(lons>180),which(lons<180))])
	Pattern4 <- t(EATC_patterns[m,4,,c(which(lons>180),which(lons<180))])

	if(scale_indices) {
		Index1 <- scale(EATC_indices[m,,1])
		Index2 <- scale(EATC_indices[m,,2])
		Index3 <- scale(EATC_indices[m,,3])
		Index4 <- scale(EATC_indices[m,,4])
	} else { 	
		Index1 <- EATC_indices[m,,1]
		Index2 <- EATC_indices[m,,2]
		Index3 <- EATC_indices[m,,3]
		Index4 <- EATC_indices[m,,4]
	}

	metadata <- list(Pattern1 = list(units=''))
	attr(Pattern1,'variables') <- metadata
	names(dim(Pattern1)) <- c('longitude',  'latitude')
	metadata <- list(Pattern2 = list(units=''))
	attr(Pattern2,'variables') <- metadata
	names(dim(Pattern2)) <- c('longitude',  'latitude')
	metadata <- list(Pattern3 = list(units=''))
	attr(Pattern3,'variables') <- metadata
	names(dim(Pattern3)) <- c('longitude',  'latitude')
	metadata <- list(Pattern4 = list(units=''))
	attr(Pattern4,'variables') <- metadata
	names(dim(Pattern4)) <- c('longitude',  'latitude')

 	dim(Index1) <- length(Index1)
	metadata <- list(Index1 = list(units=''))
	attr(Index1,'variables') <- metadata
	names(dim(Index1)) <- c('time')
 	dim(Index2) <- length(Index2)
	metadata <- list(Index2 = list(units=''))
	attr(Index2,'variables') <- metadata
	names(dim(Index2)) <- c('time')
 	dim(Index3) <- length(Index3)
	metadata <- list(Index3 = list(units=''))
	attr(Index3,'variables') <- metadata
	names(dim(Index3)) <- c('time')
 	dim(Index4) <- length(Index4)
	metadata <- list(Index4 = list(units=''))
	attr(Index4,'variables') <- metadata
	names(dim(Index4)) <- c('time')

	latitude <- lats
	dim(latitude) <- length(lats)
	metadata <- list(latitude = list(units = 'degrees_north'))
	attr(latitude, 'variables') <- metadata
	names(dim(latitude)) <- 'latitude'

	longitude <- lons1
	dim(longitude) <- length(lons1)
	metadata <- list(longitude = list(units = 'degrees_east'))
	attr(longitude, 'variables') <- metadata
	names(dim(longitude)) <- 'longitude'

	time <- 1981:2016
	dim(time) <- length(time)
	metadata <- list(time=list(units = 'years'))
	attr(time, 'variables') <- metadata
	names(dim(time)) <- 'Year'

	ArrayToNetCDF(list(Pattern1,Pattern2,Pattern3,Pattern4,Index1,Index2,Index3,Index4,latitude,longitude,time), paste0('/esarchive/scratch/llledo/S2S4E/EATC/ERA_Interim_EATC_',month.name[m],'_1981_2016.nc'))
}

# Load monthly sfcWind and CFs 
for (m in 1:12) {
	cf1 <- array(NA,dim=c(year=36,longitude=512,latitude=256))
	cf2 <- array(NA,dim=c(year=36,longitude=512,latitude=256))
	cf3 <- array(NA,dim=c(year=36,longitude=512,latitude=256))
	i <- 1
	for (y in 1981:2016) {
		cf1[i,,] <- Mean1Dim(readRDS(paste0("/esarchive/scratch/llledo/daily_CF_ERAI/ERAI_CF_iec1_",y,sprintf("%02d",m),".RDS")),1)
		cf2[i,,] <- Mean1Dim(readRDS(paste0("/esarchive/scratch/llledo/daily_CF_ERAI/ERAI_CF_iec2_",y,sprintf("%02d",m),".RDS")),1)
		cf3[i,,] <- Mean1Dim(readRDS(paste0("/esarchive/scratch/llledo/daily_CF_ERAI/ERAI_CF_iec3_",y,sprintf("%02d",m),".RDS")),1)
		i <- i+1
	}

	sfcwind <- Load('sfcWind', exp = NULL,
		obs = 'erainterim',
		sdates = paste0(1981:2016, sprintf("%02d",m),"01"),
		storefreq = 'monthly',
		leadtimemin = 1, leadtimemax = 1,
		output = 'lonlat')
	sfcWind <- aperm(drop(sfcwind$obs),c(1,3,2))
	data.ls <- list()
	data.ls[[1]] <- sfcWind
	data.ls[[2]] <- cf1
	data.ls[[3]] <- cf2
	data.ls[[4]] <- cf3

	results <- lapply(1:4,function(i) {
		lapply(1:4,function(j) { 
			apply(data.ls[[j]],c(2,3),function(x) cor(x,EATC_indices[m,,i]))
		})
	})

	lats <- sfcwind$lat
	lons <- sfcwind$lon
	lons2 <- ifelse(lons<180,lons,lons-360)
	idxlat <- (lats<80 & lats>20)
	idxlon <- (lons2<70 & lons2>-60)
	idxlat2 <- (lats<72 & lats>32)
	idxlon2 <- (lons2<40 & lons2>-18)
	PlotLayout(PlotEquiMap,c(1,2),unlist(results,recursive=F),lons2,lats,filled.continents=F,brks=seq(-1,1,0.1),toptitle=paste0("Impact of EuroAtlantic Teleconnections on wind and capacity factor\n in ",month.name[m]," 1981-2016 (ERA-Interim)"),row_titles=c("NAO","EA","EA/WR","SCA"),col_titles=c("sfcWind","CF iec1","CF iec2","CF iec3"),title_scale=0.7,bar_scale=0.5,title_margin_scale=3,fileout=paste0("EATC_impact_wind_and_CF_",month.name[m],"_global.pdf"),width=11,height=12)
	PlotLayout(PlotEquiMap,c(2,3),abind(unlist(results,recursive=F),along=0)[,idxlon,idxlat],lons2[idxlon],lats[idxlat],filled.continents=F,brks=seq(-1,1,0.1),toptitle=paste0("Impact of EuroAtlantic Teleconnections on wind and capacity factor\n in ",month.name[m]," 1981-2016 (ERA-Interim)"),row_titles=c("NAO","EA","EA/WR","SCA"),col_titles=c("sfcWind","CF iec1","CF iec2","CF iec3"),title_scale=0.7,bar_scale=0.5,title_margin_scale=3,fileout=paste0("EATC_impact_wind_and_CF_",month.name[m],".pdf"),width=11,height=12)
	PlotLayout(PlotEquiMap,c(2,3),abind(unlist(results,recursive=F),along=0)[,idxlon2,idxlat2],lons2[idxlon2],lats[idxlat2],filled.continents=F,brks=seq(-1,1,0.1),toptitle=paste0("Impact of EuroAtlantic Teleconnections on wind and capacity factor\n in ",month.name[m]," 1981-2016 (ERA-Interim)"),row_titles=c("NAO","EA","EA/WR","SCA"),col_titles=c("sfcWind","CF iec1","CF iec2","CF iec3"),title_scale=0.7,bar_scale=0.5,title_margin_scale=3,fileout=paste0("EATC_impact_wind_and_CF_",month.name[m],"_zoom.pdf"),width=11,height=12)
}
