
library(s2dverification)  # for the function Load()
library(abind)
source('/home/Earth/ncortesi/Downloads/scripts/Rfunctions.R') # for the calendar functions

# available reanalysis for the geopotential (g500) or the geopotential height (z500 = g500/9.81):
ERAint <- list(path = '/esnas/reconstructions/ecmwf/eraint/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc')
ERAint.name <- "ERA-Interim"
JRA55 <- list(path = '/esnas/reconstructions/jma/jra-55/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc')
JRA55.name <- "JRA-55"

workdir <- "/scratch/Earth/ncortesi/RESILIENCE/Regimes"    # working dir with the input files with the WT classifications for each MSLP grid point
mapdir <- "/scratch/Earth/ncortesi/RESILIENCE/Regimes_maps"  # output dir with seasonal and yearly maps

rean <- ERAint  # choose one of the two above reanalysis 
rean.name <- ERAint.name

psl <- "g500"                 # pressure variable to use from the chosen reanalysis
var.name <- "tas" #"sfcWind"  # name of the 'predictand' variable of the chosen reanalysis

year.start <- 1979
year.end <- 2013

period = 13 # (winter)  # 1-12: Jan-Dec; 13-16: Winter-Spring-Summer-Autumn; 17: Year

# Coordinates of the box enclosing the study region (the Northern Atlantic in this case):
lat.max <- 70
lat.min <- 30
lon.max <- 40
lon.min <- 280 # put a positive number here because the geopotential has only positive values of longitud!

#############################################################################################################################

# load only 1 day of geopot to detect the minimum and maximum lat and lon values to identify only the North Atlantic
domain <- Load(var = psl, exp = NULL, obs = list(rean), paste0(year.start,'0101'), storefreq = 'daily', leadtimemax = 1, output = 'lonlat', nprocs=1)
pos.lat.max <- which(abs(domain$lat-lat.max)==min(abs(domain$lat-lat.max)))
pos.lat.min <- which(abs(domain$lat-lat.min)==min(abs(domain$lat-lat.min)))
pos.lon.max <- which(abs(domain$lon-lon.max)==min(abs(domain$lon-lon.max)))
pos.lon.min <- which(abs(domain$lon-lon.min)==min(abs(domain$lon-lon.min)))
pos.lat <- if(pos.lat.min < pos.lat.max) {pos.lat.min:pos.lat.max} else { pos.lat.max:pos.lat.min}
pos.lon <- c(1:pos.lon.max,pos.lon.min:length(domain$lon)) # beware that it only consider the case where the study region cross the meridian of Greenwhich!
n.pos.lat <- length(pos.lat)
n.pos.lon <- length(pos.lon)
lat <- domain$lat[pos.lat] # lat of chosen area only (it is smaller than the whole spatial domain)
lon <- domain$lon[pos.lon] # lon of chosen area only

z500euFull <-array(NA,c(n.days.in.a.yearly.period(year.start,year.end), n.pos.lat, n.pos.lon))

## for (y in year.start:year.end){
##     var <- Load(var = 'z500', exp = NULL, obs = list(rean), paste0(y,'0101'), storefreq = 'daily', leadtimemax = n.days.in.a.year(y), output = 'lonlat')
##     z500eu <- var$obs[,,,,pos.lat,pos.lon]
##     z500euFull[seq.days.in.a.future.year(year.start, y),,] <- z500eu
##     rm(z500eu)
##     gc()
## }

if(domain$lat[pos.lat.min] >= domain$lat[pos.lat.max]) stop("lat.min cannot be higher than lat.max")

for (y in year.start:year.end){
    var <- Load(var = psl, exp = NULL, obs = list(rean), paste0(y,'0101'), storefreq = 'daily', leadtimemax = n.days.in.a.year(y), output = 'lonlat',
                latmin = domain$lat[pos.lat.min], latmax = domain$lat[pos.lat.max], lonmin = domain$lon[pos.lon.min], lonmax = domain$lon[pos.lon.max])
    z500euFull[seq.days.in.a.future.year(year.start, y),,] <- var$obs
    rm(var)
    gc()
}


# each time you want to change variable 'period' and/or 'var', set period to 13 and load the winter dataset where all input data is stored:
load(file=paste0(workdir,"/weather_regimes_",rean.name,"_",var.name,"_",year.start,"-",year.end,"_", my.period[period],".RData"))

days.period <- NA
for(y in year.start:year.end) days.period <- c(days.period, n.days.in.a.future.year(year.start, y) + pos.period(y,period))
days.period <- days.period[-1] # remove the NA at the begin that was introduced only to be able to execture the above command
n.days.period <- length(days.period)

z500 <- z500euFull[days.period,,] # select only days in the chosen period (i.e: winter)

z500mat <- z500

dim(z500mat) <- c(n.days.period, n.pos.lat*n.pos.lon) # convert array in a matrix!

my.seq <- seq(1, n.pos.lat*n.pos.lon, 9) # select only 1 point of 9

z500cut <- z500mat[,my.seq]

my.PCA <- princomp(z500cut,cor=FALSE)

head(cumsum(my.PCA$sdev^2/sum(my.PCA$sdev^2)),20)  # check how many PCAs to keep basing on the sum of their explained variance

my.cluster <- kmeans(my.PCA$scores[,1:8], 4)  # 4 is the number of clusters, 7 the number of EOFs which explains ~80% of variance

rm(z500mat)
gc()

# Load wind data:

vareuFull <-array(NA,c(n.days.in.a.yearly.period(year.start,year.end),n.pos.lat,n.pos.lon))
    
## for (y in year.start:year.end){
##     var <- Load(var = var.name, exp = NULL, obs = list(rean), paste0(y,'0101'), storefreq = 'daily', leadtimemax = n.days.in.a.year(y), output = 'lonlat')
##     vareu <- var$obs[,,,,pos.lat,pos.lon]
##     vareuFull[seq.days.in.a.future.year(year.start, y),,] <- vareu
## }

for (y in year.start:year.end){
    var <- Load(var = var.name, exp = NULL, obs = list(rean), paste0(y,'0101'), storefreq = 'daily', leadtimemax = n.days.in.a.year(y), output = 'lonlat',
           latmin = domain$lat[pos.lat.min], latmax = domain$lat[pos.lat.max], lonmin = domain$lon[pos.lon.min], lonmax = domain$lon[pos.lon.max])
    vareuFull[seq.days.in.a.future.year(year.start, y),,] <- var$obs
    rm(var)
    gc()
}

save.image(file=paste0(workdir,"/weather_regimes_",rean.name,"_",var.name,"_",year.start,"-",year.end,"_", my.period[period],".RData"))

varPeriod <- vareuFull[days.period,,] # select only var data during the chosen period

varPeriodClim <- apply(varPeriod,c(2,3),mean,na.rm=T)
varPeriodClim2 <- InsertDim(varPeriodClim, 1, n.days.period)

varPeriodAnom <- varPeriod - varPeriodClim2
rm(varPeriod, varPeriodClim2)
gc()

wr1 <- which(my.cluster$cluster==1)
wr2 <- which(my.cluster$cluster==2)
wr3 <- which(my.cluster$cluster==3)
wr4 <- which(my.cluster$cluster==4)

varPeriodAnom1 <- varPeriodAnom[wr1,,]
varPeriodAnom2 <- varPeriodAnom[wr2,,]
varPeriodAnom3 <- varPeriodAnom[wr3,,]
varPeriodAnom4 <- varPeriodAnom[wr4,,]

varPeriodAnom1mean <- apply(varPeriodAnom1,c(2,3),mean,na.rm=T)
varPeriodAnom2mean <- apply(varPeriodAnom2,c(2,3),mean,na.rm=T)
varPeriodAnom3mean <- apply(varPeriodAnom3,c(2,3),mean,na.rm=T)
varPeriodAnom4mean <- apply(varPeriodAnom4,c(2,3),mean,na.rm=T)

varPeriodAnomBoth1 <- abind(varPeriodAnom, varPeriodAnom1, along = 1)
varPeriodAnomBoth2 <- abind(varPeriodAnom, varPeriodAnom2, along = 1)
varPeriodAnomBoth3 <- abind(varPeriodAnom, varPeriodAnom3, along = 1)
varPeriodAnomBoth4 <- abind(varPeriodAnom, varPeriodAnom4, along = 1)

pvalue1 <- apply(varPeriodAnomBoth1, c(2,3), function(x) t.test(x[1:n.days.period],x[(n.days.period+1):length(x)])$p.value)
pvalue2 <- apply(varPeriodAnomBoth2, c(2,3), function(x) t.test(x[1:n.days.period],x[(n.days.period+1):length(x)])$p.value)
pvalue3 <- apply(varPeriodAnomBoth3, c(2,3), function(x) t.test(x[1:n.days.period],x[(n.days.period+1):length(x)])$p.value)
pvalue4 <- apply(varPeriodAnomBoth4, c(2,3), function(x) t.test(x[1:n.days.period],x[(n.days.period+1):length(x)])$p.value)
rm(varPeriodAnom, varPeriodAnomBoth1, varPeriodAnomBoth2, varPeriodAnomBoth3, varPeriodAnomBoth4)
gc()

z500wr1 <- z500[wr1,,]
z500wr2 <- z500[wr2,,]
z500wr3 <- z500[wr3,,]
z500wr4 <- z500[wr4,,]

z500wr1mean <- apply(z500wr1,c(2,3),mean,na.rm=T)
z500wr2mean <- apply(z500wr2,c(2,3),mean,na.rm=T)
z500wr3mean <- apply(z500wr3,c(2,3),mean,na.rm=T)
z500wr4mean <- apply(z500wr4,c(2,3),mean,na.rm=T)

# Mean z500 maps:

my.brks <- c(48000, seq(48501,57100,1), 60000) # % Mean anomaly of a WR
my.brks2 <- c(48000, seq(48500,58000,500)) # % Mean anomaly of a WR
my.cols <-  colorRampPalette((brewer.pal(9,"PuBu")))(length(my.brks)-1)    # blue--white--red colors

# First check each cluster to find to which regime it corrsponds:
PlotEquiMap(z500wr1mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, contours=t(z500wr1mean), brks2=my.brks2) 
PlotEquiMap(z500wr2mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, contours=t(z500wr2mean), brks2=my.brks2) 
PlotEquiMap(z500wr3mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, contours=t(z500wr3mean), brks2=my.brks2) 
PlotEquiMap(z500wr4mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, contours=t(z500wr4mean), brks2=my.brks2) 

# Assign to each cluster the name of its regime:
regime1.name <- "Cluster #1"
regime2.name <- "Cluster #2"
regime3.name <- "Cluster #3"
regime4.name <- "Cluster #4"

# Visualize and save average geopotential associated to each regime:
png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime1.name,".png"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(z500wr1mean, lon, lat, filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle=regime1.name,  contours=t(z500wr1mean), brks2=my.brks2, drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime2.name,".png"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(z500wr2mean, lon, lat, filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle=regime2.name,  contours=t(z500wr2mean), brks2=my.brks2, drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime3.name,".png"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(z500wr3mean, lon, lat, filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle=regime3.name,  contours=t(z500wr3mean), brks2=my.brks2, drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime4.name,".png"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(z500wr4mean, lon, lat, filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle=regime4.name,  contours=t(z500wr4mean), brks2=my.brks2, drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
dev.off()


# Visualize and save average wind anomalies associated to each regime:
my.brks.var <- c(seq(-10,-3,1),seq(-2.9,2.9,0.1),seq(3,10,1)) # % Mean anomaly of a WR
#my.cols <-  colorRampPalette(as.character(read.csv("/home/Earth/vtorralb/rgbhex.csv",header=F)[,1]))(length(my.brks)-1)  # blue--yellow-red colors
my.cols.var <-  colorRampPalette(rev(brewer.pal(11,"RdBu")))(length(my.brks.var)-1)    # blue--white--red colors

png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime1.name,"_anomalies.png"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom1mean, lon, lat, filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=paste(regime1.name, "Anomalies"), drawleg=F, dots=t(pvalue1 < 0.05))
ColorBar(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_", my.period[period],"_",regime2.name,"_anomalies.png"), width=1000, height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom2mean,lon,lat,filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=paste(regime2.name, "Anomalies"), drawleg=F, dots=t(pvalue2 < 0.05))
ColorBar(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_", my.period[period],"_",regime3.name,"_anomalies.png"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom3mean,lon,lat,filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=paste(regime3.name, "Anomalies"), drawleg=F, dots=t(pvalue3 < 0.05))
ColorBar(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_", my.period[period],"_",regime4.name,"_anomalies.png"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom4mean,lon,lat,filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=paste(regime4.name, "Anomalies"), drawleg=F, dots=t(pvalue4 < 0.05))
ColorBar(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()





# save as .pdf instead of .png:
pdf(file=paste0(mapdir,"/",rean,"_",var.name,"_", my.period[period],"NAO-.pdf"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom1mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle="NAO-",drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

pdf(file=paste0(mapdir,"/",rean,"_",var.name,"_", my.period[period],"Blocking.pdf"), width=1000, height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom2mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle="Blocking", drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

pdf(file=paste0(mapdir,"/",rean,"_",var.name,"_", my.period[period],"Atlantic_Ridge.pdf"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom3mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle="Atlantic Ridge", drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

pdf(file=paste0(mapdir,"/",rean,"_",var.name,"_", my.period[period],"NAO+.pdf"),width=1000,height=700)
layout(matrix(c(rep(1,9),2), 10, 1, byrow = TRUE))
PlotEquiMap(varPeriodAnom4mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle="NAO+", drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

rm(varPeriodeuFull,varPeriodeu,varPeriod,var)
