###################################################################################
## Author: E. Duzenli
## Description: Example script for seasonal downscaling with full cross valuidation
###################################################################################
setwd("/esarchive/scratch/ptrascas/R/dev-test_bigpredidata/sunset/sunset") # edit please
#setwd("/esarchive/scratch/eduzenli/gitlab/amd_gitlab/sunset/") 
# Load modules
source("modules/Loading/Loading.R")
source("modules/Downscaling/Crossval_downscaling.R")
source("modules/Downscaling/Downscaling.R")
#source("modules/Saving/Saving.R")
source("modules/Visualization/Visualization.R")
# Input
var <- "tas"
REF <- "ERA5"
#target_grid <- "/esarchive/recon/ecmwf/era5land/monthly_mean/tas_f1h/tas_200202.nc" # a path to ref dataset
target_grid <- "/esarchive/recon/ecmwf/era5/monthly_mean/tas_f1h/tas_202401.nc"
filesystem <- "esarchive"
System <- "ECMWF-SEAS5.1"
Horizon <- "seasonal"
freq <- "monthly_mean" 
target_month <- 1 # 1:12
lt <- 0 # lead time
hcst_start <- 2000
hcst_end <- 2005
fcst_year <- 2020
anomaly_calc <- TRUE # for now, keep always TRUE, for analogs. 
                     # To obtain downscaled absolute values as the output set ABS_VALUE TRUE.
ABS_VALUE <- TRUE
dwn_type <- 'intbc'   ### 'intbc', 'intlr', 'analogs', 'logreg'
int_method <- "bil"  # "NULL",'con', 'bil', 'bic', 'nn', 'con2',"dis","laf"
bc_method <- "bias" # # NULL, "quantile_mapping", "bias", "evmos", "mse_min", "crps_min", "rpc-based", "qm" 
lr_method <- NULL           # 'NULL' 'basic', 'large-scale', '9nn'
log_reg_method <- NULL # "NULL" ,"ens_mean", "ens_mean_sd", "sorted_members"
nanalogs <- 1
latmin = 36 
latmax = 44 
lonmin = -10
lonmax = 5 
categories <- list(c(1/3, 2/3), c(1/10, 9/10))

# Read recipe
if (filesystem == "gpfs") {
  recipe_file <- "recipes/atomic_recipes/recipe_seasonal_downscaling_gpfs.yml"
} else {
  recipe_file <- "recipes/recipe_seasonal_downscaling.yml"
}
recipe <- prepare_outputs(recipe_file)

# Edit recipe
recipe$Analysis$Variables$name <- var
recipe$Analysis$Datasets$Reference$name <- REF
recipe$Analysis$Datasets$System$name <- System
recipe$Analysis$Horizon <- Horizon
recipe$Analysis$Variables$freq <- freq
recipe$Analysis$Time$ftime_min <- lt+1
recipe$Analysis$Time$ftime_max <- lt+1
recipe$Analysis$Time$hcst_start <- hcst_start
recipe$Analysis$Time$hcst_end <- hcst_end
recipe$Analysis$Time$fcst_year <- fcst_year
diffmon <- target_month - lt
smonth <- paste0(sprintf('%02d', ifelse(diffmon > 0, diffmon, 12 + diffmon)), '01')
recipe$Analysis$Time$sdate <-smonth
recipe$Analysis$Workflow$Anomalies$compute <- anomaly_calc
recipe$Analysis$Workflow$Downscaling$type           <- dwn_type
recipe$Analysis$Workflow$Downscaling$int_method     <- int_method
recipe$Analysis$Workflow$Downscaling$bc_method      <- bc_method
recipe$Analysis$Workflow$Downscaling$lr_method      <- lr_method
recipe$Analysis$Workflow$Downscaling$log_reg_method <- log_reg_method
recipe$Analysis$Workflow$Downscaling$nanalogs       <- nanalogs
recipe$Analysis$Workflow$Downscaling$target_grid <- target_grid

recipe$Analysis$Workflow$Probabilities$percentiles <- NULL
for (i in 1:length(categories)) {
  recipe$Analysis$Workflow$Probabilities$percentiles[[i]] <- categories[[i]]
}
recipe$Run$filesystem <- filesystem 
# Load datasets
data <- Loading(recipe)

# Downscale datasets
dwn_data <- Crossval_downscaling(recipe, data)

if (ABS_VALUE) {
  # TODO it is better to use window or loess-wise 
  # functions for the daily climatology calculation to include the effect of neighbouring days.
  for (i in 1:dwn_data$obs$dims[["syear"]]) {
    tmp <- Subset (dwn_data$obs$data, along = "syear", indices = -i, drop = FALSE) 
    dum <- Clim(tmp, tmp, time_dim = "syear", memb_dim = "ensemble", 
                dat_dim = "dat", na.rm = TRUE)$clim_obs
    assign (paste0("clim_", i), dum)
    rm(dum, tmp)
  }

  Clim <- abind(mget(ls(pattern = ("clim_"))), along = length(dim(clim_1)) + 1 )
  names(dim(Clim)) <- c(names(dim(clim_1)), "syear")
  Clim <- Subset (Clim, along = "ensemble", indices = 1, drop = "selected")
  rm(list=ls(pattern = ("clim_")))

  abs_vals <- Apply (list(dwn_data$hcst$data, Clim), target_dims = c("latitude","longitude"),
                     fun = function (x,y) return(x+y))$output1
}

source("modules/Crossval/Crossval_metrics.R")
skill_metrics <- Crossval_metrics(recipe = recipe, data_crossval = dwn_data,
                                fair = FALSE, nmemb = NULL, nmemb_ref = NULL)




if (!is.null(dwn_data$fcst)) {
  # Required to plot a forecast:
  data$fcst <- dwn_data$fcst
  tmp_probs <- list(probs_fcst = dwn_data$probs$fcst[[1]])
  #nlats <- as.numeric(dim(data$hcst$data)['latitude'])
  #nlons <- as.numeric(dim(data$hcst$data)['longitude'])
  #ntimes <- as.numeric(dim(data$hcst$data)['time'])

  #   TODO: consider more than terciles
  #dim(tmp_probs$probs_fcst) <- c(bin = 3, syear = 1, var = 1, time = ntimes,
  #                            latitude = nlats, longitude = nlons)


  Visualization(recipe = recipe, data = dwn_data, skill_metrics = skill_metrics,
                significance = TRUE, probabilities = tmp_probs)
} else {
    Visualization(recipe = recipe, data = data, skill_metrics = skill_metrics,
                  significance = TRUE)
}


#all (as.numeric(dwn_data$hcst$data) == as.numeric(dwn_data2$hcst$data))



