# This script runs the calibration function using "compute"

# 1. Load libraries

library(startR)
library(s2dv)
library(multiApply)
library(easyVerification)
library(s2dverification)
library(CSTools)

# 2. Load data

# Select the number of cores needed 
ncores <- 8

# Select "retrieve <- FALSE" when using "compute"
retrieve <- FALSE

# Specify the longitude and latitude boundaries of the region of interest
lons.min <- 0
lons.max <- 49
lats.min <- -40
lats.max <- 10

# Select variable
var_name = 'tasmax'

# Array containing the calibration methods
methods <- c('raw','bias','evmos','mse_min','crps_min','rpc-based')

# Directory where the forecast files are kept
exp_path <- '/esarchive/exp/ecmwf/system5c3s/monthly_mean/$var$_f24h/$var$_$sdate$01.nc'


# Specify number of categories for the RPSS calculation
n_cat <- 3

# Specify threshold for statistical significance
alpha <- 0.05
time_dim = 'sdate'
member_dim = 'ensemble'


# HOMOGENEOUS GRID. 2 options:

#  Option 1:
# Specify one file of the forecast (System5). The grid of the reanalysis (ERA5)
# will be converted to the forecast's grid to allow the calibration,
# as both datasets must have the same number and structure of grid points.
#grid <- '/esarchive/exp/ecmwf/system5c3s/monthly_mean/prlr_s0-24h/prlr_20000501.nc'

#  Option 2:
# uncomment the line below to transform onto a specific grid:
grid <- 'r144x72' 

# Load the forecast
exp <- startR::Start(
                     # Select the path of the forecast
                     dat = exp_path,
                     # Variable of interest
                     var = var_name,
                     # Start dates, years that we are using to calibrate the forecast with reanalysis
                     sdate = c("200011","200111","200211","200311"),
                     # Select all ensemble members
                     ensemble = 'all',
                     # Forecast time
                     time = indices(1:6),
                     latitude = values(list(lats.min, lats.max)),
                     latitude_reorder = Sort(decreasing = T),
                     longitude = values(list(lons.min, lons.max)),
                     # Reorder longitude points from [0,360] to [-180, 180]
                     longitude_reorder = CircularSort(-180,180),
                     transform = CDORemapper,
		             transform_extra_cells = 8,
                     transform_params = list(grid = grid,
                                             method = 'conservative',
                                             crop = c(lons.min, lons.max, lats.min, lats.max)),
                     transform_vars = c('latitude', 'longitude'),
                     synonims = list(latitude = c('lat', 'latitude'),
                                     longitude = c('lon', 'longitude')),
                     return_vars = list(latitude = 'dat',
                                        longitude = 'dat',
                                        time = 'sdate'),
                    retrieve = retrieve)


# Select the dates of interest that you will extract from the reanalysis.
# They will be the same as in the forecast.
dates <- attr(exp, 'Variables')$common$time
obs_path <- '/esarchive/recon/ecmwf/era5/monthly_mean/$var$_f24h/$var$_$date$.nc'
#dates <- format(dates, '%Y%m')
dates_file <- format(dates, '%Y%m')

# Select the number of years (sdate) and the forecast time (time) to extract from the reanalysis
# They must be the same as in the forecast
dim(dates_file) <- list(sdate = 4,
                  time = 6)
obs <- startR::Start(# load observational (reanalysis) data
             dat = obs_path,
             # Variable of interest
             var = var_name,
             date = dates_file,
             # time = values(dates),
             latitude = values(list(lats.min, lats.max)),
             latitude_reorder = Sort(decreasing = T),
             longitude = values(list(lons.min, lons.max)),
             longitude_reorder = CircularSort(-180,180),
             synonims = list(longitude = c('lon', 'longitude'),
                             latitude = c('lat', 'latitude')),
             # Convert the grid of the reanalysis onto the grid of the forecast
             # to make them comparable
             transform = CDORemapper,
             transform_extra_cells = 8,
             transform_params = list(grid = grid,
                                     method = 'conservative',
                                     crop = c(lons.min, lons.max, lats.min, lats.max)),
             transform_vars = c('latitude', 'longitude'),
             split_multiselected_dims = TRUE,
             return_vars = list(time = 'date',
                                latitude = 'dat',
                                longitude = 'dat'),
            retrieve = retrieve)

# 3. RUN THE CALIBRATION FUNCTION

# Calibration function. Needs to be fed with:
    # exp = the forecast that we previously loaded
    # obs = the reanalysis that we previously loaded
    # methods = a list with the calibration methods of interest
    # alpha = Threshold for statistical significance
    # n_cat = Number of categories for the RPSS

skill_calibration <- function(exp, obs, methods, alpha, n_cat, ncores) {
    library(s2dv)
    library(SpecsVerification)
    library(CSTools)
    library(easyVerification)
    time_dim = 'sdate'
    member_dim = 'ensemble'

   # The function "correlation_eno" calculates the anomalous correlation coefficient (ACC),
    # the difference in correlation between the forecast and observations in the calibrated vs. non calibrated forecast
  correlation_eno <- function(exp, obs, time_dim = time_dim, member_dim = member_dim, alpha = 0.05){
    exp <- s2dv::MeanDims(data = exp, dims = member_dim, na.rm = FALSE)
    cor = NULL
    cor$r = cor(exp,obs) # Correlation coefficient
    
    n_eff = s2dv::Eno(data = obs, time_dim = time_dim, na.action = na.pass, ncores = 1)
    t_alpha2_n2 = qt(p=alpha/2, df = n_eff-2, lower.tail = FALSE)
    t = abs(cor$r) * sqrt(n_eff-2) / sqrt(1-cor$r^2)
    
    if (anyNA(c(t,t_alpha2_n2)) == FALSE & t >= t_alpha2_n2){
      cor$sign = TRUE
    } else {
      cor$sign = FALSE
    }
    
    return(cor)
  }
  
  exp <- s2dv::MeanDims(data = exp, dims = 'time', na.rm = FALSE)
  obs <- s2dv::MeanDims(data = obs, dims = 'time', na.rm = FALSE)
  
  # Create empty lists:
  metrics  <- list(acc = NULL) # add the anomaly correlation coefficients and statistical significance of each method
  calibrated <- list() # add calibrated forecasts
  my.EnsRpss <-list() # add RPSS
  crpss <- list() # add CRPSS

  # Loop over different calibration methods  
  for (method in methods){
    
    if (!identical(method,'raw')){
      # Calibration function
      calibrated[[method]] <- CSTools::Calibration(exp = exp, obs = obs, cal.method = method, eval.method = "leave-one-out",
                                                   multi.model = FALSE, na.fill = TRUE, na.rm = FALSE, memb_dim = 'ensemble', 
                                                   sdate_dim = time_dim, apply_to = 'sign', alpha = 0.1, ncores = ncores)
         
    } else { # if "raw" is selected, return the raw forecast
      calibrated[[method]] <- exp
    }

    # Anomaly correlation coefficient (ACC)
    metrics$acc[[method]] <- correlation_eno(exp = calibrated[[method]], obs = obs, time_dim = time_dim, 
                                             member_dim = 'ensemble', alpha = alpha)
    
    # CRPSS
    crpss[[method]] <-easyVerification::veriApply("EnsCrpss",fcst=calibrated[[method]],obs=obs,
                                       tdim=which(names(dim(exp)) == time_dim), ensdim=which(names(dim(exp)) == member_dim), na.rm=TRUE)[[1]]

    # RPSS
    my.EnsRpss[[method]] <-easyVerification::veriApply("EnsRpss", fcst=calibrated[[method]], obs=obs,
                           prob=(1:(n_cat-1))/n_cat, tdim=which(names(dim(exp)) == time_dim), ensdim=which(names(dim(exp)) == member_dim), 
                            na.rm=TRUE)[[1]] 

      
}      
  # Return list with the calibrated data, it's ACC, RPSS and CRPSS
  return(list(raw_calib = s2dv::Reorder(data = calibrated$raw, order = c(time_dim,member_dim)),
	      acc_raw_r = metrics$acc$raw$r,
              acc_raw_sign = metrics$acc$raw$sign,
              raw_rpss = my.EnsRpss$raw,
              raw_crpss = crpss$raw,
              bias_calib = s2dv::Reorder(data = calibrated$bias, order = c(time_dim,member_dim)),
	      acc_bias_r = metrics$acc$bias$r,
              acc_bias_sign = metrics$acc$bias$sign,
              bias_rpss = my.EnsRpss$bias,
              bias_crpss = crpss$bias,
              evmos_calib = s2dv::Reorder(data = calibrated$evmos, order = c(time_dim,member_dim)),
	      acc_evmos_r = metrics$acc$evmos$r,
              acc_evmos_sign = metrics$acc$evmos$sign,
              evmos_rpss = my.EnsRpss$evmos,
              evmos_crpss = crpss$evmos,
              mse_min_calib = s2dv::Reorder(data = calibrated$mse_min,order = c(time_dim,member_dim)),
              acc_mse_min_r = metrics$acc$mse_min$r,
              acc_mse_min_sign = metrics$acc$mse_min$sign,
              mse_min_rpss = my.EnsRpss$mse_min,
              mse_min_crpss = crpss$mse_min,
              crps_min_calib = s2dv::Reorder(data = calibrated$crps_min,order = c(time_dim,member_dim)),
	      acc_crps_min_r = metrics$acc$crps_min$r,
              acc_crps_min_sign = metrics$acc$crps_min$sign,
              crps_min_rpss = my.EnsRpss$crps_min,
              crps_min_crpss = crpss$crps_min,
              rpc_based_calib = s2dv::Reorder(data = calibrated[['rpc-based']],order = c(time_dim,member_dim)),
	      acc_rpc_based_r = metrics$acc[['rpc-based']]$r,
              acc_rpc_based_sign = metrics$acc[['rpc-based']]$sign,
              rpc_based_rpss = my.EnsRpss[['rpc-based']],
              rpc_based_crpss = crpss[['rpc-based']]))
    
}


# Specify number of categories for the RPSS calculation
n_cat <- 3

# Specify threshold for statistical significance
alpha <- 0.05

step <- startR::Step(fun = skill_calibration,
                       target_dims = list(exp = c('sdate','ensemble','time'), obs = c('sdate','time')),
                       # Especificar output dims:
                       output_dims = list(raw_calib = c('sdate','ensemble'),
                                          acc_raw_r = NULL, 
                                          acc_raw_sign = NULL,
                                          raw_rpss = NULL,
                                          raw_crpss = NULL,
                                          bias_calib = c('sdate','ensemble'),
                                          acc_bias_r = NULL, 
                                          acc_bias_sign = NULL,
                                          bias_rpss = NULL,
                                          bias_crpss = NULL,
                                          evmos_calib = c('sdate','ensemble'),
                                          acc_evmos_r = NULL, 
                                          acc_evmos_sign = NULL,
                                          evmos_rpss = NULL,
                                          evmos_crpss = NULL,
                                          mse_min_calib = c('sdate','ensemble'),
                                          acc_mse_min_r = NULL, 
                                          acc_mse_min_sign= NULL,
                                          mse_min_rpss = NULL,
                                          mse_min_crpss = NULL,
                                          crps_min_calib = c('sdate','ensemble'),
                                          acc_crps_min_r = NULL, 
                                          acc_crps_min_sign = NULL,
                                          crps_min_rpss = NULL,
                                          crps_min_crpss = NULL,
                                          rpc_based_calib = c('sdate','ensemble'),
                                          acc_rpc_based_r = NULL, 
                                          acc_rpc_based_sign = NULL,
                                          rpc_based_rpss = NULL,
                                          rpc_based_crpss = NULL))
 



wf <- startR::AddStep(inputs = list(exp = exp, obs = obs), 
                        methods=methods, 
                        n_cat = n_cat,
                        alpha = alpha,
			ncores = ncores,
                        step_fun = step)[[1]]

output <- startR::Compute(workflow = wf,
                             chunks = list(longitude = 2),
                             threads_load = 1,
                             threads_compute = ncores,
                             cluster = list(queue_host = 'nord4',
                                            queue_type = 'slurm',
                                            temp_dir = '/gpfs/scratch/bsc32/bsc32413/startR_hpc/', #############################3
                                            CDO_module = 'CDO/1.9.8-foss-2019b',
                                            cores_per_job = ncores,
                                            job_wallclock = '02:00:00',
                                            max_jobs = 10,
                                            bidirectional = FALSE,
                                            polling_period = 10),
                             ecflow_suite_dir = "/home/Earth/ptrascas/Desktop/startR_local/", ################################
                             wait = TRUE)

lat <- as.numeric(attr(obs,'Variables')$dat1$latitude)
lon <- as.numeric(attr(obs,'Variables')$dat1$longitude)


# Save list in a file
saveRDS(object = list(output = output, lat = lat, lon = lon),
        file = 'calibration_compute_FOCUS.rds')
