
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 from where to load the input psl and var data
rean.name <- ERAint.name

psl <- "g500"                 # pressure variable to use from the chosen reanalysis

var.name <- "tas" #"sfcWind" #"tas"                              # name of the 'predictand' variable of the chosen reanalysis
var.name.full <- "Temperature" #"Wind Speed" #"Temperature"         # full name of the 'predictand' variable to put in the title of the graphs
var.unit <- "ºC" #"m/s"                                  # unit of measure (for drawing color scales)

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

variance.explained <- 0.8  # the user can set here the minimum % of variance explained by the PCs (typically 0.8 which is equal to 80%)

# 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
if(domain$lat[pos.lat.min] >= domain$lat[pos.lat.max]) stop("lat.min cannot be higher than lat.max")

# Load psl data:
psleuFull <-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 = 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])
    psleuFull[seq.days.in.a.future.year(year.start, y),,] <- var$obs
    rm(var)
    gc()
}

if(psl=="g500") psleuFull <- psleuFull/9.81   # convert geopotential to geopotential height
    
# Load var 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',
           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,".RData"))

# each time you want to change only variable 'period', start from here loading the saved data:
load(file=paste0(workdir,"/weather_regimes_",rean.name,"_",var.name,"_",year.start,"-",year.end,".RData"))

for(period in 13:16){ # 1-12: Jan-Dec; 13-16: Winter-Spring-Summer-Autumn; 17: Year

    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 execure the above command
    n.days.period <- length(days.period)

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

    pslmat <- pslPeriod

    dim(pslmat) <- 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

    pslcut <- pslmat[,my.seq]

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

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

    n.pcs <- head(as.numeric(which(tot.variance > variance.explained)),1) # select only the pcs that explains at least 80% of variance

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

    rm(pslmat)
    gc()

    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)

    period.length <- n.days.in.a.period(period,1999)  # using year 1999 introduce a small error in winter season because of bisestile years

    wr1y <- wr2y <- wr3y <-wr4y <- c()
    for(y in year.start:year.end){
        wr1y[y-year.start+1] <- length(which(my.cluster$cluster[(y-year.start)*period.length+(1:period.length)] == 1))
        wr2y[y-year.start+1] <- length(which(my.cluster$cluster[(y-year.start)*period.length+(1:period.length)] == 2))
        wr3y[y-year.start+1] <- length(which(my.cluster$cluster[(y-year.start)*period.length+(1:period.length)] == 3))
        wr4y[y-year.start+1] <- length(which(my.cluster$cluster[(y-year.start)*period.length+(1:period.length)] == 4))
    }

    # convert to frequencies in %:
    wr1y <- wr1y/period.length
    wr2y <- wr2y/period.length
    wr3y <- wr3y/period.length
    wr4y <- wr4y/period.length

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

    pslPeriodClim <- apply(pslPeriod, c(2,3), mean, na.rm=T)
    pslPeriodClim2 <- InsertDim(pslPeriodClim, 1, n.days.period)

    pslPeriodAnom <- pslPeriod - pslPeriodClim2
    rm(pslPeriod, pslPeriodClim, pslPeriodClim2)
    gc()

    pslwr1 <- pslPeriodAnom[wr1,,]
    pslwr2 <- pslPeriodAnom[wr2,,]
    pslwr3 <- pslPeriodAnom[wr3,,]
    pslwr4 <- pslPeriodAnom[wr4,,]

    pslwr1mean <- apply(pslwr1,c(2,3),mean,na.rm=T)
    pslwr2mean <- apply(pslwr2,c(2,3),mean,na.rm=T)
    pslwr3mean <- apply(pslwr3,c(2,3),mean,na.rm=T)
    pslwr4mean <- apply(pslwr4,c(2,3),mean,na.rm=T)

    rm(pslwr1,pslwr2,pslwr3,pslwr4)
    
    # breaks and colors of the geopotential fields:
    #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.brks <- c(-300,-199:200,300) # % Mean anomaly of a WR
    my.brks2 <- c(-300,seq(-200,200,10),300) # % Mean anomaly of a WR
    #my.cols <- colorRampPalette((brewer.pal(9,"PuBu")))(length(my.brks)-1)  
    my.cols <- colorRampPalette(rev(brewer.pal(11,"RdBu")))(length(my.brks)-1)    # blue--white--red colors

    # breaks and colors of the impact maps:
    my.brks.var <- c(-20,seq(-10,-3,1),seq(-2.9,2.9,0.1),seq(3,10,1),20) # % Mean anomaly of a WR
    #my.brks.var <- c(-20,seq(-10,-3,1),seq(-2.8,2.8,0.2),seq(3,10,1),20) # % 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

    regime1.name <- "Cluster #1"
    regime2.name <- "Cluster #2"
    regime3.name <- "Cluster #3"
    regime4.name <- "Cluster #4"

    # when you want to add the right regime name to the clusters:    
    period=13
    load(file=paste0(workdir,"/",rean.name,"_",var.name,"_",my.period[period],"_mapdata.RData"))
    if(period==13){
    regime1.name="Atlantic Ridge"
    regime2.name="Blocking"
    regime3.name="NAO+"
    regime4.name="NAO-"
    }
    if(period==14){
    regime1.name="Blocking"
    regime2.name="NAO-"
    regime3.name="Atlantic Ridge"
    regime4.name="NAO+"
    if(period==15){
    regime1.name="NAO+"
    regime2.name="Blocking"
    regime3.name="Atlantic Ridge"
    regime4.name="NAO-"
    }
    if(period==16){
    regime1.name="NAO-"
    regime2.name="Blocking"
    regime3.name="NAO+"
    regime4.name="Atlantic Ridge"  
    }
    
    # Visualize the composition with the 3 graphs for each single regime: its pressure fields, its impact on var and its frequency serie:
    png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime1.name,".png"),width=1600,height=600)
    layout(matrix(c(rep(1,9),2,rep(3,9),4,rep(5,10),6,rep(7,8),8), 10, 4, byrow = FALSE), widths=c(2,2,0.2,2))
    #layout.show(8)
    EU <- c(1:54,130:161) # position of long values of Europe only (without the Atlantic Sea and America)
    title1 <- paste(regime1.name, "geopotential height")
    PlotEquiMap(pslwr1mean[,EU], lon[EU], lat, filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle=title1,  contours=t(pslwr1mean[,EU]), brks2=my.brks2, drawleg=F)
    ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
    mtext(side=4," m", cex=1.5)
    title1 <- paste(regime1.name, "impact on", var.name.full)
    PlotEquiMap(varPeriodAnom1mean[,EU], lon[EU], lat, filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=title1, drawleg=F, dots=t(pvalue1[,EU] < 0.05))
    ColorBar2(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=10, cex=1.2, my.ticks=c(0,10,20,30,40,50,60,70,80), my.labels=c(-20,-10,-6,-3,0,3,6,10,20))
    mtext(side=4,paste0(" ",var.unit), cex=1.5)
    plot(0, type="n", axes=F, xlab="", ylab="")
    plot(0, type="n", axes=F, xlab="", ylab="")  # plot 2 empty graphs
    title(paste0(regime1.name, " Frequency (", round(100*mean(wr1y),1), "%)"), cex.main=1.7, line=-1)
    barplot.freq(100*wr1y, year.start, year.end, cex.y=1.1, cex.x=1.1, freq.max=80)
    plot(0, type="n", axes=F,  xlab="", ylab="")
    title("Year", cex.main=1.3, line=-1)
    dev.off()

    png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime2.name,".png"),width=1600,height=600)
    layout(matrix(c(rep(1,9),2,rep(3,9),4,rep(5,10),6,rep(7,8),8), 10, 4, byrow = FALSE), widths=c(2,2,0.2,2))
    #layout.show(8)
    EU <- c(1:54,130:161) # position of long values of Europe
    title2 <- paste(regime2.name, "geopotential height")
    PlotEquiMap(pslwr2mean[,EU], lon[EU], lat, filled.continents=FALSE, brks=my.brks, cols=my.cols, toptitle=title2,  contours=t(pslwr2mean[,EU]), brks2=my.brks2, drawleg=F)
    ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
    mtext(side=4," m", cex=1.5)
    title2 <- paste(regime2.name, "impact on", var.name.full)
    PlotEquiMap(varPeriodAnom2mean[,EU], lon[EU], lat, filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=title2, drawleg=F, dots=t(pvalue2[,EU] < 0.05))
    ColorBar2(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=10, cex=1.2, my.ticks=c(0,10,20,30,40,50,60,70,80), my.labels=c(-20,-10,-6,-3,0,3,6,10,20))
    mtext(side=4,paste0(" ",var.unit), cex=1.5)
    plot(0, type="n", axes=F, xlab="", ylab="")
    plot(0, type="n", axes=F, xlab="", ylab="")
    title(paste0(regime2.name, " Frequency (", round(100*mean(wr2y),1), "%)"), cex.main=1.7, line=-1)
    barplot.freq(100*wr2y, year.start, year.end, cex.y=1.1, cex.x=1.1, freq.max=80)
    plot(0, type="n", axes=F,  xlab="", ylab="")
    title("Year", cex.main=1.3, line=-1)
    dev.off()

    png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime3.name,".png"),width=1600,height=600)
    layout(matrix(c(rep(1,9),2,rep(3,9),4,rep(5,10),6,rep(7,8),8), 10, 4, byrow = FALSE), widths=c(2,2,0.2,2))
    #layout.show(7)
    EU <- c(1:54,130:161) # position of long values of Europe
    title3 <- paste(regime3.name, "geopotential height")
    PlotEquiMap(pslwr3mean[,EU], lon[EU], lat, filled.continents = FALSE, brks=my.brks, cols=my.cols, toptitle=title3,  contours=t(pslwr3mean[,EU]), brks2=my.brks2, drawleg=F)
    ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
    mtext(side=4," m", cex=1.5)
    title3 <- paste(regime3.name, "impact on", var.name.full)
    PlotEquiMap(varPeriodAnom3mean[,EU], lon[EU], lat, filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=title3, drawleg=F, dots=t(pvalue3[,EU] < 0.05))
    ColorBar2(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=10, cex=1.2, my.ticks=c(0,10,20,30,40,50,60,70,80), my.labels=c(-20,-10,-6,-3,0,3,6,10,20))
    mtext(side=4,paste0(" ",var.unit), cex=1.5)
    plot(0, type="n", axes=F, xlab="", ylab="")
    plot(0, type="n", axes=F,  xlab="", ylab="")
    title(paste0(regime3.name, " Frequency (", round(100*mean(wr3y),1), "%)"), cex.main=1.7, line=-1)
    barplot.freq(100*wr3y, year.start, year.end, cex.y=1.1, cex.x=1.1, freq.max=80)
    plot(0, type="n", axes=F,  xlab="", ylab="")
    title("Year", cex.main=1.3, line=-1)
    dev.off()

    png(filename=paste0(mapdir,"/",rean.name,"_",var.name,"_",my.period[period],"_",regime4.name,".png"),width=1600,height=600)
    layout(matrix(c(rep(1,9),2,rep(3,9),4,rep(5,10),6,rep(7,8),8), 10, 4, byrow = FALSE), widths=c(2,2,0.2,2))
    #layout.show(8)
    EU <- c(1:54,130:161) # position of long values of Europe
    title4 <- paste(regime4.name, "geopotential height")
    PlotEquiMap(pslwr4mean[,EU], lon[EU], lat, filled.continents=FALSE, brks=my.brks, cols=my.cols, toptitle=title4, contours=t(pslwr4mean[,EU]), brks2=my.brks2, drawleg=F)
    ColorBar(my.brks, cols=my.cols, vert=FALSE, subsampleg=100, cex=1.2)
    mtext(side=4," m", cex=1.5)
    title4 <- paste(regime4.name, "impact on", var.name.full)
    PlotEquiMap(varPeriodAnom4mean[,EU], lon[EU], lat, filled.continents = FALSE, brks=my.brks.var, cols=my.cols.var, toptitle=title4, drawleg=F, dots=t(pvalue4[,EU] < 0.05))
    ColorBar2(my.brks.var, cols=my.cols.var, vert=FALSE, subsampleg=10, cex=1.2, my.ticks=c(0,10,20,30,40,50,60,70,80), my.labels=c(-20,-10,-6,-3,0,3,6,10,20))
    mtext(side=4,paste0(" ",var.unit), cex=1.5)
    plot(0, type="n", axes=F, xlab="", ylab="")
    plot(0, type="n", axes=F, xlab="", ylab="")
    title(paste0(regime4.name, " Frequency (", round(100*mean(wr4y),1), "%)"), cex.main=1.7, line=-1)
    barplot.freq(100*wr4y, year.start, year.end, cex.y=1.1, cex.x=1.1, freq.max=80)
    plot(0, type="n", axes=F,  xlab="", ylab="")
    title("Year", cex.main=1.3, line=-1)
    dev.off()

    # save all the data necessary to redraw the graphs when we know the right regime:
    save(workdir,rean.name,var.name,var.name.full,var.unit,year.start,year.end,period,my.period,lon,lat,my.brks,my.brks2,my.brks.var,my.cols,my.cols.var,pslwr1mean,pslwr2mean,pslwr3mean,pslwr4mean, varPeriodAnom1mean, varPeriodAnom2mean, varPeriodAnom3mean,varPeriodAnom4mean,wr1y,wr2y,wr3y,wr4y,pvalue1,pvalue2,pvalue3,pvalue4,regime1.name,regime2.name,regime3.name,regime4.name,file=paste0(workdir,"/",rean.name,"_",var.name,"_",my.period[period],"_mapdata.RData"))
    
    rm(pvalue1,pvalue2,pvalue3,pvalue4)
    rm(pslwr1mean,pslwr2mean,pslwr3mean,pslwr4mean)
    rm(varPeriodAnom1mean,varPeriodAnom2mean,varPeriodAnom3mean,varPeriodAnom4mean)
    gc()    
    
} # close for loop on period

