
library(s2dverification)  # for the function Load()
source('/home/Earth/ncortesi/Downloads/scripts/Rfunctions.R') # for the calendar functions
# reanalysis available:
ERAint <- list(path = '/esnas/reconstructions/ecmwf/eraint/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc')
JRA55 <- list(path = '/esnas/reconstructions/jma/jra-55/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc')

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

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

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

var.name <- "tas" #"sfcWind"  # name of the 'predictand' variable of the chosen reanalysis

# 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 Z500 has only positive values of longitud!

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

# load only 1 day of z500 to detect the minimum and maximum lat and lon values to identify only the North Atlantic
domain <- Load(var = 'z500', 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)

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()
}

lat <- var$lat[pos.lat]
lon <- var$lon[pos.lon]

period = 13 # (winter)

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(head(dim(z500mat),1), n.pos.lat*172) # convert array in a matrix!

my.seq <- seq(1,9976,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:7], 4)  # 4 is the number of clusters, 7 the number of EOFs which explains in our case the 79% of variance

rm(z500euFull,z500eu,z500mat,var)
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
}

save.image(file=paste0(workdir,"/weather_regimes_",rean,"_",var.name,".RData"))
load(file=paste0(workdir,"/weather_regimes_",rean,"_",var.name,".RData"))

varPeriod <- vareuFull[days.period,,]

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

varPeriodAnom <- varPeriod - varPeriodClim2

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)

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(seq(48800,57100,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

PlotEquiMap(z500wr1mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols)  # NAO-
PlotEquiMap(z500wr2mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols)  # blocking
PlotEquiMap(z500wr3mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols)  # atlantic ridge
PlotEquiMap(z500wr4mean,lon,lat,filled.continents = FALSE, brks=my.brks, cols=my.cols)  # NAO+


# Average wind anomalies:

my.brks <- 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

png(filename=paste0(mapdir,"/",rean,"_",var.name,"_",my.period[period],"NAO-.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, cols=my.cols, toptitle="NAO-",drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean,"_",var.name,"_", my.period[period],"Blocking.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, cols=my.cols, toptitle="Blocking", drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean,"_",var.name,"_", my.period[period],"Atlantic Ridge.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, cols=my.cols, toptitle="Atlantic Ridge", drawleg=F)
ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=1, cex=1.2)
dev.off()

png(filename=paste0(mapdir,"/",rean,"_",var.name,"_", my.period[period],"NAO+.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, cols=my.cols, toptitle="NAO+", drawleg=F)
ColorBar(my.brks, cols=my.cols, 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)
