# Veronica Torralba (CMCC)
# 10th November 2022
# Medcof training - AEMET

# cleaning:
rm(list=ls())

# To run this example several libraries are required
# install.packages('CSTools', 'zeallot', 's2dv','ClimProjDiags')

# load libraries:
library(CSTools)
library(s2dv)
library(ClimProjDiags)
library(zeallot)

# 1) Load data
# In this case, the seasonal forecasts of temperature 
# initialized in May from the CMCC and ECMWF seasonal forecast systems will 
# be used. ERA5 reanalysis will be used as observational reference 
# for calibration

# define paths:
path <- '/home/tr100/CST_Calibration/data_storage'

exp_path <- list(
  list(name = 'cmcc',
       path = paste0(
         path,
         '/exp/cmcc/sps35/monthly_mean/$VAR_NAME$/$VAR_NAME$_$START_DATE$.nc')),
  list(name = 'ecmwf',
       path = paste0(
         path,
         '/exp/ecmwf/seas5/monthly_mean/$VAR_NAME$/$VAR_NAME$_$START_DATE$.nc')))

obs_path <- list(
  list(name = 'era5',
       path = paste0(
         path,
         '/recon/ecmwf/era5/monthly_mean/$VAR_NAME$/$VAR_NAME$_$YEAR$$MONTH$.nc')))

# The seasonal forecasts from the two systems (CMCC_SPS35, ECMWF_SEAS5) 
# and the observational reference (ERA5) are stored in two objects: exp and obs
# CST_Load provides corresponding forecasts and observations
# for a common  period and grid 

c(exp,obs) %<-% CST_Load(var = 'tas',
                         exp=exp_path,
                         obs=obs_path,
                         nmember=25,
                         sdates = paste0(2000:2016, '0501'),
                         storefreq ='monthly',
                         leadtimemin = 2, leadtimemax = 4,
                         latmin = 27, latmax = 48,
                         lonmin = -12, lonmax = 40,
                         output = 'lonlat',nprocs=1)

#### to be used in case the Load function is slow #########
# data <- readRDS('~/CST_Calibration/CST_Calibration_data_2000_2016.RDS')
# exp <- data$exp
# obs <- data$obs

# Explore the dimensions
print(dim(exp$data))
print(dim(obs$data))

# 2) Seasonal mean
# The seasonal mean of both forecasts and observations
# are computed by averaging over the 'ftime' dimension.
exp$data <- MeanDims(exp$data, 'ftime',na.rm = T)
obs$data <- MeanDims(obs$data, 'ftime',na.rm = T)

# 3) Calibration
# The forecasts from each seasonal prediction system are calibrated 
# with the ERA5 observational reference. 

# We will do this calibration separately for each model. 
exp_cmcc <- exp_ecmwf <- exp
exp_cmcc$data <- Subset(exp$data,along='dataset',indices=1)
exp_ecmwf$data <- Subset(exp$data,along='dataset',indices=2)

exp_cal_cmcc <- CST_Calibration(exp = exp_cmcc,
                  obs = obs,
                  cal.method = 'mse_min')

exp_cal_ecmwf <- CST_Calibration(exp = exp_ecmwf,
                  obs = obs,
                  cal.method = 'mse_min')

# 4) Visualisation

# Bias of the raw forecasts
bias_raw_cmcc <- Bias(exp=exp_cmcc$data,obs=obs$data,
                      memb_dim='member',na.rm=T)
bias_raw_ecmwf <- Bias(exp=exp_ecmwf$data,obs=obs$data,
                       memb_dim='member',na.rm=T)

# Bias of the calibrated forecasts
bias_cal_cmcc <- Bias(exp=exp_cal_cmcc$data,obs=obs$data,
                      memb_dim='member',na.rm=T)
bias_cal_ecmwf <- Bias(exp=exp_cal_ecmwf$data,obs=obs$data,
                       memb_dim='member',na.rm=T)

# Plot biases 
PlotLayout('PlotEquiMap',plot_dims = c("lat", "lon"),
           var=list(bias_raw_cmcc,
                    bias_raw_ecmwf,
                    bias_cal_cmcc,
                    bias_cal_ecmwf),
           lon=exp$lon,lat=exp$lat,drawleg="S",
           titles = c('CMCC-SPS35',
                      'ECMWF-SEAS5',
                      'CMCC-SPS35 calibrated',
                      'ECMWF-SEAS5 calibrated'),
           filled.continents = F,brks=seq(-5.5,5.5,by=1),
           color_fun = clim.palette("bluered"),
           draw_separators = T,
           bar_label_scale = 1.5,
           toptitle='Mean bias (K) in JJA 2000-2016',
           title_scale = 0.8,
           bar_extra_margin = c(2,0,0,0),
           fileout = '~/CST_Calibration/bias.ps',
           country.borders=TRUE)



# ---------------
# Bonus PLOT 
# ---------------

cal_boxplot <- function(data_exp,data_cal,data_obs,la,lo,years=NULL,
                        tit='',
                        yl=NULL,ylabel='Temperature (K)'){
  # Arguments:
  # ----------
  # data_exp: array of experimental data, data_cal: array of calibrated data
  # data_obs: array of observational data
  # la: latitude index , lo: longitude index,  years: years for the timeseries

  par(cex.lab=1.2)
  col1 <- adjustcolor( "red", alpha.f = 0.2)
  col2 <- adjustcolor( "blue", alpha.f = 0.2)
  if (is.null(years)){
    years <- 1:dim(data_exp)['sdate']
  }
  ts_obs <- drop(Subset(data_obs,along=c('lat','lon'),indices=list(la,lo)))
  ts_exp <- drop(Subset(data_exp,along=c('lat','lon'),indices=list(la,lo)))
  ts_cal <- drop(Subset(data_cal,along=c('lat','lon'),indices=list(la,lo)))
  
  boxplot(ts_exp,names=years,main=tit,
          las=2,col=col1,ylim=yl, ylab=ylabel)
  boxplot(ts_cal,axes=F,
          las=2,col=col2,ylim=yl,ylab=ylabel,add=T)
  lines(ts_obs,col='black',lwd=2)
  
  legend("top", fill =c(col1,col2,NA),
         border=c(col1,col2,NA),
         cex=1, bty='n',legend = c("raw","calibrated",'obs'),
         horiz = TRUE, lty = c(NA,NA,1), col=c(NA,NA,"black"))
}

# Plot the time evolution of both raw and calibrated 
# forecasts compared to observations for one specific grid point

cairo_ps('~/CST_Calibration/timeseries_CMCC_ECMWF_calibrated.ps',width = 8,height = 12)
layout(matrix(c(1,2),ncol=1,byrow=T))
cal_boxplot(data_exp=exp_cmcc$data,
            tit = 'CMCC-SPS35',
            data_cal=exp_cal_cmcc$data,
            data_obs=obs$data,
            la=3, lo=10, years=2000:2016,yl=c(285,305))
cal_boxplot(data_exp=exp_ecmwf$data,
            data_cal=exp_cal_ecmwf$data,
            data_obs=obs$data,
            tit = 'ECMWF-SEAS5',
            la=3, lo=10, years=2000:2016,yl=c(285,305))
dev.off()

