###################################################################################
## 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/Units/Units.R")
source("modules/Visualization/Visualization.R")
# Input

#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"
target_month <- 1 # 1:12
#lt <- 0 # lead time
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
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
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)

#recipe$Analysis$Time$ftime_min <- lt+1
#recipe$Analysis$Time$ftime_max <- lt+1
#diffmon <- target_month - lt
#smonth <- paste0(sprintf('%02d', ifelse(diffmon > 0, diffmon, 12 + diffmon)), '01')
recipe$Analysis$Workflow$Anomalies$compute <- anomaly_calc
# Bilinear interpolation for tas, rsds, sfcWind, conservative for prlr

#if (recipe$Analysis$Variables$name == "prlr") {
#    int_method <- 'conservative'
#} else {
#    int_method <- 'bilinear'
#}

#if (recipe$Analysis$Variables$name == "prlr") {
#    int_method <- 'con'
#} else {
#    int_method <- 'bil'
#}

#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)

# Check units

data <- Units(recipe, data)

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


# Shortcut for analogues
#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
#}

if (recipe$Analysis$Region$name == "Iberia") {
    lat_indices <- which(dwn_data$fcst$coords$latitude >= 36 & dwn_data$fcst$coords$latitude <= 44)
    lon_indices <- which(dwn_data$fcst$coords$longitude >= -10 & dwn_data$fcst$coords$longitude <= 5)
} else if (recipe$Analysis$Region$name == "EU") {
    lat_indices <- which(dwn_data$fcst$coords$latitude >= 20 & dwn_data$fcst$coords$latitude <= 80)
    lon_indices <- which(dwn_data$fcst$coords$longitude >= -20 & dwn_data$fcst$coords$longitude <= 40)
}


 
dwn_data$hcst <- CST_Subset(dwn_data$hcst, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices))
dwn_data$fcst <- CST_Subset(dwn_data$fcst, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices)) 
dwn_data$obs <- CST_Subset(dwn_data$obs, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices))
dwn_data$hcst.full_val <- CST_Subset(dwn_data$hcst.full_val, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices))
dwn_data$obs.full_val <- CST_Subset(dwn_data$obs.full_val, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices))

 # Apply Subset() to each numeric array inside the lists
dwn_data$probs$hcst_ev <- lapply(dwn_data$probs$hcst_ev, function(x) {
   Subset(x, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices)) })
 
dwn_data$probs$obs_ev <- lapply(dwn_data$probs$obs_ev, function(x) {
   Subset(x, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices)) })
 
dwn_data$probs$fcst <- lapply(dwn_data$probs$fcst, function(x) {
   Subset(x, along = c('latitude', 'longitude'), indices = list(lat_indices, lon_indices)) })




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 = data, skill_metrics = skill_metrics,
                significance = TRUE, probabilities = tmp_probs)
} else {
    Visualization(recipe = recipe, data = dwn_data, skill_metrics = skill_metrics,
                  significance = TRUE)
}


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



