# =========================================================================

# From the beginning to end
# load data
# compute anomalies
# compute weights based on BOCM and correlations
# compute weights based on CAWM and correlations
# compute adjusted p-values
# plot tas or prlr

# =========================================================================


#source("YOUR_PATH/funcs/loadData.R") ##### BUILT_IN FUNCTION INSIDE funcs FOLDER

# use above function to load data
# requires following arguments:
#                             climate variable -> "tas" or "prlr"
#                                        month -> startdate 
#                                                 (MUST BE of type ***STRING***)


NovStartData <- loadData(clim_var = "tas", mth = "11") 
MayStartData <- loadData(clim_var = "tas", mth = "05")


PrNovStartData <- loadData(clim_var = "prlr", mth = "11")
PrMayStartData <- loadData(clim_var = "prlr", mth = "05") 


Lat <- NovStartData$lat
Lon <- NovStartData$lon
PrLat <- PrNovStartData$lat
PrLon <- PrNovStartData$lon


# --------------------------------------------------------------------------------------------




#source("YOUR_PATH/funcs/ano_corr.R") ##### BUILT_IN FUNCTION INSIDE funcs FOLDER

# use above function to calculate 
#                                 1) anomalies
#                                 2) correlation of individual, 
#                                 3) correlation of MMM
#                                 3) correlation of WMM
# requires following arguments:
#                              Data -> must be an R object from s2dverification's load fnc
#                                      i.e. must have Data$mod and Data$obs 


# tas
AnomDJF <- ano_corr(NovStartData)[[1]]
AnomJJA <- ano_corr(MayStartData)[[1]]

AnomPrDJF <- ano_corr(PrNovStartData)[[1]]
AnomPrJJA <- ano_corr(PrMayStartData)[[1]]


# --------------------------------------------------------------------------------------------




# source("YOUR_PATH"/funcs/corr_WMM_BOCM.R") ##### BUILT_IN FUNCTION INSIDE funcs FOLDER

# use above function to calculate 
#                                 1) predictions based on correlation used as weights in cv mode
#                                 2) correlation of the above prediction with obs dataset
# requires following arguments:
#                              Data -> must be an R object from s2dverification's anomalies fnc
#                                      i.e. must have Data$ano_exp and Data$ano_obs


corrDJF_WMM_BOCM <- corr_WMM_BOCM(AnomDJF)[[2]]
corrJJA_WMM_BOCM <- corr_WMM_BOCM(AnomJJA)[[2]]

corrPrDJF_WMM_BOCM <- corr_WMM_BOCM(AnomPrDJF)[[2]]
corrPrJJA_WMM_BOCM <- corr_WMM_BOCM(AnomPrJJA)[[2]]


# --------------------------------------------------------------------------------------------




#source("YOUR_PATH/funcs/corr_WMM_CorrWgts.R") ##### BUILT_IN FUNCTION INSIDE funcs FOLDER

# use above function to calculate 
#                                 1) predictions based on correlation used as weights in cv mode
#                                 2) correlation of the above prediction with obs dataset
# requires following arguments:
#                              Data -> must be an R object from s2dverification's anomalies fnc
#                                      i.e. must have Data$ano_exp and Data$ano_obs


corrDJF_WMM_CAWM <- corr_WMM_CorrWgts(AnomDJF)[[2]]
corrJJA_WMM_CAWM <- corr_WMM_CorrWgts(AnomJJA)[[2]]


corrPrDJF_WMM_CAWM <- corr_WMM_CorrWgts(AnomPrDJF)[[2]]
corrPrJJA_WMM_CAWM <- corr_WMM_CorrWgts(AnomPrJJA)[[2]]


# --------------------------------------------------------------------------------------------




# SPECIFY WHICH CLIMATE VARIABLE BELOW
clim_var <- "tas" 
#clim_var <- "prlr"



if (clim_var == "tas") {
  n <- "8"
  name <- "Temperature"
  
  corrDJF_MMM <- corrDJF_MMM[1,1,,,]
  corrDJF_MMM <- corrDJF_MMM[1,1,,,] 
  
  corrDJF_WMM_bocm <- corrDJF_WMM_BOCM[1,1,,,] 
  corrJJA_WMM_bocm <- corrJJA_WMM_BOCM[1,1,,,] 
  
  corrDJF_WMM_cawm <- corrDJF_WMM_CAWM[1,1,,,] 
  corrJJA_WMM_cawm <- corrJJA_WMM_CAWM[1,1,,,] 
  
  Lat <- Lat
  Lon <- Lon
  
} else {
  n <- "9"
  name <- "Precipitation"
  
  corrDJF_MMM <- corrPrDJF_MMM[1,1,,,]
  corrDJF_MMM <- corrPrDJF_MMM[1,1,,,] 
  
  corrDJF_WMM_bocm <- corrPrDJF_WMM_BOCM[1,1,,,] 
  corrJJA_WMM_bocm <- corrPrJJA_WMM_BOCM[1,1,,,] 
  
  corrDJF_WMM_cawm <- corrPrDJF_WMM_CAWM[1,1,,,] 
  corrJJA_WMM_cawm <- corrPrJJA_WMM_CAWM[1,1,,,] 
  
  Lat <- PrLat
  Lon <- PrLon
  
}


# -------------------------------------------------------------------------




# function to obtain adjusted p-values

adj_pval <- function(matrix, nRow, nCol){
  vec <- c(matrix)
  N <- length(vec) - sum(is.na(matrix))
  pval <- p.adjust(vec, "fdr", n = N)
  matrix <- matrix(pval, nRow, nCol)
  return(matrix)
}
NROW <- length(Lat)
NCOL <- length(Lon)


# -------------------------------------------------------------------------


# rearrange dataset
# equally weighted, BOCM weights, CAWM weights (in that order)

corrDJF <- abind::abind(corrDJF_MMM, corrDJF_WMM_bocm, corrDJF_WMM_cawm, along = 0) # winter
corrJJA <- abind::abind(corrJJA_MMM, corrJJA_WMM_bocm, corrJJA_WMM_cawm, along = 0) # summer




# -------------------------------------------------------------------------


# plot

# colors
int <- seq(-1, 1, 0.1)
ncols <- length(int) -1
int.col <- c(hcl(240, l=seq(20,99,length=ncols/2), c=seq(70,30,length=ncols/2)),
             hcl(10, l=seq(99,20,length=ncols/2), c=seq(30,70,length=ncols/2)))


#setwd("YOUR_DIRECTORY_TO_SAVE_OUTPUT") 

postscript(paste(n, "_MMM_vs_WMM_", clim_var, ".ps", sep = ''))

par(mar=c(0,1,0,0))
layout(matrix(c(1,1,1,1,
                2,3,4,5,
                6,7,8,9,
                10,11,12,13,
                14,15,15,15), 5, 4, byrow = TRUE),
       respect = TRUE, widths=c(.75,3,3,3), heights=c(.65,.65, 3,3,1))


#heading
plot.new()
title=paste("Correlation of Weighted Multi-Models for Seasonal", name, "", sep = ' ')
text(0.5,0.5, title, cex=1.5,font=2)

# subheading  
plot.new()
text(0.5,0.5," ", cex=1.25,font=2)
plot.new()
text(0.5,0.5,"MME \n(Equally Weighted)", cex=1.25,font=2)
plot.new()
text(0.5,0.5,"BOCM \n(Unequally Weighted)", cex=1.25,font=2)
plot.new()
text(0.5,0.5,"CAWM \n(Unequally Weighted)", cex=1.25,font=2)


library(s2dverification)

# winter
plot.new()
text(0.5,0.5,"W\nI\nN\nT\nE\nR", cex=1.25,font=2)

#loop over the models
for (imod in 1:3)
{
  # dot to mark significance
  dot <- as.matrix(corrDJF[imod,2,,]>adj_pval(corrDJF[imod,4,,], NROW, NCOL))
  dot[is.na(dot)] <- FALSE
  
  # plot
  PlotEquiMap(corrDJF[imod,2,,], Lon, Lat,
              brks = int, cols = int.col, colNA = 'white',
              drawleg = FALSE, numbfig = 6, axelab = FALSE,
              filled.continents = FALSE, square = TRUE, dots=dot)
}


# summer
plot.new()
text(0.5,0.5,"S\nU\nM\nM\nE\nR", cex=1.25,font=2)

for (imod in 1:3)
{
  # dot to mark significance
  dot <- as.matrix(corrJJA[imod,2,,]>adj_pval(corrJJA[imod,4,,], NROW, NCOL))
  dot[is.na(dot)] <- FALSE
  
  # plot
  PlotEquiMap(corrJJA[imod,2,,], Lon, Lat,
              brks = int, cols = int.col, colNA = 'white',
              drawleg = FALSE, numbfig = 6, axelab = FALSE,
              filled.continents = FALSE, square = TRUE, dots=dot)
}


# color bar
plot.new()
text(0.5,0.5," ", cex=1.25,font=2)
#source("YOUR_PATH/3.0.ColorBar.R")
ColorBar(int, cols=int.col, vert=FALSE, cex = .5, tick=FALSE, 
         upperspace=0.6, lowerspace = .6) 

dev.off()
