
library(s2dverification)
library(multiApply)
library(abind)
library(future) # for availableCores() function
source('/esarchive/scratch/ncortesi/scripts/Rfunctions.R')
source('/esarchive/scratch/ncortesi/scripts/Vero/VeroRegimes.R')

work.dir <- "/esarchive/scratch/ncortesi/Regimes/87_WRs_subseasonal"

#####################################################################################################
#   WR classification for extended winter (October to March) and summer (Abril to September)        #
#####################################################################################################

##
## like the WR vignette but without the forecast data (faster to load),
## and loading all year instead than one month only and with LOESS filtering to measure clim instead of simple mean:
##

year.start <- 1979
year.end <- 2017

work.dir <- paste0("/esarchive/scratch/ncortesi/Regimes/90_WRs_winter_and_summer/validation_", year.end + 1)

sdates <- paste0(year.start:year.end, '0101')

data366 <- Load('psl', exp = NULL,
             obs = 'erainterim',
             sdates = sdates,
             storefreq='daily',
             leadtimemin = 1, leadtimemax = 366,
             latmin = 27, latmax = 81,
             lonmin = 274.5, lonmax = 45, output = 'lonlat', nprocs=1)

## Climatology with simple mean:
## clim <- Clim(var_exp = data$obs, var_obs = data$obs, memb=F)
## clim_smoothed_obs <- aperm(apply(clim$clim_obs,c(1:length(dim(clim$clim_obs)))[-which(names(dim(clim$clim_obs))=='ftime')], Loess, loess_span=1),c(2,1,3,4))
## anom_obs <- Ano(data$obs, clim_smoothed_obs)

## WR_obs <- WeatherRegime(data = anom_obs, EOFS = FALSE, lat = data$lat, ncenters = 4)
## names(dim(WR_obs$composite)) <- c("lat", "lon", "cluster", "", "")

lat <- data366$lat
lon <- data366$lon
n.pos.lat <- length(lat)
n.pos.lon <- length(lon)

## remove bisestile days to have all arrays with the same dimensions:
cat("Removing bisestile days. Please wait...\n")
n.years <- year.end - year.start + 1
data <- array(NA, c(1, 1, n.years, 365, n.pos.lat, n.pos.lon))
data[,,,,,] <- data366$obs[,,,1:365,,]

for(y in year.start:year.end){
    y2 <- y - year.start + 1
    if(n.days.in.a.year(y) == 366) data[,,y2,60:365,,] <- data366$obs[,,y2,61:366,,] # take the march to december period removing the 29th of February
}

data366$obs <- NULL ## remove unuseful data 
gc()

## convert psl in daily anomalies with the LOESS filter:
cat("Calculating anomalies. Please wait...\n")    
dataClim <- apply(data, c(1,2,4,5,6), mean, na.rm=T)

my.span <- 0.15
dataClimLoess <- array(NA, dim(dataClim))
for(i in 1:n.pos.lat){
    for(j in 1:n.pos.lon){
        my.data <- data.frame(ens.mean = rep(dataClim[1,1,,i,j],3), day=1:(365*3))
        my.loess <- loess(ens.mean ~ day, my.data, span=my.span)
        dataClimLoess[1,1,,i,j] <- predict(my.loess)[366:730]
    }
}

## check if the my.span parameter is okey:
## plot(dataClimLoess[1,1,,1,1])
## plot(dataClim[1,1,,1,1])

dataClim2 <- InsertDim(dataClimLoess, 3, n.years)
    
data <- data - dataClim2 ## convert SLP data in anomalies

data <- data / 100   # convert MSLP in Pascal to MSLP in hPa

rm(dataClim, dataClim2)
gc()

season.name <- c("winter", "summer")
for(season in 1:2){ ## 1: winter, 2: summer
    ## select only days belonging to one of the two extended seasons:
    pos.last.day.march <- 31 + 28 + 31 ## (=90)
    pos.first.day.october <- 31 + 28 + 31 + 30 + 31 + 30 + 31 + 31 + 30 + 1 ## (=274)

    ## season 1 is winter and goes from Jan to Mar and from Oct do Dec; season 2 is summer and goes from Apr to Sep
    if(season == 1) dd <- c(1:pos.last.day.march, pos.first.day.october:365) else dd <- (pos.last.day.march + 1):(pos.first.day.october - 1)
                                       
    dataPeriod <- data[1,1,, dd,,] # select only days in the chosen period   

    lat.weighted <- cos(lat*pi/180)
    lat.weighted.array <- InsertDim(InsertDim(InsertDim(lat.weighted, 2, n.pos.lon), 1, length(dd)), 1, n.years)

    pslPeriod.weighted <- dataPeriod * lat.weighted.array

    rm(lat.weighted.array)
    gc()

    
    ## Standard R kmeans clustering algorithm is identical to the one below with Vero's function 'WeatherRegime()', so it was replaced by it:
    ##
    ## psl.kmeans <- aperm(pslPeriod.weighted, c(2,1,3,4))
    ## my.cluster <- kmeans(psl.kmeans, centers = n.cl, iter.max = my.iter.max, nstart = my.nstart, trace = FALSE)
    n.cl <- 4; my.iter.max <- 200; my.nstart <- 50
    WR_obs <- WeatherRegime(data = pslPeriod.weighted, EOFS = FALSE, lat = lat, ncenters = n.cl, iter.max = my.iter.max, nstart = my.nstart)
    
    PlotLayout(PlotEquiMap, c(1,2), lon=lon, lat=lat, var=WR_obs$composite, filled.continents=F, axelab=F, draw_separators = T, subsampleg = 1, brks=c(-16, seq(-10, 10, by=2), 16), bar_extra_labels = c(2,0,0,0), ncol=1, fileout=paste0(work.dir,"/", season.name[season],"_patterns/observed_regimes.png"), width=5, height=15)

    ## save weather regimes for season:
    save(WR_obs, file=paste0(work.dir,"/", season.name[season],"/observed_regimes.RData"))

} ## close if on season


#####################################################################################################
#           WR classification at Subseasonal scale with different methodologies                     #
#####################################################################################################

rean <- '/esarchive/recon/ecmwf/erainterim/$STORE_FREQ$_mean/$VAR_NAME$_f6h/$VAR_NAME$_$YEAR$$MONTH$.nc'
fields <- rean
psl <- "psl"

year.start <- 1979 ##1981
year.end <- 2017

## choose the type of clustering to apply to the sequences of 7 days: (from 1 to 364), or more if you want to do a running clustering:
## 1: 1 week only; 3: 3 weeks; 5: 5 weeks ; 0: monthly clustering + first week composite maps
running <- 5

## Set the first Tuesday of the chosen forecast year (it depends on the forecast year. You can also set a different day than Tuesday):
first.tue <- 1 ##5

lat.min <- 27
lat.max <- 81
lon.min <- 274.5
lon.max <- 45

n.cl <- 4
my.iter.max <- 200
my.nstart <- 50

# value of the 'span' parameter of the loess() function. It MUST be updated each time the length of the input data change significantly!!!
my.span <- 0.15            

################################################################################################
n.years <- year.end - year.start + 1

## Load rean psl data:
psleuFull366 <-  Load(var = psl, exp = NULL, obs = list(list(path=fields)), paste0(year.start:year.end,'0101'), storefreq = 'daily', leadtimemax = 366, output = 'lonlat', latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max, nprocs=max(1, availableCores()-1))

## fast load for getting only lat and lon:
##psleuFull366 <-  Load(var = psl, exp = NULL, obs = list(list(path=fields)), paste0(year.start,'0101'), storefreq = 'daily', leadtimemax = 1, output = 'lonlat', latmin = lat.min, latmax = lat.max, lonmin = lon.min, lonmax = lon.max, nprocs=1)

lat <- psleuFull366$lat
lon <- psleuFull366$lon

n.pos.lat <- length(lat)
n.pos.lon <- length(lon)

## remove bisestile days to have all arrays with the same dimensions:
cat("Removing bisestile days. Please wait...\n")
psleuFull <- array(NA,c(1, 1, n.years, 365, n.pos.lat, n.pos.lon))
psleuFull[,,,,,] <- psleuFull366$obs[,,,1:365,,]

for(y in year.start:year.end){
    y2 <- y - year.start + 1
    if(n.days.in.a.year(y) == 366) psleuFull[,,y2,60:365,,] <- psleuFull366$obs[,,y2,61:366,,] # take the march to december period removing the 29th of February
}

psleuFull366$obs <- NULL
gc()


## if((running.cluster || running.15)) {        
##     ## the months after the last month of the last year loaded are automatically filled with NA by Load(); the clustering algorithm doesn't want NA, so
##     ## we have to replace them with some values (in this case the mean of the previous years not to affect the analysis too much)
##     ## if we want to be able to compute the running clustering.
##     ## WARNING: in case your last month is December, for the running clustering it will use the january data of the last year loaded, not of the next year:
##     ##psleuFull[,,n.years,305:334,,] <-  psleuFull[,,n.years-1,305:334,,]
    
##     month.NA <- psleuFull366$not_found_files[1] # filename of the file with the first month with NA (string value with full file path)
##     first.month.NA <- as.integer(substr(month.NA,nchar(month.NA)-4,nchar(month.NA)-3)) # first month with NA
##     first.day.NA <- pos.period(1,first.month.NA)[1] # first day with NA (counting from the start of the year)

##     psleuFullnolast <- psleuFull[,,1:(n.years-1),first.day.NA:365,,]
##     psleuFullMean <- apply(psleuFullnolast, c(2,3,4), mean, na.rm=TRUE) # mean of the previous years
        
##     psleuFull[,,n.years,first.day.NA:365,,] <-  psleuFullMean
        
##     rm(psleuFullnolast, psleuFullMean)

## }
         
## convert psl in daily anomalies with the LOESS filter:
cat("Calculating anomalies. Please wait...\n")
year.end.clim <- year.end # - 1 # last year for computation of rean climatology (unless year.end has data available for all months, it's better to set it to the previous year)
n.years.clim <- year.end.clim - year.start + 1

psleuFull_no_last_year <- psleuFull[1,1, 1:n.years.clim,,,,drop=FALSE]
    
pslPeriodClim <- apply(psleuFull_no_last_year, c(1,2,4,5,6), mean, na.rm=T)
    
rm(psleuFull_no_last_year)
gc()
    

pslPeriodClimLoess <- array(NA, dim(pslPeriodClim))

for(i in 1:n.pos.lat){
    for(j in 1:n.pos.lon){
        my.data <- data.frame(ens.mean=rep(pslPeriodClim[1,1,,i,j],3), day=1:(365*3))
        my.loess <- loess(ens.mean ~ day, my.data, span=my.span)
        pslPeriodClimLoess[1,1,,i,j] <- predict(my.loess)[366:730]
    }
}

rm(pslPeriodClim)
gc()

pslPeriodClim2 <- InsertDim(pslPeriodClimLoess, 3, n.years)
    
## s4.data <- data.frame(s4=pslPeriodClim[1,], day=1:216)
## s4.loess <- loess(s4 ~ day, s4.data, span=0.35)
## s4.pred <- predict(s4.loess)

psleuFull <- psleuFull - pslPeriodClim2

rm(pslPeriodClim2)
gc()

if(psl=="psl") psleuFull <- psleuFull/100   # convert MSLP in Pascal to MSLP in hPa
gc()


## add the days before the 1st of January and after the 31th of december to be able to deal with the running window:
psleuFull.triple <- abind(list(psleuFull, psleuFull, psleuFull), along=4)

WR.sp.patterns <- list()
## compute weather regimes
for(w in 1:52){ ## weeks of the year, starting from the first Tuesday of the forecast year ('first.tue' variable can also be bypassed to start from any day)
    print(w)
    if(running == 1) dd <- first.tue - 1 + (w-1)*7 + (1:7) ## 32:38 ## days corresponding to the week w of the year
    if(running == 3) dd <- first.tue - 1 + (w-1)*7 + (-6:14) ## 25:45 ## days corresponoding to a running windows of 3 weeks centered on week w
    if(running == 5) dd <- first.tue - 1 + (w-1)*7 + (-13:21) ## 18:52 ## days corresponoding to a running windows of 5 weeks centered on week w
    if(running == 0) dd <- pos.month(1, seq.months.in.a.year(1)[(w-1)*7 + 4])  ## daily sequence of all the days of the month belonging to the same central day of week w (eg: 32:59 are the days corresponding to all Feb. days)
    
    pslPeriod <- psleuFull.triple[1,1,, 365 + dd,,] # select only days in the chosen period   

    lat.weighted <- cos(lat*pi/180)
    lat.weighted.array <- InsertDim(InsertDim(InsertDim(lat.weighted, 2, n.pos.lon), 1, length(dd)), 1, n.years)

    pslPeriod.weighted <- pslPeriod * lat.weighted.array

    rm(lat.weighted.array)
    gc()

    ##cat(paste0("Fast formatting pressure for clustering with regimes. Please wait...\n"))
    ##psl.kmeans <- array(aperm(pslPeriod.weighted[,,,, drop=FALSE], c(2,1,3,4)), c(dim(pslPeriod)[1] * dim(pslPeriod)[2], dim(pslPeriod)[3] * dim(pslPeriod)[4]))
    cat(paste0("Clustering regimes for week ", w, "/52. Please wait...\n"))

    ## Standard R kmeans clustering algorithm is identical to the one below with Vero's function 'WeatherRegime()', so it was replaced by it:
    ##
    ## psl.kmeans <- aperm(pslPeriod.weighted, c(2,1,3,4))
    ## my.cluster <- kmeans(psl.kmeans, centers = n.cl, iter.max = my.iter.max, nstart = my.nstart, trace = FALSE)
    WR_obs <- WeatherRegime(data = pslPeriod.weighted, EOFS = FALSE, lat = lat, ncenters = n.cl, iter.max = my.iter.max, nstart = my.nstart)

    WR.sp.patterns[[w]] <- WR_obs$composite
    
    PlotLayout(PlotEquiMap, c(1,2), lon=lon, lat=lat, var=WR_obs$composite, filled.continents=F, axelab=F, draw_separators = T, subsampleg = 1, brks=c(-16, seq(-10, 10, by=2), 16), bar_extra_labels = c(2,0,0,0), ncol=1) ##, fileout=paste0(work.dir,'/weekly/observed_regimes_week_', w, '.png'), width=5, height=15)

    save(WR_obs, file=paste0(work.dir,'/weekly/observed_regimes_week_', w, '.RData'))

    ##par(fig=c(0,1,0,0.15))
    ##ColorBar(brks=seq(-10, 10, by=2), vertical = FALSE, triangle_ends = c(TRUE, TRUE))
    
} ## close for on week w

## ## WR.sp.pat is like WR.sp.patterns but is an array, not a list
## WR.sp.pat <- array(NA, c(52, dim(WR.sp.patterns[[1]])))
## for(w in 1:52) WR.sp.pat[w,,,] <- WR.sp.patterns[[w]]

## PlotLayout(PlotEquiMap, c(2,3), lon=lon, lat=lat, var=WR.sp.pat, filled.continents=F) #, axelab=F, draw_separators = T, subsampleg = 1, brks=seq(-16,16,by=2), bar_extra_labels = c(2,0,0,0), nrow=4, ncol=52, fileout=paste0('~/observed_regimes_week_', w, '.png')) ##, width=5, height=15)


## if you like, before measuring the correlation matrix you can reorder the matrix basing on the quasi-stable classification:
reordering <- TRUE
if(reordering){
    cl <- array(NA, c(52,4))
    cl[1,] <- c(3,2,1,4); cl[2,] <- c(1,3,4,2); cl[3,] <- c(4,1,2,3); cl[4,] <- c(2,3,4,1); cl[5,] <- c(3,2,1,4); cl[6,] <- c(4,3,1,2); cl[7,] <- c(4,2,3,1); cl[8,] <- c(1,3,2,4)
    cl[9,] <- c(3,4,2,1); cl[10,] <- c(1,4,3,2); cl[11,] <- c(2,1,3,4); cl[12,] <- c(3,2,4,1); cl[13,] <- c(3,4,2,1); cl[14,] <- c(3,4,1,2); cl[15,] <- c(1,3,4,2); cl[16,] <- c(2,4,1,3)
    cl[17,] <- c(2,1,4,3); cl[18,] <- c(1,2,4,3); cl[19,] <- c(2,3,1,4); cl[20,] <- c(3,1,2,4); cl[21,] <- c(2,3,1,4); cl[22,] <- c(4,3,2,1); cl[23,] <- c(4,2,1,3); cl[24,] <- c(4,1,3,2)
    cl[25,] <- c(4,3,1,2); cl[26,] <- c(3,2,1,4); cl[27,] <- c(3,2,4,1); cl[28,] <- c(4,2,1,3); cl[29,] <- c(3,1,4,2); cl[30,] <- c(4,1,3,2); cl[31,] <- c(4,3,2,1); cl[32,] <- c(3,4,1,2)
    cl[33,] <- c(1,3,2,4); cl[34,] <- c(1,4,3,2); cl[35,] <- c(2,1,4,3); cl[36,] <- c(3,4,2,1); cl[37,] <- c(2,3,4,1); cl[38,] <- c(1,4,2,3); cl[39,] <- c(3,1,4,2); cl[40,] <- c(3,1,2,4)
    cl[41,] <- c(1,3,4,2); cl[42,] <- c(2,4,3,1); cl[43,] <- c(3,4,1,2); cl[44,] <- c(3,2,4,1); cl[45,] <- c(1,2,3,4); cl[46,] <- c(4,3,1,2); cl[47,] <- c(4,3,2,1); cl[48,] <- c(2,3,1,4)
    cl[49,] <- c(4,1,3,2); cl[50,] <- c(3,1,4,2); cl[51,] <- c(2,3,4,1); cl[52,] <- c(3,4,1,2)
} else {
    ## if you are not reordening the correlation matrix, use this null reordering instead:
    cl <- t(array(1:4, c(4, 52)))
}

#
# correlation matrix:
#

## matrix with the spatial correlations of the 4 WRs for each week of the year:
WR.sp.corr.bigmatrix <- array(NA, c(4*52, 4*52))
WR.sp.patterns <- list()

for(w1 in 1:52){
    print(paste0("Correlating week ", w1))
    
    for(w2 in 1:w1){
        WR.sp.corr.matrix <- array(NA, c(4,4))

        ## load WR patterns of week w:
        load(file=paste0(work.dir,'/weekly_1979-2017_starting_from_first_of_January_of_2014_ERA-Interim/observed_regimes_week_', w1, '.RData'))
        WR.sp.patterns[[w1]] <-  WR_obs$composite

        ## load WR patterns of week w2:
        load(file=paste0(work.dir,'/weekly_1979-2017_starting_from_first_of_January_of_2014_ERA-Interim/observed_regimes_week_', w2, '.RData'))
        WR.sp.patterns[[w2]] <-  WR_obs$composite

        ## if you want, you can check the spatial patterns of the WRs of weeks w1 and w2:
        ## PlotEquiMap(WR.sp.patterns[[w1]][,, cl[w1, 1]], lon, lat, filled.continents=F, brks=seq(-10, 10, 2))
        
        WR.sp.corr.matrix[1,1] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 1]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 1]]))
        if(w2 < w1) WR.sp.corr.matrix[1,2] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 1]]), as.vector(WR.sp.patterns[[w2]][,,  cl[w2, 2]]))
        if(w2 < w1) WR.sp.corr.matrix[1,3] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 1]]), as.vector(WR.sp.patterns[[w2]][,,  cl[w2, 3]]))
        if(w2 < w1) WR.sp.corr.matrix[1,4] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 1]]), as.vector(WR.sp.patterns[[w2]][,,  cl[w2, 4]]))
        
        WR.sp.corr.matrix[2,1] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 2]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 1]]))
        WR.sp.corr.matrix[2,2] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 2]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 2]]))
        if(w2 < w1) WR.sp.corr.matrix[2,3] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 2]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 3]]))
        if(w2 < w1) WR.sp.corr.matrix[2,4] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 2]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 4]]))
        
        WR.sp.corr.matrix[3,1] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 3]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 1]]))
        WR.sp.corr.matrix[3,2] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 3]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 2]]))
        WR.sp.corr.matrix[3,3] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 3]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 3]]))
        if(w2 < w1) WR.sp.corr.matrix[3,4] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 3]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 4]]))
        
        WR.sp.corr.matrix[4,1] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 4]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 1]]))
        WR.sp.corr.matrix[4,2] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 4]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 2]]))
        WR.sp.corr.matrix[4,3] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 4]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 3]]))
        WR.sp.corr.matrix[4,4] <- cor(as.vector(WR.sp.patterns[[w1]][,, cl[w1, 4]]), as.vector(WR.sp.patterns[[w2]][,, cl[w2, 4]]))
        
        WR.sp.corr.bigmatrix[(w1 - 1) * 4 + 1:4, (w2 - 1) * 4 + 1:4] <- WR.sp.corr.matrix
    }
}

##edit(round(WR.sp.corr.bigmatrix,2))

## save the correlation matrix (original or reordered):
save(WR.sp.corr.bigmatrix, file=paste0(work.dir,"/weekly_1979-2017_starting_from_first_of_January_of_2014_ERA-Interim/correlation_matrix.RData"))

## if you want, you can restrict correlations to values greater than 0.8:
WR.sp.corr.bigmatrix2 <- WR.sp.corr.bigmatrix
WR.sp.corr.bigmatrix2[which(WR.sp.corr.bigmatrix2 < 0.8)] <- NA
edit(round(WR.sp.corr.bigmatrix2,2))

## add the climatological frequency of the WRs over each figure:
clim_frequencies <- paste0('freq = ',round(Mean1Dim(WR_obs$frequency,1),1),'%')
    
if(running != 1) {
    length(dd)
    cluster.array <- matrix(WR_obs$cluster, 35, length(dd) )  ## matrix with cluster sequence for all cluster period; format: years x days (7 or 21)
    cl.freq <- c()
    
    for(cl in 1:n.cl){
        ## set to NA days not belonging to regime cl:
        ss <- which(cluster.array != cl, arr.ind=T)
        psl.kmeans2 <- psl.kmeans        
        for(i in 1:dim(ss)[1]) psl.kmeans2[ss[i,2],ss[i,1],,] <- NA

        ## select only days in the central week:
        if(running == 3) psl.kmeans3 <- psl.kmeans2[8:14,,,]
        if(running == 5) psl.kmeans3 <- psl.kmeans2[15:21,,,]
        if(running == 0) psl.kmeans3 <- psl.kmeans2[1:7,,,]
        
        ## average psl data to get WR pattern:
        WR_obs$composite[,,cl] <- apply(psl.kmeans3, c(3,4), mean, na.rm=TRUE)

        ## measure frequency of occurrence:
        cl.freq[cl] <- 100 * (dim(psl.kmeans3)[1] * dim(psl.kmeans3)[2] - length(which(is.na(psl.kmeans3[,,1,1])))) /  (dim(psl.kmeans3)[1] * dim(psl.kmeans3)[2])
        
    }

    clim_frequencies <- paste0('freq = ',round(cl.freq,1),'%')
}

PlotLayout(PlotEquiMap, c(1,2), lon=lon, lat=lat, var=WR_obs$composite,
           titles=paste0(paste0('Cluster ',1:4),' (',clim_frequencies,' )'), filled.continents=F,
           axelab=F, draw_separators = T, subsampleg = 1, brks=seq(-16,16,by=2),
           bar_extra_labels = c(2,0,0,0)) ##,fileout='~/observed_regimes.png')

##
## Plot the correlation matrix 
##

## load correlation matrix "WR.sp.corr.bigmatrix":
load(file=paste0(work.dir,"/weekly/correlation_matrix.RData"))

##my.cols <- c('white','#fef0d9','#fdd49e','#fdbb84','#fc8d59','#e34a33','#b30000', '#7f0000')
##my.cols <- c('#fff7fb','#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')
##my.cols <- c('#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')
my.cols <- rev(c('#67001f','#b2182b','#d6604d','#f4a582','#fddbc7','#d1e5f0','#92c5de','#4393c3','#2166ac','#053061'))
##my.brks <- c(-0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.30, 1)
##my.brks <- c(0,20,30,40,50,60,70,80,90,100)
my.brks <- seq(-1,1,0.2)

my.file <- paste0(work.dir,"/correlation_matrix.png")
png(file=my.file, width=700*5, height=700*5)  
plot.new()

##par(mar = c(0, 0, 2, 0), fig = c(0.32, 0.57, 0.8, 0.94), new=TRUE)
par(fig = c(0, 0.90, 0, 1))
plot(1, 1, type = "n", xaxt = "n", yaxt = "n", ylim = c(1, 208.04), xlim = c(-0.23, 206.2), ann = F, bty="n") ## to remove black border

## Add axes labels:
mtext(side = 1, text = "                   Week", line = 3, cex = 1.2, font = 3)
mtext(side = 2, text = "Week", line = 3, cex = 1.2, font = 3)

## Add figure title:
my.title = "Correlation matrix between spatial patterns of weekly WRs\n(source: ERA-Interim)"
mtext(side = 3, text = my.title, cex = 1.3)

## old x-axis that cannot be rotated of 45 degrees; only used to draw ticks:
axis(1, at=seq(0,207), las=3, cex.axis=1, labels=FALSE, mgp=c(3,0.7,0)) 

## Plot x-axis labels rotated of 45 degrees at default tick marks:
##labels <- paste0(" ",1:31)
##text(0:207, par("usr")[3] - 1.55, srt = 0, adj = 1, labels = 1:208, xpd = TRUE, cex=.8) 

## Draw subtitle:
##mtext(1, text = "                                     Reference dataset: Era-Interim", line = 7)

## Plot y-axis:
axis(2, at=208:1, las=2, labels=1:208, cex.axis=.8)

## create an array similar to array.pers but with colors instead of frequencies:                
n.cols <- length(my.cols)
array.colors <- array(my.cols[n.cols], c(208,208))
for(int in (n.cols-1):1) array.colors[WR.sp.corr.bigmatrix < my.brks[int+1]] <- my.cols[int]

##set to gray the NAs:
array.colors[which(is.na(WR.sp.corr.bigmatrix))] <- "#ece7f2"

## fill with colors the cells in the figure:
for(p in 1:208){ for(l in 0:(208-1)){ polygon(c(0.5+l-1, 0.5+l-1, 1.5+l-1, 1.5+l-1), c(-0.5+209-p, 0.5+209-p, 0.5+209-p, -0.5+209-p), col=array.colors[p,1+l], border="black") }}

## add color legend:
par(fig = c(0.88, 1, 0.30, 0.75), new=TRUE)
ColorBar(brks=my.brks, cols=my.cols, vert=TRUE, label_scale=5.3, subsample=1) #my.ticks=-0.5 + 1:length(my.brks.var2), my.labels=my.brks.var2) subset=my.subset2)
    
dev.off()    


##
## Plot the correlation matrix only for each regime individually
##

## load correlation matrix "WR.sp.corr.bigmatrix":
load(file=paste0(work.dir,"/weekly_1979-2017_starting_from_first_of_January_of_2014_ERA-Interim/correlation_matrix.RData"))

## select only correlations between different weeks of the same regime:
wr.pos <- wr.corr <- list()
for(wr in 1:4) wr.pos[[wr]] <- seq(wr, 208, 4)
for(wr in 1:4) wr.corr[[wr]] <- WR.sp.corr.bigmatrix[wr.pos[[wr]], wr.pos[[wr]]]

##my.cols <- c('white','#fef0d9','#fdd49e','#fdbb84','#fc8d59','#e34a33','#b30000', '#7f0000')
##my.cols <- c('#fff7fb','#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')
##my.cols <- c('#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')
my.cols <- rev(c('#67001f','#b2182b','#d6604d','#f4a582','#fddbc7','#d1e5f0','#92c5de','#4393c3','#2166ac','#053061'))
##my.brks <- c(-0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.30, 1)
##my.brks <- c(0,20,30,40,50,60,70,80,90,100)
my.brks <- seq(-1,1,0.2)

for(wr in 1:4){
    my.file <- paste0(work.dir, "/weekly_1979-2017_starting_from_first_of_January_of_2014_ERA-Interim/correlation_matrix_regime_", wr,".png")
    png(file=my.file, width=700*5, height=700*5)  
    plot.new()

    ##par(mar = c(0, 0, 2, 0), fig = c(0.32, 0.57, 0.8, 0.94), new=TRUE)
    par(fig = c(0, 0.90, 0, 1))
    plot(1, 1, type = "n", xaxt = "n", yaxt = "n", ylim = c(1, 52.04), xlim = c(-0.23, 50.2), ann = F, bty="n") ## to remove black border

    ## Add axes labels:
    mtext(side = 1, text = "                   Week", line = 3, cex = 4.2, font = 3)
    mtext(side = 2, text = "Week", line = 1, cex = 4.2, font = 3)

    ## Add figure title:
    my.title = paste0("Correlation matrix between spatial patterns of weekly WR #", wr," \n(source: ERA-Interim)")
    mtext(side = 3, text = my.title, cex = 5.3, line = -6)

    ## old x-axis that cannot be rotated of 45 degrees:
    axis(1, at=seq(0, 51), las=3, cex.axis=4, labels=1:52, mgp=c(3,0.7,0), line = -5) 

    ## Plot x-axis labels rotated of 45 degrees at default tick marks:
    ##labels <- paste0(" ",1:31)
    ##text(0:207, par("usr")[3] - 1.55, srt = 0, adj = 1, labels = 1:208, xpd = TRUE, cex=.8) 

    ## Draw subtitle:
    ##mtext(1, text = "                                     Reference dataset: Era-Interim", line = 7)

    ## Plot y-axis:
    axis(2, at=52:1, las=2, labels=1:52, cex.axis=4.8, line = -5)

    ## create an array similar to array.pers but with colors instead of frequencies:                
    n.cols <- length(my.cols)
    array.colors <- array(my.cols[n.cols], c(52, 52))
    for(int in (n.cols-1):1) array.colors[wr.corr[[wr]] < my.brks[int+1]] <- my.cols[int]

    ##set to gray the NAs:
    array.colors[which(is.na(wr.corr[[wr]]))] <- "#ece7f2"

    ## fill with colors the cells in the figure:
    for(p in 1:52){ for(l in 0:(52-1)){ polygon(c(0.5+l-1, 0.5+l-1, 1.5+l-1, 1.5+l-1), c(-0.5+53-p, 0.5+53-p, 0.5+53-p, -0.5+53-p), col=array.colors[p,1+l], border="black") }}

    ## add color legend:
    par(fig = c(0.88, 1, 0.30, 0.75), new=TRUE)
    ColorBar(brks=my.brks, cols=my.cols, vert=TRUE, label_scale=5.3, subsample=1) #my.ticks=-0.5 + 1:length(my.brks.var2), my.labels=my.brks.var2) subset=my.subset2)
    
    dev.off()    

} ## close for on wr



#####################################################################################################
#   WR forecast and obs. comparison at subseasonal scale with 2 methods, assign and spatial corr.   #
#####################################################################################################

yr1.obs <- 1979 # first year of psl data from the reanalysis
yr2.obs <- 2017 # last year of psl data from the reanalysis 
    
forecast.year <- 2014 ##2016             # starting year of the weekly sequence of the forecasts
mes <- 1                                                 # starting forecast month (1: january, 2: february, etc.)
day <- 2 #7                                            # starting forecast day (usually the first Thursday of January of the forecast year,
                                                                # which in case of forecasts of 2014 is the 2 of January and in case of the forecasts of 2016 is the 7 of January)
                                                                # (for 2014 forecasts, also Tuesday forecasts were downloaded. They start on the 7th of January)
yr1.hind <- 1994 #1996                       # first hindcast year
yr2.hind <- 2013 #2015                                   # last hindcast year (usually the forecast year -1)

## 6-hourly data (not used unless you have no daily means data available on /esarchive and you have to averege them by yourself):
forecast.dir <- paste0('/esarchive/exp/ecmwf/monthly_ensforhc/6hourly/$VAR_NAME$/', forecast.year,'$MONTH$$DAY$00/$VAR_NAME$_$START_DATE$00.nc') ## forecast year: 2014
##forecast.dir <- paste0('/esarchive/exp/ecmwf/monthly_ensforhc/6hourly/$VAR_NAME$/', forecast.year,'$MONTH$$DAY$/$VAR_NAME$_$START_DATE$.nc')  ## forecast year: 2016

## daily data:
##forecast.dir <- paste0('/esarchive/exp/ecmwf/monthly_ensforhc/daily_mean/$VAR_NAME$_f6h/', forecast.year,'$MONTH$$DAY$00/$VAR_NAME$_$START_DATE$.nc') ## forecast year: 2014

n.clusters <- 4

##################################################################################################################################
n.yrs.hind <- yr2.hind - yr1.hind + 1
n.yrs.obs <- yr2.obs - yr1.obs + 1

## do it only once to save the observed monthly cluster values:
for(month in 1:12){    
    ## you have to load reanalysis data separately from forecasts because it has more years of data available, so you can classify observed WRs using a longer period
    ## you also have to load the data of the whole month (not weekly) to be able to classify WRs once:
    ## notice that leap days are not loaded!!!
    month2 <- ifelse(month >= 10, month, paste0('0',month))  ## add a 0 to the left of the month number if the month is a number lower than 10
    sdates <- paste0(yr1.obs:yr2.obs, paste0(month2,'01')) ## [1,1,39 (yrs), 31 (days), 77, 186]
    data.rean <- Load('psl', exp = NULL,
                 obs = 'erainterim',
                 sdates = sdates,
                 storefreq='daily',
                 leadtimemin = 1, leadtimemax = n.days.in.a.month(month,2001),
                 latmin = 27, latmax = 81,
                 lonmin = 274.5, lonmax = 45,output = 'lonlat', nproc=availableCores()-1) ## it is 2 times faster loading with 10 cpus instead of only one. even if it is not interpolating!!!

    clim <- Clim(var_exp = data.rean$obs,var_obs = data.rean$obs,memb=F)
    clim$clim_exp <- NULL ## remove the clim_exp variable because it was itnroduced only to be able to use the Clim() function
    clim_smoothed_obs <- aperm(apply(clim$clim_obs,c(1:length(dim(clim$clim_obs)))[-which(names(dim(clim$clim_obs))=='ftime')],Loess,loess_span=1),c(2,1,3,4)) ## apply Loess filter
    anom_obs <- Ano(data.rean$obs,clim_smoothed_obs)

    ## cluster the observed psl data:
    WR_obs <- WeatherRegime(data=anom_obs, EOFS=FALSE, lat=data.rean$lat, ncenters = 4)

    ## wait for X windows:
    ##clim_frequencies<-paste0('freq = ',round(Mean1Dim(WR_obs$frequency,1),1),'%')
    ##PlotLayout(PlotEquiMap,c(2,1),lon=data$lon,lat=data$lat,var=WR_obs$composite/100,
    ## titles=paste0(paste0('Cluster ',1:4),' (',clim_frequencies,' )'),filled.continents=F,
    ## axelab=F,draw_separators = T,subsampleg = 1,brks=seq(-16,16,by=2),
    ## bar_extra_labels = c(2,0,0,0),fileout='observed_regimes.png')
    ##PlotEquiMap(WR_obs$composite[,,1,1,1], lon=data$lon, lat=data$lat)

    reference <- drop(WR_obs$composite)
    names(dim(reference)) <- c('lat','lon','nclust')
    
    WR.seq.obs <- drop(WR_obs$cluster)

    ## save the clusters for each month:
    save(reference, WR.seq.obs, file=paste0(work.dir,"/WRs_observed_",period.name[month],".RData"))

    ## rm(data.rean)
}


##
## Assign method:
##
## Load Subseasonal SLP hindcast data for one start date and lead time (1 week), interpolate it to the same grid of the reanalysis 
## and finally load SLP monthly reanalysis data to assign to each day an observed WR
##

## detect lon and lat of reanalysis to be able to interpolate the forecasts to the reanalysis grid:
data.grid <- Load('psl', exp = NULL, 
                 obs = 'erainterim',
                 sdates = '20010101',
                 storefreq='daily',
                 leadtimemin = 1, leadtimemax = 1,
                 output = 'lonlat', nproc=1) ## load 1 day only

lons <- data.grid$lon
lats <- data.grid$lat
n.lon.rean <- length(lons) 
n.lat.rean <- length(lats)
my.grid <- paste0('r',n.lon.rean,'x',n.lat.rean)

sdates.seq <- weekly.seq(forecast.year, mes, day)  # load the sequence of dates corresponding to all the thursday of the year
pos.bis <- which(sdates.seq == paste0(forecast.year,"0229")) # find if there is a startdate at the 29th of february and return its position in the vector sdates.seq
if(length(pos.bis) != 0) sdates.seq <- sdates.seq[-pos.bis] # if there is a start date during the 29th of  February, remove it
n.sdates <- length(sdates.seq) ## number of start dates

WR.seq.perc <- array(0, c(n.sdates, 32)) ## array which will store the % or same WRs forecast and observed, for each day in the 31 lead times
for(sdate in 1:n.sdates){
    print(paste0("Start date num. ",sdate, "/",n.sdates))
    
    sdates <- paste0(yr1.hind:yr2.hind,substr(sdates.seq[sdate],5,6),substr(sdates.seq[sdate],7,8)) ## same start date but for all hindcast years
    
    ## load forecasts for that all the years of start date 'sdate' and ALL FORECAST DAYS associated to it (32) and interpolate them to the same resolution of the reanalysis:
    ## note that the first forecast day is not equal to startdate+1 but to startdate!!!
    data <- Load('psl', exp = list(list(path=forecast.dir)),
                 obs = NULL,
                 sdates = sdates,
                 storefreq='daily',
                 leadtimemin = 1, leadtimemax = 129, ## replace 129 with 32 when you'll load daily data directly
                 latmin = 27, latmax = 81,
                 lonmin = 274.5, lonmax = 45,output = 'lonlat',
                 grid=my.grid, method='bilinear', nprocs=max(1, availableCores()-1))

    n.members <- dim(data$mod)[2]
    n.lat <-  dim(data$mod)[5]
    n.lon <-  dim(data$mod)[6]

    ## In case you are loading 6-hourly data, you have to convert them to daily means before:
    data.daily <- array(NA,c(1,n.members,n.yrs.hind,32,n.lat,n.lon)) ## [dataset x members x sdate x time x lat x lon]
    my.names <- names(dim(data$mod))

    ##for(dd in 1:32) { print(paste0("Day ", dd, " of 32")); data.daily[1,1:n.members,1:n.yrs.hind, dd, 1:n.lat,1:n.lon] <-  apply(data$mod, c(1,2,3,5,6), function(x){mean(x[(dd-1)*4 + 1:4])})}
    
    ## it is twice faster with .Internal(mean), but ther must be no NA values:
    for(dd in 1:32) { print(paste0("Day ", dd, " of 32")); data.daily[1,1:n.members,1:n.yrs.hind, dd, 1:n.lat,1:n.lon] <-  apply(data$mod, c(1,2,3,5,6), function(x){.Internal(mean(x[(dd-1)*4 + 1:4]))})}

    data$mod <- data.daily
    names(dim(data$mod)) <- my.names ## re-assign the dimension names because if they were deleted by the previous command

    rm(data.daily)

    ## remove the eventual forecasts of the 29th of February: (they are present for the forecasts of 2016, but not of 2014):

    
    ## measure forecast climatology for the 32 days loaded:
    clim <- Clim(var_exp = data$mod, var_obs = data$mod, memb=F)

    ## apply the LOESS filter to the 32-days climatology:
    clim_smoothed_exp <- aperm(apply(clim$clim_exp, c(1,3,4), Loess, loess_span=1), c(2,1,3,4))

    ## measure forecast anomalies:
    anom_exp <- Ano(data$mod, clim_smoothed_exp)

    ## for each forecast day of the start date 'sdate', identify to which month it belongs:
    forecast.daily.seq <- as.numeric(substr(sdates.seq[sdate],7,8)) + 1:32 - 1 ## forecast days starting counting from the SAME month of the start date
    
    ## month of the currently selected start date:
    month.sdate <- as.numeric(substr(sdates.seq[sdate],5,6))
    
    ## get the number of days BEFORE the month of the start date (for no leap years):
    if(month.sdate > 1) n.days.previous <- n.days.in.a.monthly.period(1,month.sdate-1,2001) else n.days.previous <- 0

    ## get the sequence of forecast days counting from the START of the YEAR:
    days.seq <- n.days.previous + forecast.daily.seq

    ## get the sequence of month numbers corresponding to the sequence of forecast days:
    monthly.seq <- seq.months.in.a.year(2001)[days.seq]
    if(length(which(is.na(monthly.seq))) > 0) monthly.seq[which(is.na(monthly.seq))] <- 1## if the forecast days extend to the following year, assign them to January
                
    ## check all the sequences (for the debug):
    cbind(1:length(days.seq),days.seq,monthly.seq) 

    ## identify which month(s) are to load:
    months2load <- unique(monthly.seq)

    WR.seq.for <- array(NA, dim(anom_exp)[1:4]) ## array which will store the cluster sequence of the predicted WRs for all the 32 forecast days of the selected sdate
    WR.seq.rean <- array(NA,c(n.yrs.hind, 32)) ## array which will store the cluster sequence of the observed WRs for all the 32 forecast days of the selected sdate    
    for(month in months2load){
        ## load 'reference' and 'WR.seq.obs' variables with the cluster data of the same month of the forecast day considered:    
        load(paste0(work.dir,"/WRs_observed_",period.name[month],".RData"))
        
        ## extract only forecast anomalies of days belonging to that month:
        pos.month <- which(monthly.seq == month) ## pos.month is the number of the forecast days which belongs to the selected month
        anom_exp_same_month <- anom_exp[,,,pos.month,,]
                                        
        ## measure distance of each forecast day to the reference month(s):
        WR.seq.for[,,,pos.month] <- RegimesAssign(var_ano = anom_exp_same_month, ref_maps = reference, lats = data$lat, method = 'distance')$cluster
        ##   WR_exp <- RegimesAssign(var_ano=anom_exp,ref_maps=reference,lats=data$lat,method='distance')

        ##  ## select only the days in the selected month:
        ## WR.seq.for2 <-  WR.seq.for[1,,,pos.month]

        ## reshape the daily sequence of observed monthly WRs in an array [years, days]:
        WR.seq.obs.array <- t(array(WR.seq.obs, c(n.days.in.a.month(month,2001), n.yrs.obs)))

        ## select only the reanalysis years within the hindcasts years to be able to compare the obsered WRs with the forecasted ones:
        yr1 <- yr1.hind - yr1.obs + 1
        yr2 <- yr2.hind - yr1.obs + 1

        WR.seq.obs.array2 <- WR.seq.obs.array[yr1:(yr2+1),] ## load one year more to be able to assign obs WRs for forecast days beyond the 31th of December (see below)

        ndm <- n.days.in.a.month(month,2001) ## number of days in 'month'
        
        if(length(months2load) == 2) {
            if (month == months2load[1]) pos.fd <- (ndm - length(which(monthly.seq == month)) + 1) : ndm
            if (month == months2load[2]) pos.fd <- 1: length(which(monthly.seq == month))
        }
        
        if(length(months2load) == 3) {
            if (month == months2load[1]) pos.fd <- (ndm - length(which(monthly.seq == month)) + 1) : ndm
            if (month == months2load[2]) pos.fd <- 1:ndm
            if (month == months2load[3]) pos.fd <- 1:length(which(monthly.seq == month))
        }
            
        ## Select only the part of WR.seq.obs.array inside the forecast days of the selected month.
        ## Note that WR.seq.rean[,X] indicates the observed WRs corresponding to the forecast days, while WR.seq.obs.array indicates the observed WRs
        ## corresponding to the calendar days, so you have to convert the second to the first with the command below:
        WR.seq.rean[,pos.month] <- WR.seq.obs.array2[1:n.yrs.hind, pos.fd]

        ## correct in case the forecast days enter in the following year:
        if(month != months2load[1] && months2load[2] < months2load[1]) WR.seq.rean[,pos.month] <- WR.seq.obs.array2[-1, pos.fd]
            
            
    } ## close for on month

    ## compare the forecasted WR sequence with the observed one:
    for(d in 1:32){
        for(y in 1:n.yrs.hind){
            WR.seq.perc[sdate,d] <- WR.seq.perc[sdate,d] + length(which(WR.seq.for[1,,y,d] == WR.seq.rean[y,d]))
        }
    }

    print(WR.seq.perc)    
} ## close for on sdate

## convert it to a %:
WR.seq.perc <- 100 * WR.seq.perc / (n.yrs.hind * n.members)

save(WR.seq.perc, file=paste0(work.dir,"/Assign_table.Rdata"))


##
## Plot the figure
##

load(paste0(work.dir,"/Assign_table.Rdata"))

##my.cols <- c('white','#fef0d9','#fdd49e','#fdbb84','#fc8d59','#e34a33','#b30000', '#7f0000')
##my.cols <- c('#fff7fb','#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')
my.cols <- c('#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')

##my.brks <- c(-0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.30, 1)
##my.brks <- c(0,20,30,40,50,60,70,80,90,100)
my.brks <- c(0,30,40,50,60,70,80,90,100)

## ColorTable <- function(matrix, cols, brks, fileout = NULL, wifth = 8, height = 5, size_units = "in", res = 100){
##     if (!is.null(fileout)) {
##         saveToFile(fileout)
##     } else {
##         dev.new(units = size_units, res = res, width = width, height = height)
##     }
## }


my.file <- paste0(work.dir,"/Assign_figure.png")
png(file=my.file, width=530, height=700)  
plot.new()

##par(mar = c(0, 0, 2, 0), fig = c(0.32, 0.57, 0.8, 0.94), new=TRUE)
par(fig = c(0, 0.88, 0, 1))
plot(1, 1, type = "n", xaxt = "n", yaxt = "n", ylim = c(1, 52.04), xlim = c(-0.23, 30.2), ann = F, bty="n") ## to remove black border

## Add axes labels:
mtext(side = 1, text = "                   Forecast day", line = 3, cex = 1.2, font = 3)
mtext(side = 2, text = "Start date", line = 3, cex = 1.2, font = 3)

## Add figure title:
my.title = "Coincident regimes (in %) between ECMWF-MPS and ERA-Interim\nwith the assign method"
mtext(side = 3, text = my.title, cex = 1.3)

## old x-axis that cannot be rotated of 45 degrees; only used to draw ticks:
axis(1, at=seq(0,31), las=3, cex.axis=1, labels=FALSE, mgp=c(3,0.7,0)) 

## Plot x-axis labels rotated of 45 degrees at default tick marks:
##labels <- paste0(" ",1:31)
text(0:31, par("usr")[3] - 1.55, srt = 0, adj = 1, labels = 1:32, xpd = TRUE, cex=.8) 

## Draw subtitle:
##mtext(1, text = "                                     Reference dataset: Era-Interim", line = 7)

## Plot y-axis:
axis(2, at=52:1, las=2, labels=1:52, cex.axis=.8)

## create an array similar to array.pers but with colors instead of frequencies:                
n.cols <- length(my.cols)
array.colors <- array(my.cols[n.cols], c(52,32))
for(int in (n.cols-1):1) array.colors[WR.seq.perc < my.brks[int+1]] <- my.cols[int]

##set to gray the last column:
##array.colors[,32] <- "#ece7f2"

## fill with colors the cells in the figure:
for(p in 1:52){ for(l in 0:(32-1)){ polygon(c(0.5+l-1, 0.5+l-1, 1.5+l-1, 1.5+l-1), c(-0.5+53-p, 0.5+53-p, 0.5+53-p, -0.5+53-p), col=array.colors[p,1+l], border="black") }}

## add color legend:
par(fig = c(0.83, 1, 0.30, 0.75), new=TRUE)
ColorBar(brks=my.brks, cols=my.cols, vert=TRUE, label_scale=1.3, subsample=1) #my.ticks=-0.5 + 1:length(my.brks.var2), my.labels=my.brks.var2) subset=my.subset2)
    
dev.off()    




##
##  Spatial corr method:
##  Classify WR from the forecasted psl directly and then related them to the observed WR taking the highest spatial correlations of the regime patterns
##
##

## detect lon and lat of reanalysis to be able to interpolate the forecasts to the reanalysis grid:
data.grid <- Load('psl', exp = NULL, 
                 obs = 'erainterim',
                 sdates = '20010101',
                 storefreq='daily',
                 leadtimemin = 1, leadtimemax = 1,
                 output = 'lonlat', nproc=1) ## load 1 day only

lons.rean <- data.grid$lon
lats.rean <- data.grid$lat
n.lon.rean <- length(lons.rean) 
n.lat.rean <- length(lats.rean)
my.grid <- paste0('r',n.lon.rean,'x',n.lat.rean)


sdates.seq <- weekly.seq(forecast.year, mes, day)  # load the sequence of dates corresponding to all the thursday of the year
pos.bis <- which(sdates.seq == paste0(forecast.year,"0229")) # find if there is a startdate at the 29th of february and return its position in the vector sdates.seq
if(length(pos.bis) != 0) sdates.seq <- sdates.seq[-pos.bis] # if there is a start date during the 29th of  February, remove it
n.sdates <- length(sdates.seq) ## number of start dates

WR.seq.perc2 <- array(0, c(n.sdates,32)) ## array which will store the % or same WRs forecast and observed, for each day in the 31 lead times
for(sdate in 1:n.sdates){
    print(paste0("Start date num. ",sdate, "/",n.sdates))
    
    sdates <- paste0(yr1.hind:yr2.hind,substr(sdates.seq[sdate],5,6),substr(sdates.seq[sdate],7,8)) ## same start date but for all hindcast years
    
    ## load forecasts for that all the years of start date 'sdate' and ALL FORECAST DAYS associated to it (32) and interpolate them to the same resolution of the reanalysis:
    ## note that the first forecast day is not equal to startdate+1 but to startdate!!!
    data <- Load('psl', exp = list(list(path=forecast.dir)),
                 obs = NULL,
                 sdates = sdates,
                 storefreq='daily',
                 leadtimemin = 1, leadtimemax = 129, ## replace 129 with 32 when you'll load daily data directly
                 latmin = 27, latmax = 81,
                 lonmin = 274.5, lonmax = 45,output = 'lonlat',
                 grid=my.grid, method='bilinear', nprocs=max(1, availableCores()-1))

    n.members <- dim(data$mod)[2]
    n.lat <-  dim(data$mod)[5]
    n.lon <-  dim(data$mod)[6]

    ## In case you are loading 6-hourly data, you have to convert them to daily means before:
    data.daily <- array(NA,c(1,n.members,n.yrs.hind,32,n.lat,n.lon)) ## [dataset x members x sdate x time x lat x lon]
    my.names <- names(dim(data$mod))
    for(dd in 1:32) { print(paste0("Day ", dd, " of 32"));  data.daily[1,1:n.members,1:n.yrs.hind, dd, 1:n.lat,1:n.lon] <-  apply(data$mod, c(1,2,3,5,6), function(x){mean(x[(dd-1)*4+1:4])})}
    data$mod <- data.daily
    names(dim(data$mod)) <- my.names ## re-assign the dimension names because if they were deleted by the previous command
    rm(data.daily)

    ## ## using Apply() instead of appl() with 9 cores, it is MUCH SLOWER!!!!
    ## library(multiApply)
    ##for(dd in 1:32) { print(paste0("Day ", dd, " of 32"));  data.daily[1,1:n.members,1:n.yrs.hind, dd, 1:n.lat,1:n.lon] <-  Apply(data$mod, margins=c(1,2,3,5,6), fun=function(x){mean(x[1+(dd-1)*4+1:4])}, ncores=9)}
    ##
    ## comparison between apply and Apply:
    ##
    ## system.time(Apply(data=data$mod, margins=c(1,2,3,5,6), fun=function(x){mean(x[1:4])}, ncores=9)) ## 27 s
    ## system.time(apply(X=data$mod, MARGIN=c(1,2,3,5,6), FUN=function(x){mean(x[1:4])})) ## 16 s
    ##
    ## other examples:
    ##
    ## my.array <- array(rnorm(10000000), c(1000,1000,100))
    ## system.time(apply(X=my.array, MARGIN=c(1,2), FUN=mean)) ## 19.6 s
    ## system.time(Apply(data=my.array, margins=c(1,2), fun=mean, ncores=9)) ## 19.2 s
    ##
    ## my.array <- array(rnorm(10000000), c(1000,1000,100))
    ## system.time(apply(X=my.array, MARGIN=c(1,2), FUN=function(x){max(x+5*x^2)})) ## 8.1 s
    ## system.time(Apply(data=my.array, margins=c(1,2), fun=function(x){max(x+5*x^2)}, ncores=9)) ## 16.3 s
    ##
    ## ## using parApply() it TAKES ALMOST THE SAME TIME!!! However, we can't use the variable dd inside function(x), which makes parApply totally unuseful!!!
    ## system.time(test <- parApply(cl, X=data$mod, MARGIN=c(1,2,3,5,6), FUN=function(x){mean(x[1+(10-1)*4+1:4])}))
    ## system.time(test <- apply(X=data$mod, MARGIN=c(1,2,3,5,6), FUN=function(x){mean(x[1+(10-1)*4+1:4])}))
    ##
    ## library(parallel)
    ## cl <- makeCluster(9)
    ## system.time(parLapply(cl, rep(10,10000000), function(x) c(x, x^2, x^3)))
    ## stopCluster(cl)
    ## system.time(lapply( rep(10,10000000), function(x) c(x, x^2, x^3))) ## confront computational time with serial version

    
    ## remove the eventual forecasts of the 29th of February: (they are present when the forecasts are of 2016)

    
    ## measure forecast climatology for the 32 days loaded:
    clim <- Clim(var_exp = data$mod,var_obs = data$mod,memb=F)

    ## apply the LOESS filter to the 32-days climatology:
    clim_smoothed_exp <- aperm(apply(clim$clim_exp,c(1,3,4),Loess,loess_span=1),c(2,1,3,4))

    ## measure forecast anomalies:
    anom_exp <- Ano(data$mod, clim_smoothed_exp)

    ## for each forecast day of the start date 'sdate', identify to which month it belongs:
    forecast.daily.seq <- as.numeric(substr(sdates.seq[sdate],7,8)) + 1:32 - 1 ## forecast days starting counting from the SAME month of the start date
    
    ## month of the currently selected start date:
    month.sdate <- as.numeric(substr(sdates.seq[sdate],5,6))
    
    ## get the number of days BEFORE the month of the start date (for no leap years):
    if(month.sdate > 1) n.days.previous <- n.days.in.a.monthly.period(1,month.sdate-1,2001) else n.days.previous <- 0

    ## get the sequence of forecast days counting from the START of the YEAR:
    days.seq <- n.days.previous + forecast.daily.seq

    ## get the sequence of month numbers corresponding to the sequence of forecast days:
    monthly.seq <- seq.months.in.a.year(2001)[days.seq]
    if(length(which(is.na(monthly.seq))) > 0) monthly.seq[which(is.na(monthly.seq))] <- 1## if the forecast days extend to the following year, assign them to January
                
    ## check all the sequences for the debug:
    cbind(1:length(days.seq),days.seq,monthly.seq) 

    ## identify which month(s) are to load:
    months2load <- unique(monthly.seq)

    ## classify forecasted psl in WRs:
    lat.weighted <- cos(data$lat*pi/180)
    lat.weighted.array <- InsertDim(InsertDim(InsertDim(InsertDim(var=lat.weighted,2,n.lon),1, 32), 1, n.yrs.hind), 1, n.members)
    anom.weighted <- drop(anom_exp)* lat.weighted.array

    ## convert the array of anomalies in a matrix of type Member + Year + LeadDay ~ Lat + Lon;
    ## i.e: the new matrix follows the row order: member1year1forecastday1, member1year1forecastday2, ..., member1year2forecastday1,member1year2forecastday2, ... member2year1forecast1, ...    
    psl.kmeans <- array(aperm(anom.weighted,c(3,2,1,5,4)), c(n.members * n.yrs.hind * 32, n.lon*n.lat))

    ## psl.kmeans2 <- anom.weighted
    ## dim(psl.kmeans2) <- c(32 * n.yrs.hind * n.members, n.lon * n.lat)
    
    ## for the debug:
    ## test <- array(1:(2*3*4*2*2), c(2,3,4,2,2))
    ## test2 <- array(aperm(test,c(3,2,1,5,4)), c(2*3*4, 2*2))

    ## apply clustering:
    WR.for <-  kmeans(psl.kmeans, centers=n.clusters, iter.max=200, nstart=50, trace=FALSE)

    ## forecasted WR daily sequence as an array: [days x years x members]
    WR.seq.for2 <- array(WR.for$cluster, c(32, n.yrs.hind, n.members))
    
    ## get WR composites:
    ## Method A: without grouping daily sequence in an array:
    anom1 <- array(aperm(drop(anom_exp), c(3,2,1,5,4)), c(n.members * n.yrs.hind * 32, n.lon*n.lat))
    anom.cl <- array(NA, c(n.lon*n.lat, n.clusters)); for(cl in 1: n.clusters) anom.cl[,cl] <- colMeans(anom1[which(WR.for$cluster == cl),])
    WR.for.composite <- array(anom.cl, c( n.lon, n.lat, n.clusters))
                               
    ## ## Method B: with arrays (not working properly)
    ## WR.for.arr <- array(WR.for$cluster, c(32, n.yrs.hind, n.members)) # it must be in reverse order compared to the definition of psl.kmeans
    ## WR.for.arr2 <- aperm(WR.for.arr, c(3,2,1)) ## move the array to have the same order of the first three dimensions of anom.weighted
    ## WR.for.arr3 <- InsertDim(InsertDim(WR.for.arr2,4,n.lat),5,n.lon)
    ## WR.for.na <- list(); for(cl in 1:n.clusters) {  WR.for.na[[cl]] <- WR.for.arr3; WR.for.na[[cl]][which(WR.for.arr2 != cl)] <- NA; WR.for.na[[cl]][which(WR.for.arr3 == cl)] <- 1} ## replace WR number with NA or 1
    ## anom.for.cl <- list(); for(cl in 1:n.clusters) anom.for.cl[[cl]] <- drop(anom_exp)* WR.for.na[[cl]]
    ## WR.for.composite <- array(NA, c(n.lat, n.lon, n.clusters)); for(cl in 1:n.clusters) WR.for.composite[,,cl] <- apply(anom.for.cl[[cl]], c(4,5), mean, na.rm=T)
        
    ## for debug:
    ## mean(WR.for.composite[[4]])
    ## PlotLayout(PlotEquiMap,c(1,2),lon=1:186,lat=1:77,var=array(rnorm(1000),c(186,77,4)))
    ## PlotEquiMap(array(rnorm(1000),c(186,77)), lon=1:186, lat=1:77)

    ## save forecasted regimes:
    ## If PlotLayout gives a "var lon must have same number of elements ..." error, try loading s2dverification again with library(s2dverifiation)
    ##
    ## PlotLayout(PlotEquiMap,c(1,2),lon=data$lon,lat=data$lat,var=WR.for.composite/100,
    ##        titles=paste0(paste0('Cluster ',1:n.clusters)),filled.continents=F,
    ##        axelab=F,draw_separators = T,subsampleg = 1,brks=seq(-16,16,by=2),
    ##        bar_extra_labels = c(2,0,0,0),fileout=paste0(work.dir,'/forecasted_regimes_for_startdate_', sdate,'.png')) ## introduce when .SelecDevice() will be fixed

    ## spatial correlations between the 4 forecasted WRs and the observed ones:
    m <- 1
    for.linked <- array(NA, c(length(months2load), n.clusters))
    for(month in months2load){
        ## load 'reference' (=WR_obs$composite) and 'WR.seq.obs' variables with the cluster data of the same month of the forecast day considered:    
        load(paste0(work.dir,"/WRs_observed_",period.name[month],".RData"))
        
        WR.for.composite <- aperm(WR.for.composite,c(2,1,3)) ## [lat, lon, member]
        WR.obs.composite <- reference ## [lat, lon, member]
        
        ## PlotLayout(PlotEquiMap,c(1,2),lon=data$lon,lat=data$lat,var=WR.obs.composite/100,
        ##            titles=paste0(paste0('Cluster ',1:n.clusters)),filled.continents=F,
        ##            axelab=F,draw_separators = T,subsampleg = 1,brks=seq(-16,16,by=2),
        ##            bar_extra_labels = c(2,0,0,0),fileout=paste0(work.dir,'/forecasted_regimes_for_startdate_', sdate,'.png')) ## introduce 'fileout' when .SelectDevice() will be fixed in Power9

        ## correlation matrix for each month [4 observed WRs (rows) versus 4 forecasted WRs (columns)]:
        cor.mes <- array(NA, c(n.clusters, n.clusters, length(months2load)))
        for(i in 1:n.clusters) for(j in 1:n.clusters) cor.mes[i,j,m] <- cor(as.vector(WR.obs.composite[,,i]), as.vector(WR.for.composite[,,j]))
        
        ## associate each simulated cluster to one observed regime:
        max.cor <- array(NA,c(n.clusters,2))

        cor.mes2 <- cor.mes
        for(cl in 1:n.clusters){
            max.cor[cl,] <- which(cor.mes2[,,m] == max(cor.mes2[,,m]), arr.ind=T)
            cor.mes2[max.cor[cl,1],, m] <- -999 # set this row to negative values so it won't be found as a maximum anymore
            cor.mes2[,max.cor[cl,2], m] <- -999 # set this column to negative values so it won't be found as a maximum anymore
        }

        
        ## for each month, found which are the forecasted WR cluster numbers associated to the observed WR cluster numbers:
        for(cl in 1:n.clusters) for.linked[m, cl] <- max.cor[which(max.cor[,1] == cl), 2]
        m <- m + 1
        
    } ## close for on month
    
    
    WR.seq.for <- array(NA, dim(anom_exp)[1:4]) ## array which will store the cluster sequence of the predicted WRs for all the 31 forecast days of the selected sdate
    WR.seq.rean <- array(NA,c(n.yrs.hind, 32)) ## array which will store the cluster sequence of the observed WRs for all the 31 forecast days of the selected sdate
    WR.seq.rean2 <- array(NA,c(n.yrs.hind, 32)) ## as before but with the numbers of the observed WRs replaced with the numbers of the associated forecasted WRs
    m <- 1
    for(month in months2load){
        ## load 'reference' (=WR_obs$composite) and 'WR.seq.obs' variables with the cluster data of the same month of the forecast day considered:    
        load(paste0(work.dir,"/WRs_observed_",period.name[month],".RData"))
        
        ## extract only forecast anomalies of days belonging to that month:
        pos.month <- which(monthly.seq == month) ## pos.month is the number of the forecast days which belongs to the selected month
                                        
        ## reshape the daily sequence of observed monthly WRs in an array [years, days]:
        WR.seq.obs.array <- t(array(WR.seq.obs, c(n.days.in.a.month(month,2001), n.yrs.obs)))

        ## select only the reanalysis years within the hindcasts years to be able to compare the obsered WRs with the forecasted ones:
        yr1 <- yr1.hind - yr1.obs + 1
        yr2 <- yr2.hind - yr1.obs + 1

        ## load one year more to be able to assign obs WRs for forecast days beyond the 31th of December (see below)
        WR.seq.obs.array2 <- WR.seq.obs.array[yr1:(yr2+1),] 

        ## ## replace the observed cluster numbers with the forecasted cluster numbers associated to the observed WRs:
        ## WR.seq.obs.array3 <- WR.seq.obs.array2 
        ## for(cl in 1:n.clusters) WR.seq.obs.array3[,which(WR.seq.obs.array2 == cl)] <- obs2for[,,month]
        
        ndm <- n.days.in.a.month(month,2001) ## number of days in 'month'
        
        if(length(months2load) == 2) {
            if (month == months2load[1]) pos.fd <- (ndm - length(which(monthly.seq == month)) + 1) : ndm
            if (month == months2load[2]) pos.fd <- 1: length(which(monthly.seq == month))
        }
        
        if(length(months2load) == 3) {
            if (month == months2load[1]) pos.fd <- (ndm - length(which(monthly.seq == month)) + 1) : ndm
            if (month == months2load[2]) pos.fd <- 1:ndm
            if (month == months2load[3]) pos.fd <- 1: length(which(monthly.seq == month))
        }
            
        ## select only the part of WR.seq.obs.array2 inside the forecast days of the selected month:
        WR.seq.rean[,pos.month] <- WR.seq.obs.array2[1:n.yrs.hind, pos.fd]

        ## correct in case the forecast days enter in the following year:
        if(month != months2load[1] && months2load[2] < months2load[1]) WR.seq.rean[,pos.month] <- WR.seq.obs.array2[-1,pos.fd]

        for(cl in 1:n.clusters) WR.seq.rean2[,pos.month][which(WR.seq.rean[, pos.month] == cl)] <- for.linked[m, cl]
        m <- m + 1
    } ## close for on month



    ##rm(psl.kmeans, anom.+weighted, lat.weighted.array)

    ## compare the forecasted WR sequence with the observed one:
    for(d in 1:32){
        for(y in 1:n.yrs.hind){
            WR.seq.perc2[sdate,d] <- WR.seq.perc2[sdate,d] + length(which(WR.seq.for2[d,y,] == WR.seq.rean2[y,d]))
        }
    }

    
    print(WR.seq.perc2)    
} ## close for on sdate


save(WR.seq.perc2, file=paste0(work.dir,"/Forecast_table.Rdata"))

##
## Plot the figure
##

work.dir <- "/esarchive/scratch/ncortesi/Regimes/87_WRs_subseasonal"

load(paste0(work.dir,"/Forecast_table.Rdata"))

##my.cols <- c('white','#fef0d9','#fdd49e','#fdbb84','#fc8d59','#e34a33','#b30000', '#7f0000')
##my.cols <- c('#fff7fb','#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')
my.cols <- c('#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')

##my.brks <- c(-0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.30, 1)
##my.brks <- c(0,20,30,40,50,60,70,80,90,100)
my.brks <- c(0,30,40,50,60,70,80,90,100)

## ColorTable <- function(matrix, cols, brks, fileout = NULL, wifth = 8, height = 5, size_units = "in", res = 100){
##     if (!is.null(fileout)) {
##         saveToFile(fileout)
##     } else {
##         dev.new(units = size_units, res = res, width = width, height = height)
##     }
## }


my.file <- paste0(work.dir,"/Forecast_figure.png")
png(file=my.file, width=530, height=700)  
plot.new()

##par(mar = c(0, 0, 2, 0), fig = c(0.32, 0.57, 0.8, 0.94), new=TRUE)
par(fig = c(0, 0.88, 0, 1))
plot(1, 1, type = "n", xaxt = "n", yaxt = "n", ylim = c(1, 52.04), xlim = c(-0.23, 30.2), ann = F, bty="n") ## to remove black border

## Add axes labels:
mtext(side = 1, text = "                   Forecast day", line = 3, cex = 1.2, font = 3)
mtext(side = 2, text = "Start date", line = 3, cex = 1.2, font = 3)

## Add figure title:
my.title = "Coincident regimes (in %) between ECMWF-MPS and ERA-Interim\nwith the spatial correlation method"
mtext(side = 3, text = my.title, cex = 1.3)

## old x-axis that cannot be rotated of 45 degrees; only used to draw ticks:
axis(1, at=seq(0,31), las=3, cex.axis=1, labels=FALSE, mgp=c(3,0.7,0)) 

## Plot x-axis labels rotated of 45 degrees at default tick marks:
##labels <- paste0(" ",1:31)
text(0:31, par("usr")[3] - 1.55, srt = 0, adj = 1, labels = 1:32, xpd = TRUE, cex=.8) 

## Draw subtitle:
##mtext(1, text = "                                     Reference dataset: Era-Interim", line = 7)

## Plot y-axis:
axis(2, at=52:1, las=2, labels=1:52, cex.axis=.8)

## create an array similar to array.pers but with colors instead of frequencies:                
n.cols <- length(my.cols)
array.colors <- array(my.cols[n.cols], c(52,32))
for(int in (n.cols-1):1) array.colors[WR.seq.perc2 < my.brks[int+1]] <- my.cols[int]

##set to gray the last column:
##array.colors[,32] <- "#ece7f2"

## fill with colors the cells in the figure:
for(p in 1:52){ for(l in 0:(32-1)){ polygon(c(0.5+l-1, 0.5+l-1, 1.5+l-1, 1.5+l-1), c(-0.5+53-p, 0.5+53-p, 0.5+53-p, -0.5+53-p), col=array.colors[p,1+l], border="black") }}

## add color legend:
par(fig = c(0.83, 1, 0.30, 0.75), new=TRUE)
ColorBar(brks=my.brks, cols=my.cols, vert=TRUE, label_scale=1.3, subsample=1) #my.ticks=-0.5 + 1:length(my.brks.var2), my.labels=my.brks.var2) subset=my.subset2)
    
dev.off()    





#######################################################################################################
##                                                                                                   ##
##   Validation of WR forecasts at subseasonal scale with the assign method (final methodology)      ##
##                                                                                                   ##
#######################################################################################################

## important: even if Load() function is able to interpolate grid, when loading 1 data chunk, Load() loads and interpolates the whole grid instead than
## only 1 chunk, so it is 50-100 times slower than what should be.

yr1.obs <- 1979 # first year of psl data from the reanalysis
yr2.obs <- 2017 # last year of psl data from the reanalysis 
    
forecast.year <- 2014 ##2016             # starting year of the weekly sequence of the forecasts
mes <- 1                                 # starting forecast month (1: january, 2: february, etc.)
day <- 2 #7                              # starting forecast day (usually the first Thursday of January of the forecast year,
                                         # which in case of forecasts of 2014 is the 2 of January and in case of the forecasts of 2016 is the 7 of January)
                                         # (for 2014 forecasts, also Tuesday forecasts were downloaded. They start on the 7th of January)
yr1.hind <- 1994 #1996                   # first hindcast year
yr2.hind <- 2013 #2015                   # last hindcast year (usually the forecast year -1)

## 6-hourly data (not used unless you have no daily means data available on /esarchive and you have to averege them by yourself):
forecast.dir <- paste0('/esarchive/exp/ecmwf/monthly_ensforhc/6hourly/$VAR_NAME$/', forecast.year,'$MONTH$$DAY$00/$VAR_NAME$_$START_DATE$00.nc') ## forecast year: 2014
##forecast.dir <- paste0('/esarchive/exp/ecmwf/monthly_ensforhc/6hourly/$VAR_NAME$/', forecast.year,'$MONTH$$DAY$/$VAR_NAME$_$START_DATE$.nc')  ## forecast year: 2016

## daily data:
##forecast.dir <- paste0('/esarchive/exp/ecmwf/monthly_ensforhc/daily_mean/$VAR_NAME$_f6h/', forecast.year,'$MONTH$$DAY$00/$VAR_NAME$_$START_DATE$.nc') ## forecast year: 2014

n.clusters <- 4

##################################################################################################################################
## Load Subseasonal SLP hindcast data for one start date and lead time (1 week), interpolate it to the same grid of the reanalysis 
## and finally load SLP monthly reanalysis data to assign to each day an observed WR

n.yrs.hind <- yr2.hind - yr1.hind + 1
n.yrs.obs <- yr2.obs - yr1.obs + 1

## detect lon and lat of reanalysis to be able to interpolate the forecasts to the reanalysis grid (forecasts of ECMWF are 320x640, reanalysis (ERA-Int) is 256x512):
data.grid <- Load('psl', exp = NULL, 
                 obs = 'erainterim',
                 sdates = '20010101',
                 storefreq='daily',
                 leadtimemin = 1, leadtimemax = 1,
                 output = 'lonlat', nproc=1) ## load 1 day only

lon.rean <- data.grid$lon
lat.rean <- data.grid$lat
n.lon.rean <- length(lon.rean) 
n.lat.rean <- length(lat.rean)
grid.rean <- paste0('r', n.lon.rean,'x',n.lat.rean) ## [256 x 512 in case of era-interim]

sdates.seq <- weekly.seq(forecast.year, mes, day)  # load the sequence of dates corresponding to all the thursday of the year (for forecast data)
pos.bis <- which(sdates.seq == paste0(forecast.year, "0229")) # find if there is a startdate at the 29th of february and return its position in the vector sdates.seq
if(length(pos.bis) != 0) sdates.seq <- sdates.seq[-pos.bis] # if there is a start date that starts exatly on the 29th of February, remove it
n.sdates <- length(sdates.seq) ## number of start dates

WR.seq.perc <- array(0, c(n.sdates, 32)) ## array which will store the % or same WRs forecast and observed, for each day in the 31 lead times (days)
for(sdate in 1:n.sdates){
    ## note that sdate is the forecast start date, not the observed start date for the reference WR classification that starts on the 1st of January, so they differ
    print(paste0("Start date num. ",sdate, "/",n.sdates))
    
    sdates <- paste0(yr1.hind:yr2.hind,substr(sdates.seq[sdate],5,6),substr(sdates.seq[sdate],7,8)) ## same start date but for all hindcast years
    
    ## load forecasts for all the years of start date 'sdate' and ALL FORECAST DAYS associated to it (32) and interpolate them to the same resolution of the reanalysis:
    ## note that the first forecast day is not equal to startdate+1 but to startdate!!!

    system.time(data <- Load('psl', exp = list(list(path=forecast.dir)),
                             obs = NULL,
                             sdates = sdates,
                             storefreq='daily',
                             leadtimemin = 1, leadtimemax = 129, ## replace 129 with 32 when you'll load daily data directly
                             latmin = 27, latmax = 81,
                             lonmin = 274.5, lonmax = 45, output = 'lonlat',
                             grid=grid.rean, method='bilinear', nprocs=max(1, availableCores())))

    
    n.members <- dim(data$mod)[2]
    n.lat <-  dim(data$mod)[5]
    n.lon <-  dim(data$mod)[6]

    ## In case you are loading 6-hourly data, you have to convert them to daily means before:
    data.daily <- array(NA, c(1, n.members, n.yrs.hind, 32, n.lat, n.lon)) ## [dataset x members x sdate x time x lat x lon]
    my.names <- names(dim(data$mod))

    ##for(dd in 1:32) { print(paste0("Day ", dd, " of 32")); data.daily[1,1:n.members,1:n.yrs.hind, dd, 1:n.lat,1:n.lon] <-  apply(data$mod, c(1,2,3,5,6), function(x){mean(x[(dd-1)*4 + 1:4])})}
    
    ## above loop is twice faster with .Internal(mean), but there must be no NA values:
    for(dd in 1:32) { print(paste0("Day ", dd, " of 32")); data.daily[1,1:n.members,1:n.yrs.hind, dd, 1:n.lat,1:n.lon] <-  apply(data$mod, c(1,2,3,5,6), function(x){.Internal(mean(x[(dd-1)*4 + 1:4]))})}

    data$mod <- data.daily
    names(dim(data$mod)) <- my.names ## re-assign the dimension names because if they were deleted by the previous command

    rm(data.daily)

    ## remove the eventual forecasts of the 29th of February: (they are present for the forecasts of 2016, but not of 2014):

    
    ## measure forecast climatology for the 32 days loaded:
    clim <- Clim(var_exp = data$mod, var_obs = data$mod, memb=F)

    ## apply the LOESS filter to the 32-days climatology:
    clim_smoothed_exp <- aperm(apply(clim$clim_exp, c(1,3,4), Loess, loess_span=1), c(2,1,3,4))

    ## measure forecast anomalies:
    anom_exp <- Ano(data$mod, clim_smoothed_exp)

    ## for each forecast day of the start date 'sdate', identify to which week of the year (starting from the 1st of January) it belongs:

    ## forecast days starting counting from the SAME MONTH of the start date
    forecast.daily.seq <- as.numeric(substr(sdates.seq[sdate],7,8)) + 1:32 - 1 
    
    ## month of the currently selected start date:
    month.sdate <- as.numeric(substr(sdates.seq[sdate],5,6))
    
    ## get the number of days BEFORE the month of the start date (for no leap years):
    if(month.sdate > 1) n.days.previous <- n.days.in.a.monthly.period(1, month.sdate - 1, 2001) else n.days.previous <- 0

    ## get the sequence of forecast days counting from the START of the YEAR:
    days.seq <- n.days.previous + forecast.daily.seq

    ## get the sequence of month numbers corresponding to the sequence of forecast days:
    monthly.seq <- seq.months.in.a.year(2001)[days.seq]

    ## if the forecast days extend to the following year, assign them to January:
    if(length(which(is.na(monthly.seq))) > 0) monthly.seq[which(is.na(monthly.seq))] <- 1 

    ## get the sequence of weekly numbers corresponding to the sequence of forecast days:
    ## (notice that the last week of the year is the number 53 and it is made up by 1 day ony, so it isn't used in the WR classification)
    weekly.seq <- seq.weeks.in.a.year(2001)[days.seq]
    
    ## if the forecast days extend to the following year, assign them to January's weeks:
    if(length(which(is.na(weekly.seq))) > 0) weekly.seq[which(is.na(weekly.seq))] <- rep(1:5, each=7)[1:length(which(is.na(weekly.seq)))]
    
    ## check all the sequences (for the debug):
    cbind(1:length(days.seq),days.seq,monthly.seq) 

    ## identify which month(s) are to load:
    months2load <- unique(monthly.seq)

    WR.seq.for <- array(NA, dim(anom_exp)[1:4]) ## array which will store the cluster sequence of the predicted WRs for all the 32 forecast days of the selected sdate
    WR.seq.rean <- array(NA, c(n.yrs.hind, 32)) ## array which will store the cluster sequence of the observed WRs for all the 32 forecast days of the selected sdate
    
    for(month in months2load){
        ## load 'reference' and 'WR.seq.obs' variables with the cluster data of the same month of the forecast day considered:    
        load(paste0(work.dir, "/WRs_observed_", period.name[month], ".RData"))
        
        ## extract only forecast anomalies of days belonging to that month:
        pos.month <- which(monthly.seq == month) ## pos.month is the number of the forecast days which belongs to the selected month
        anom_exp_same_month <- anom_exp[,,,pos.month,,] ## anom_exp are the forecasted anomalies for the 32 forecast days 
                                        
        ## measure distance of each forecast day to the reference month(s):
        WR.seq.for[,,,pos.month] <- RegimesAssign(var_ano = anom_exp_same_month, ref_maps = reference, lats = data$lat, method='distance')$cluster
        ##   WR_exp <- RegimesAssign(var_ano=anom_exp,ref_maps=reference,lats=data$lat,method='distance')

        ## reshape the daily sequence of observed monthly WRs in an array [years, days]:
        WR.seq.obs.array <- t(array(WR.seq.obs, c(n.days.in.a.month(month, 2001), n.yrs.obs)))

        ## select only the reanalysis years within the hindcasts years to be able to compare the obsered WRs with the forecasted ones:
        yr1 <- yr1.hind - yr1.obs + 1
        yr2 <- yr2.hind - yr1.obs + 1

        WR.seq.obs.array2 <- WR.seq.obs.array[yr1:(yr2+1),] ## load one year more to be able to assign obs WRs for forecast days beyond the 31th of December (see below)

        ndm <- n.days.in.a.month(month,2001) ## number of days in 'month'
        
        if(length(months2load) == 2) {
            if (month == months2load[1]) pos.fd <- (ndm - length(which(monthly.seq == month)) + 1) : ndm
            if (month == months2load[2]) pos.fd <- 1: length(which(monthly.seq == month))
        }
        
        if(length(months2load) == 3) {
            if (month == months2load[1]) pos.fd <- (ndm - length(which(monthly.seq == month)) + 1) : ndm
            if (month == months2load[2]) pos.fd <- 1:ndm
            if (month == months2load[3]) pos.fd <- 1:length(which(monthly.seq == month))
        }
            
        ## select only the part of WR.seq.obs.array2 inside the forecast days of the selected month:
        WR.seq.rean[,pos.month] <- WR.seq.obs.array2[1:n.yrs.hind, pos.fd]

        ## correct in case the forecast days enter in the following year:
        if(month != months2load[1] && months2load[2] < months2load[1]) WR.seq.rean[,pos.month] <- WR.seq.obs.array2[-1,pos.fd]
            
        ## check forecast day 1 to see if the forecasts correspond or not (they should be 80-90% similar):
        ## cbind(obs=WR.seq.rean[,1], for1=WR.seq.for[1,1,,1], for2=WR.seq.for[1,2,,1], for3=WR.seq.for[1,3,,1], for4=WR.seq.for[1,4,,1])
        
    } ## close for on month

    ## compare the forecasted WR sequence with the observed one:
    for(d in 1:32){
        for(y in 1:n.yrs.hind){
            WR.seq.perc[sdate,d] <- WR.seq.perc[sdate,d] + length(which(WR.seq.for[1,,y,d] == WR.seq.rean[y,d]))
        }
    }

    print(WR.seq.perc)    
} ## close for on sdate

## convert it to a %:
WR.seq.perc <- 100 * WR.seq.perc / (n.yrs.hind * n.members)

save(WR.seq.perc, file=paste0(work.dir,"/Assign_table.Rdata"))


##
## Plot the figure
##

load(paste0(work.dir,"/Assign_table.Rdata"))

##my.cols <- c('white','#fef0d9','#fdd49e','#fdbb84','#fc8d59','#e34a33','#b30000', '#7f0000')
##my.cols <- c('#fff7fb','#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')
my.cols <- c('#ece7f2','#d0d1e6','#a6bddb','#74a9cf','#3690c0','#0570b0','#045a8d','#023858')

##my.brks <- c(-0.05, 0, 0.05, 0.1, 0.15, 0.2, 0.25, 0.30, 1)
##my.brks <- c(0,20,30,40,50,60,70,80,90,100)
my.brks <- c(0,30,40,50,60,70,80,90,100)

## ColorTable <- function(matrix, cols, brks, fileout = NULL, wifth = 8, height = 5, size_units = "in", res = 100){
##     if (!is.null(fileout)) {
##         saveToFile(fileout)
##     } else {
##         dev.new(units = size_units, res = res, width = width, height = height)
##     }
## }


my.file <- paste0(work.dir,"/Assign_figure.png")
png(file=my.file, width=530, height=700)  
plot.new()

##par(mar = c(0, 0, 2, 0), fig = c(0.32, 0.57, 0.8, 0.94), new=TRUE)
par(fig = c(0, 0.88, 0, 1))
plot(1, 1, type = "n", xaxt = "n", yaxt = "n", ylim = c(1, 52.04), xlim = c(-0.23, 30.2), ann = F, bty="n") ## to remove black border

## Add axes labels:
mtext(side = 1, text = "                   Forecast day", line = 3, cex = 1.2, font = 3)
mtext(side = 2, text = "Start date", line = 3, cex = 1.2, font = 3)

## Add figure title:
my.title = "Coincident regimes (in %) between ECMWF-MPS and ERA-Interim\nwith the assign method"
mtext(side = 3, text = my.title, cex = 1.3)

## old x-axis that cannot be rotated of 45 degrees; only used to draw ticks:
axis(1, at=seq(0,31), las=3, cex.axis=1, labels=FALSE, mgp=c(3,0.7,0)) 

## Plot x-axis labels rotated of 45 degrees at default tick marks:
##labels <- paste0(" ",1:31)
text(0:31, par("usr")[3] - 1.55, srt = 0, adj = 1, labels = 1:32, xpd = TRUE, cex=.8) 

## Draw subtitle:
##mtext(1, text = "                                     Reference dataset: Era-Interim", line = 7)

## Plot y-axis:
axis(2, at=52:1, las=2, labels=1:52, cex.axis=.8)

## create an array similar to array.pers but with colors instead of frequencies:                
n.cols <- length(my.cols)
array.colors <- array(my.cols[n.cols], c(52,32))
for(int in (n.cols-1):1) array.colors[WR.seq.perc < my.brks[int+1]] <- my.cols[int]

##set to gray the last column:
##array.colors[,32] <- "#ece7f2"

## fill with colors the cells in the figure:
for(p in 1:52){ for(l in 0:(32-1)){ polygon(c(0.5+l-1, 0.5+l-1, 1.5+l-1, 1.5+l-1), c(-0.5+53-p, 0.5+53-p, 0.5+53-p, -0.5+53-p), col=array.colors[p,1+l], border="black") }}

## add color legend:
par(fig = c(0.83, 1, 0.30, 0.75), new=TRUE)
ColorBar(brks=my.brks, cols=my.cols, vert=TRUE, label_scale=1.3, subsample=1) #my.ticks=-0.5 + 1:length(my.brks.var2), my.labels=my.brks.var2) subset=my.subset2)
    
dev.off()    



#####################################################################################################
#                                          WR vignette                                              #
#####################################################################################################

sdates <- paste0(1990:2016, '1201')

data <- Load('psl', exp = 'system4_m1',
             obs = 'erainterim',
             sdates = sdates,
             storefreq='daily',
             leadtimemin = 1, leadtimemax = 31,
             latmin = 27, latmax = 81,
             lonmin = 274.5, lonmax = 45, output = 'lonlat', nprocs=max(1,availableCores()-1))

clim <- Clim(var_exp = data$mod,var_obs = data$obs,memb=F)

clim_smoothed_obs <- aperm(apply(clim$clim_obs,c(1:length(dim(clim$clim_obs)))[-which(names(dim(clim$clim_obs))=='ftime')], Loess, loess_span=1),c(2,1,3,4))
clim_smoothed_exp <- aperm(apply(clim$clim_exp,c(1:length(dim(clim$clim_exp)))[-which(names(dim(clim$clim_exp))=='ftime')], Loess, loess_span=1),c(2,1,3,4))

anom_obs <- Ano(data$obs,clim_smoothed_obs)
anom_exp <- Ano(data$mod,clim_smoothed_exp)

WR_obs <- WeatherRegime(data = anom_obs, EOFS = FALSE, lat = data$lat, ncenters = 4)
names(dim(WR_obs$composite)) <- c("lat", "lon", "cluster", "", "")

clim_frequencies <- paste0('freq = ',round(Mean1Dim(WR_obs$frequency,1),1),'%')

PlotLayout(PlotEquiMap, c(2,1), lon=data$lon, lat=data$lat, var=WR_obs$composite/100,
           titles=paste0(paste0('Cluster ',1:4),' (',clim_frequencies,' )'),filled.continents=F,
           axelab=F,draw_separators = T,subsampleg = 1,brks=seq(-16,16,by=2),
           bar_extra_labels = c(2,0,0,0)) ##,fileout='observed_regimes.png')

##
## WR vignette without the forecast data: (faster to load)
##

sdates <- paste0(1979:2017, '1001')

data <- Load('psl', exp = NULL,
             obs = 'erainterim',
             sdates = sdates,
             storefreq='daily',
             leadtimemin = 1, leadtimemax = 31,
             latmin = 27, latmax = 81,
             lonmin = 274.5, lonmax = 45, output = 'lonlat', nprocs=1)

clim <- Clim(var_exp = data$obs, var_obs = data$obs, memb=F)
clim_smoothed_obs <- aperm(apply(clim$clim_obs,c(1:length(dim(clim$clim_obs)))[-which(names(dim(clim$clim_obs))=='ftime')], Loess, loess_span=1),c(2,1,3,4))
anom_obs <- Ano(data$obs, clim_smoothed_obs)

WR_obs <- WeatherRegime(data = anom_obs, EOFS = FALSE, lat = data$lat, ncenters = 4)
names(dim(WR_obs$composite)) <- c("lat", "lon", "cluster", "", "")

## save weather regimes for Marta Almarcha:
my.month <- "October"
write.table(WR_obs$cluster, paste0("~/WRs_", my.month,"_1981-2017.txt"))
write.table(WR_obs$frequency, paste0("~/WRs_", my.month,"_1981-2017_freq.txt"))

save(WR_obs, file = paste0("~/WRs_", my.month, "_1981-2017.RData"))



