# take the output of Flor/s2s/subseasonal_loading.R

#source("../git/sunset/modules/Crossval/R/tmp/GetProbs.R")
source("modules/Crossval/R/tmp/GetProbs.R")
Crossval_Calibration <- function(recipe, data) { 

  cal_method <- recipe$Analysis$Workflow$Calibration$method
  cross.method <- recipe$Analysis$cross.method
  # TODO move check
  if (is.null(cross.method)) {
    cross.method <- 'leave-one-out'
  }
  categories <- recipe$Analysis$Workflow$Probabilities$percentiles
  categories <- lapply(categories, function (x) { 
                       sapply(x, function(y) {
                              eval(parse(text = y))})})
  ncores <- recipe$Analysis$ncores
  na.rm <- recipe$Analysis$remove_NAs 
  ## data dimensions 
  sdate_dim <- dim(data$hcst$data)['syear']
  orig_dims <- names(dim(data$hcst$data))
  # spatial dims
  if ('latitude' %in% names(dim(data$hcst$data))) {
    nlats <- dim(data$hcst$data)['latitude']
    nlons <- dim(data$hcst$data)['longitude']
    agg = 'global'
  } else if  ('region' %in% names(dim(data$hcst$data))) {
    agg = 'region'
    nregions <- dim(data$hcst$data)['region']
  }

 
  cross <- CSTools:::.make.eval.train.dexes(eval.method = cross.method,
                                            amt.points = sdate_dim,
                                            amt.points_cor = NULL) # k = ?

  cal_hcst_ev_res <- NULL
  cal_hcst_tr_res <- NULL
  #cal_obs_ev_res <- NULL
  obs_tr_res <- NULL
  # as long as probs requested in recipe:
  lims_cal_hcst_tr_res <- lapply(categories, function(X) {NULL})
  lims_obs_tr_res <- lapply(categories, function(X) {NULL})

  fcst_probs <- lapply(categories, function(x){NULL})
  hcst_probs_ev <- lapply(categories, function(x){NULL})
  obs_probs_ev <- lapply(categories, function(x){NULL})
  for (t in 1:length(cross)) {
    info(recipe$Run$logger, paste("crossval:", t))

    # subset years: Subset works at BSC not at Athos
    ## training indices
    obs_tr <- Subset(data$obs$data, along = 'syear',
   		     indices = cross[[t]]$train.dexes)
    hcst_tr <- Subset(data$hcst$data, along = 'syear', 
  		      indices = cross[[t]]$train.dexes)
    ## evaluation indices
    hcst_ev <- Subset(data$hcst$data, along = 'syear',
  		      indices = cross[[t]]$eval.dexes)
    obs_ev <- Subset(data$obs$data, along = 'syear',
  		     indices = cross[[t]]$eval.dexes)
    
    if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
source("https://earth.bsc.es/gitlab/external/cstools/-/raw/d6914a40c11d09168b9b4a191e9a0362b56a5f0c/R/CST_MergeDims.R")

      hcst_tr <- CST_MergeDims(hcst_tr, merge_dims = c('sday', 'syear'), 
                          rename_dim = 'syear', na.rm = FALSE) 
      obs_tr <- CST_MergeDims(obs_tr, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)
      hcst_ev <- CST_MergeDims(hcst_ev, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)
    }
    cal_hcst_tr <- Calibration(exp = hcst_tr, obs = obs_tr,
                               cal.method = cal_method,
                               memb_dim = 'ensemble', sdate_dim = 'syear',
                               eval.method = 'in-sample',
                               na.fill = FALSE, na.rm = TRUE,
                               apply_to = NULL,
                               alpha = NULL, ncores = ncores)

    cal_hcst_ev <- Calibration(exp = hcst_tr, obs = obs_tr, exp_cor = hcst_ev,
			       cal.method = cal_method,
                               memb_dim = 'ensemble', sdate_dim = 'syear',
			       eval.method = 'in-sample',
                               na.fill = FALSE, na.rm = TRUE,
                               apply_to = NULL,
                               alpha = NULL, ncores = ncores)


   lims_cal_hcst_tr <- Apply(cal_hcst_tr, target_dims = c("syear", "ensemble"), 
			     fun = function(x,prob_lims) {
        		test <- lapply(prob_lims, function(ps) {
           		res <- quantile(as.vector(x), ps, na.rm = T)
           		dim(res) <- c(cat = length(res))
           		return(res)
       			})
   			}, prob_lims = categories, ncores = ncores)
	   	   
    lims_obs_tr <- Apply(obs_tr, target_dims = c('syear'),#, 'ensemble'),
                             fun = function(x, prob_lims) {
                               lapply(prob_lims, function(ps) {
                              res <- quantile(as.vector(x), ps, na.rm = T)
			     dim(res) <- c(cat = length(res))
			     return(res)
			     })
			     },
                             output_dims = lapply(categories, function(x){'cat'}),
                             prob_lims = categories,
                             ncores = ncores)

    #store results
    cal_hcst_ev_res <- abind(cal_hcst_ev_res, cal_hcst_ev,
                             along = length(dim(cal_hcst_ev)) + 1)
    cal_hcst_tr_res <- abind(cal_hcst_tr_res, cal_hcst_tr,
                             along = length(dim(cal_hcst_tr)) + 1)

    # Obs ev are original obs
#    cal_obs_ev_res <- abind(ano_obs_ev_res, ano_obs_ev,
#                            along = length(dim(ano_obs_ev)) + 1)
    obs_tr_res <- abind(obs_tr_res, obs_tr,
                            along = length(dim(obs_tr)) + 1)
    for(ps in 1:length(categories)) {
      lims_cal_hcst_tr_res[[ps]] <- abind(lims_cal_hcst_tr_res[[ps]], 
					  lims_cal_hcst_tr[[ps]],
                                along = length(dim(lims_cal_hcst_tr[[ps]])) + 1)
      lims_obs_tr_res[[ps]] <- abind(lims_obs_tr_res[[ps]], 
					 lims_obs_tr[[ps]],
                                 along = length(dim(lims_obs_tr[[ps]])) + 1)
    }
  }
  info(recipe$Run$logger,
       "#### Anomalies Cross-validation loop ended #####")
  gc()
  names(dim(cal_hcst_ev_res)) <- c('dat', 'var', 'sday', 'sweek', 'unneeded', 'time',
				   'latitude', 'longitude', 'ensemble', 'syear')
  names(dim(cal_hcst_tr_res)) <- c('dat', 'var', 'sday', 'sweek', 'loop', 'time',
                                   'latitude', 'longitude', 'ensemble', 'syear')
  names(dim(obs_tr_res)) <- c('dat', 'var', 'sday', 'sweek', 'ensemble', 'time',
                                   'latitude', 'longitude', 'unneeded', 'syear')
  obs_tr_res <- Subset(obs_tr_res, along = 'unneeded',
                           indices = 1, drop = 'selected')
  cal_hcst_ev_res <- Subset(cal_hcst_ev_res, along = 'unneeded',
                            indices = 1, drop = 'selected')
  for(ps in 1:length(categories)) {
    names(dim(lims_cal_hcst_tr_res[[ps]])) <- c('cat', 'dat', 'var', 'sday', 'sweek',
						'time', 'latitude', 'longitude',
						'syear')
    names(dim(lims_obs_tr_res[[ps]])) <- c('cat', 'dat', 'var', 'sday', 'sweek',
					   'time', 'latitude', 'longitude',
					   'unneeded', 'syear')
    lims_obs_tr_res[[ps]] <- Subset(lims_obs_tr_res[[ps]],
                                    along = 'unneeded', indices = 1,
				    drop = 'selected')
  }

  # Make categories rounded number to use as names:
  categories <- recipe$Analysis$Workflow$Probabilities$percentiles
  categories <- lapply(categories, function (x) {
                       sapply(x, function(y) {
                              round(eval(parse(text = y)),2)})})
  recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
                                  "/outputs/Calibration/")

  # Forecast calibration:
  if (!is.null(data$fcst)) {
    data$fcst$data <- Calibration(exp = data$hcst$data,
				  obs = data$obs$data,
				  exp_cor = data$fcst$data, 
                                  cal.method = cal_method,
				  multi.model = FALSE, 
                                  na.fill = TRUE, na.rm = TRUE,
				  apply_to = NULL, 
                                  alpha = NULL, memb_dim = 'ensemble',
				  sdate_dim = 'syear', 
                                  dat_dim = NULL, ncores = ncores)
    hcst_cal <- Calibration(exp = data$hcst$data,
                                  obs = data$obs$data,
                                  cal.method = cal_method,
                                  eval.method = 'in-sample',
                                  multi.model = FALSE,
                                  na.fill = TRUE, na.rm = TRUE,
                                  apply_to = NULL,
                                  alpha = NULL, memb_dim = 'ensemble',
                                  sdate_dim = 'syear',
                                  dat_dim = NULL, ncores = ncores)

    # Terciles limits using the whole hindcast period:
    lims_fcst <- Apply(hcst_cal, target_dims = c('syear', 'ensemble'),
                       fun = function(x, prob_lims) {
                         lapply(prob_lims, function(ps) {
                         quantile(as.vector(x), ps, na.rm = na.rm)})},
                       output_dims = lapply(categories, function(x) {'cat'}),
                       prob_lims = categories,
                       ncores = ncores)
    lims <- Apply(data$obs$data, target_dims = c('syear', 'ensemble'),
                       fun = function(x, prob_lims) {
                         lapply(prob_lims, function(ps) {
                         quantile(as.vector(x), ps, na.rm = na.rm)})},
                       output_dims = lapply(categories, function(x) {'cat'}),
                       prob_lims = categories,
                       ncores = ncores)
    tmp_lims2 <- list()
# What to save hcst category limits or obs category limits?
    for(ps in 1:length(categories)) {
      tmp_lims3 <- drop(lims[[ps]])
      for (l in 1:dim(lims[[ps]])['cat']) {
        tmp_lims <- tmp_lims3[l,,,]
        if (!('var' %in% names(dim(tmp_lims)))) {
          dim(tmp_lims) <- c(var = 1, dim(tmp_lims))
        }
        tmp_lims2 <- append(tmp_lims2, list(tmp_lims))
        names(tmp_lims2)[length(tmp_lims2)] <- as.character(categories[[ps]][l])
      }
      if (recipe$Analysis$Workflow$Probabilities$save == 'yes') {
        save_percentiles(recipe = recipe, percentiles = tmp_lims2,
                         data_cube = data$obs,
                         agg = "global",
			 outdir = NULL) 
      }
    }
  }

  # Compute Probabilities
  for (ps in 1:length(categories)) {
    hcst_probs_ev[[ps]] <- GetProbs(cal_hcst_ev_res, time_dim = 'syear',
                              prob_thresholds = NULL,
                              bin_dim_abs = 'cat',
                              indices_for_quantiles = NULL,
                              memb_dim = 'ensemble',
                              abs_thresholds = lims_cal_hcst_tr_res[[ps]],
                              ncores = ncores)
    obs_probs_ev[[ps]] <- GetProbs(data$obs$data, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'cat',
                                   indices_for_quantiles = NULL,
                                   memb_dim = 'ensemble',
                                   abs_thresholds = lims_obs_tr_res[[ps]],
                                   ncores = ncores)
    if (!is.null(data$fcst)) {
      fcst_probs[[ps]] <- GetProbs(data$fcst$data, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'cat',
                                   indices_for_quantiles = NULL,
                                   memb_dim = 'ensemble',
                                   abs_thresholds = lims_fcst[[ps]],
                                   ncores = ncores)
    }
  }
  # Convert to s2dv_cubes the resulting calibrated
  hcst <- data$hcst
  hcst$data <- cal_hcst_ev_res 

  info(recipe$Run$logger, 
       "#### Calibrated  and Probabilities Done #####") 
  if (recipe$Analysis$Workflow$Anomalies$save != 'none') {
    info(recipe$Run$logger, "##### START SAVING CALIBRATED #####")
#    recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
#                                    "/outputs/Anomalies/")
    # Save forecast
    if ((recipe$Analysis$Workflow$Calibration$save %in% 
           c('all', 'exp_only', 'fcst_only')) && !is.null(data$fcst)) {
      save_forecast(recipe = recipe, data_cube = data$fcst, type = 'fcst')
    }
    # Save hindcast
    if (recipe$Analysis$Workflow$Anomalies$save %in%
        c('all', 'exp_only')) {
      save_forecast(recipe = recipe, data_cube = hcst, type = 'hcst')
    }
  }
   # Save probability bins
  probs_hcst <- list()
  probs_fcst <- list()
  probs_obs <- list()
  all_names <- NULL

  for (ps in 1:length(categories)) {
    for (perc in 1:(length(categories[[ps]]) + 1)) {
      if (perc == 1) {
        name_elem <- paste0("below_", categories[[ps]][perc])
      } else if (perc == length(categories[[ps]]) + 1) {
        name_elem <- paste0("above_", categories[[ps]][perc-1])
      } else {
        name_elem <- paste0("from_", categories[[ps]][perc-1], 
                            "_to_", categories[[ps]][perc])
      }
      probs_hcst <- append(list(Subset(hcst_probs_ev[[ps]],
                             along = 'cat', indices = perc, drop = 'all')),
                           probs_hcst)
      probs_obs <- append(list(Subset(obs_probs_ev[[ps]],
                              along = 'cat', indices = perc, drop = 'all')),
                          probs_obs)
      if (!is.null(data$fcst)) {
        probs_fcst <- append(list(Subset(fcst_probs[[ps]],
                              along = 'cat', indices = perc, drop = 'all')),
                             probs_fcst)
      }
      all_names <- c(all_names, name_elem)
    }
  }
  names(probs_hcst) <- all_names
  if (!('var' %in% names(dim(probs_hcst[[1]])))) {
    probs_hcst <- lapply(probs_hcst, function(x) {
                     dim(x) <- c(var = 1, dim(x))
                     return(x)})
  }
  names(probs_obs) <- all_names
  if (!('var' %in% names(dim(probs_obs[[1]])))) {
    probs_obs <- lapply(probs_obs, function(x) {
                     dim(x) <- c(var = 1, dim(x))
                     return(x)})
  }

  if (!is.null(data$fcst)) {
    names(probs_fcst) <- all_names
    if (!('var' %in% names(dim(probs_fcst[[1]])))) {
      probs_fcst <- lapply(probs_fcst, function(x) {
                     dim(x) <- c(var = 1, dim(x))
                     return(x)})
    }
    if (!('syear' %in% names(dim(probs_fcst[[1]])))) {
      probs_fcst <- lapply(probs_fcst, function(x) {
                     dim(x) <- c(syear = 1, dim(x))
                     return(x)})
    }
  }

  if (recipe$Analysis$Workflow$Probabilities$save %in%
    c('all', 'bins_only')) {
    save_probabilities(recipe = recipe, probs = probs_hcst,
                       data_cube = data$hcst, agg = agg,
                       type = "hcst")
    save_probabilities(recipe = recipe, probs = probs_obs,
                       data_cube = data$hcst, agg = agg,
                       type = "obs")
   save_probabilities(recipe = recipe, probs = probs_fcst,
                       data_cube = data$fcst, agg = agg,
                       type = "fcst")
  }


  hcst_EM <- MeanDims(hcst$data, 'ensemble', drop = T)

  if (!is.null(data$fcst)) {
    fcst_EM <- MeanDims(data$fcst$data, 'ensemble', drop = T)
  }
  return(list(hcst = hcst, obs = data$obs, fcst = data$fcst,
              hcst.full_val = data$hcst, obs.full_val = data$obs,
              hcst_EM = hcst_EM, fcst_EM = fcst_EM,
              #cat_lims = list(hcst_tr = lims_ano_hcst_tr_res,
              #                obs_tr = lims_ano_obs_tr_res),
              probs = list(hcst_ev = hcst_probs_ev,
                           obs_ev = obs_probs_ev,
			   fcst = fcst_probs),
              ref_obs_tr = obs_tr_res))

}

