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

source("modules/Crossval/R/tmp/GetProbs.R")
Crossval_Calibration <- function(recipe, data, correct_negative = FALSE) { 

  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') {

      hcst_tr <- MergeDims(hcst_tr, merge_dims = c('sday', 'syear'), 
                          rename_dim = 'syear', na.rm = FALSE) 

      obs_tr <- MergeDims(obs_tr, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)

      hcst_ev <- MergeDims(hcst_ev, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)

    }
    if (cal_method %in% c('PTF', 'DIST', 'RQUANT', 'QUANT', 'SSPLIN')) {
      cal_hcst_tr <- QuantileMapping(exp = hcst_tr, obs = obs_tr,
				     method = cal_method,
				     memb_dim = 'ensemble', 
				     sdate_dim = 'syear',
				     window_dim = NULL,
				     na.rm = TRUE, qstep = 0.1,
				     ncores = ncores)
      cal_hcst_ev <- QuantileMapping(exp = hcst_tr, obs = obs_tr,
				     exp_cor = hcst_ev,
				     method = cal_method,
                                     memb_dim = 'ensemble',
                                     sdate_dim = 'syear',
                                     window_dim = NULL,
                                     na.rm = TRUE, qstep = 0.1,
                                     ncores = ncores)
    } else {
      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)
    }
    # If precipitation is negative:
    if (correct_negative) {
      cal_hcst_tr[cal_hcst_tr < 0] <- 0
      cal_hcst_ev[cal_hcst_ev < 0] <- 0
    }
    lims_cal_hcst_tr <- Apply(cal_hcst_tr, target_dims = c('syear', 'ensemble'),
                              fun = function(x, prob_lims) {
                                lapply(prob_lims, function(ps) {
                                       as.array(quantile(as.vector(x), ps,
                                                         na.rm = na.rm))})},
                              output_dims = lapply(categories,
                                                   function(x) {'bin'}), 
                              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) {
                                      as.array(quantile(as.vector(x), ps,
                                                        na.rm = na.rm))})},
                             output_dims = lapply(categories,
                                                  function(x){'bin'}),
                             prob_lims = categories,
                             ncores = ncores)

   if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
     cal_hcst_tr <- SplitDim(cal_hcst_tr, split_dim = 'syear',
                             new_dim_name = 'sday',
                             indices = rep(1:dim(data$hcst$data)['sday'],
                                           length(cross[[t]]$train.dexes)))
     cal_hcst_ev <- SplitDim(cal_hcst_ev, split_dim = 'syear',
                             new_dim_name = 'sday',
                             indices = rep(1:dim(data$hcst$data)['sday'],
                                           length(cross[[t]]$eval.dexes)))
     obs_tr <- SplitDim(obs_tr, split_dim = 'syear', new_dim_name = 'sday',
                        indices = rep(1:dim(data$hcst$data)['sday'],
				      length(cross[[t]]$train.dexes)))
     for (ps in 1:length(categories)) {
       lims_cal_hcst_tr[[ps]] <- InsertDim(lims_cal_hcst_tr[[ps]],
				      pos = 1, len = 1, name = 'sday')
       lims_obs_tr[[ps]] <- InsertDim(lims_obs_tr[[ps]],
				      pos = 1, len = 1, name = 'sday')
     }
   }
    # check order before storing:
    cal_hcst_ev <- Reorder(cal_hcst_ev, 
			   names(dim(data$hcst$data)))
    cal_hcst_tr <- Reorder(cal_hcst_tr,
			   names(dim(data$hcst$data)))
    obs_tr <- Reorder(obs_tr, names(dim(data$obs$data)))
    for (ps in 1:length(categories)) {
      lims_cal_hcst_tr[[ps]] <- Reorder(lims_cal_hcst_tr[[ps]],
			  c('bin', names(dim(data$obs$data))[-c(5,9)]))
      lims_obs_tr[[ps]] <- Reorder(lims_obs_tr[[ps]],
                                 c('bin', names(dim(data$obs$data))[-5]))
    }
    #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,
       "#### Calibration 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('bin', 'dat', 'var', 'sday',
                                                'sweek', 'time', 
				                'latitude', 'longitude',
						'syear')
    names(dim(lims_obs_tr_res[[ps]])) <- c('bin', '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')
  }

  # Forecast calibration:
  if (!is.null(data$fcst)) {
    if (tolower(recipe$Analysis$Horizon) %in% c('seasonal', 'decadal')) {
      if (cal_method %in% c('PTF', 'DIST', 'RQUANT', 'QUANT', 'SSPLIN')) {
        data$fcst$data <- QuantileMapping(exp = data$hcst$data,
					  obs = data$obs$data,
					  exp_cor = data$fcst$data,
                                          method = cal_method,
                                          memb_dim = 'ensemble',
                                          sdate_dim = 'syear',
                                          window_dim = NULL,
                                          na.rm = TRUE, qstep = 0.1,
                                          ncores = ncores)
        hcst_cal <- QuantileMapping(exp = data$hcst$data,
				    obs = data$obs$data,
				    method = cal_method,
                                    memb_dim = 'ensemble',
                                    sdate_dim = 'syear',
                                    window_dim = NULL,
                                    na.rm = TRUE, qstep = 0.1,
                                    ncores = ncores)
      } else {
        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)
      }
      # For compatibility with subseasonal:
      obs <- data$obs$data
    } else { # if subseasonal
      # merge sample dimensions and select central week
      hcst <- MergeDims(data$hcst$data, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)
      hcst <- Subset(hcst, along = 'sweek',
                     indices = (dim(data$hcst$data)['sweek'] + 1) / 2)
      obs <- MergeDims(data$obs$data, merge_dims = c('sday', 'syear'),
                          rename_dim = 'syear', na.rm = FALSE)
      obs <- Subset(obs, along = 'sweek',
                    indices = (dim(data$obs$data)['sweek'] + 1) / 2)
      fcst <- Subset(data$fcst$data, along = 'sday', indices = 1, 
                     drop = 'selected')
      if (cal_method %in% c('PTF', 'DIST', 'RQUANT', 'QUANT', 'SSPLIN')) {
        data$fcst$data <- QuantileMapping(exp = hcst,
                                          obs = obs,
                                          exp_cor = fcst,
                                          method = cal_method,
                                          memb_dim = 'ensemble',
                                          sdate_dim = 'syear',
                                          window_dim = NULL,
                                          na.rm = TRUE, qstep = 0.1,
                                          ncores = ncores)
        hcst_cal <- QuantileMapping(exp = hcst,
                                    obs = obs,
                                    method = cal_method,
                                          memb_dim = 'ensemble',
                                          sdate_dim = 'syear',
                                          window_dim = NULL,
                                          na.rm = TRUE, qstep = 0.1,
                                          ncores = ncores)
      } else {
        data$fcst$data <- Calibration(exp = hcst,
                              obs = obs,
                              exp_cor = fcst,
                              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 = hcst,
                              obs = obs,
                              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)
      }
      if (!('sday' %in% names(dim(data$fcst$data)))) {
        data$fcst$data <- InsertDim(data$fcst$data,
                                    len = 1, pos = 3, name = 'sday')
      }
    } # end condition subseasonal
        if (correct_negative) {
      hcst_cal[hcst_cal < 0] <- 0
      data$fcst$data[data$fcst$data < 0] <- 0
    }
    # 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) {
                                as.array(quantile(as.vector(x), ps,
                                                  na.rm = na.rm))})},
                       output_dims = lapply(categories, function(x) {'bin'}),
                       prob_lims = categories,
                       ncores = ncores)
    lims <- Apply(obs, target_dims = c('syear', 'ensemble'),
                  fun = function(x, prob_lims) {
                    lapply(prob_lims, function(ps) {
                           as.array(quantile(as.vector(x), ps,
                                             na.rm = na.rm))})},
                  output_dims = lapply(categories, function(x) {'bin'}),
                  prob_lims = categories,
                  ncores = ncores)
    tmp_lims2 <- list()
    
    # SAVING 'lims' which are the observed category limits:
    # TODO saving:
    recipe$Run$output_dir <- paste0(recipe$Run$output_dir,
                                  "/outputs/Calibration/")

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

    if (recipe$Analysis$Workflow$Probabilities$save == 'all') {
      for (ps in 1:length(categories)) {
        ## TODO: use .drop_dims()
        tmp_lims3 <- .drop_dims(lims[[ps]])
        for (l in 1:dim(lims[[ps]])['bin']) {
          tmp_lims <- ClimProjDiags::Subset(x = tmp_lims3,
                                            along = "bin",
                                            indices = l,
                                            drop = "selected")
          tmp_lims2 <- append(tmp_lims2, list(tmp_lims))
          names(tmp_lims2)[length(tmp_lims2)] <- 
		  paste0("p_", as.character(categories[[ps]][l]*100))
        }
      }
      info(recipe$Run$logger,
	   "SAVING OBSERVED CATEGORY LIMITS")
      save_percentiles(recipe = recipe, percentiles = tmp_lims2,
                       data_cube = data$obs,
                       agg = "global", outdir = NULL) 
    }
  }
  # Compute Probabilities
  for (ps in 1:length(categories)) {
    # Get only the probabilities of the central day in sday
    central_day <- (dim(cal_hcst_ev_res)['sday'] + 1)/2
    tmp <- Subset(cal_hcst_ev_res, along = 'sday', indices = central_day)
    hcst_probs_ev[[ps]] <- GetProbs(tmp, time_dim = 'syear',
                            prob_thresholds = NULL,
                            bin_dim_abs = 'bin',
                            indices_for_quantiles = NULL,
                            memb_dim = 'ensemble',
                            abs_thresholds = lims_cal_hcst_tr_res[[ps]],
                            ncores = ncores)
    tmp <- Subset(data$obs$data, along = 'sday', indices = central_day)
    obs_probs_ev[[ps]] <- GetProbs(tmp, time_dim = 'syear',
                                   prob_thresholds = NULL,
                                   bin_dim_abs = 'bin',
                                   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 = 'bin',
                                   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 

  if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
    # Keep only de central sday position from the dimension
    if (dim(hcst$data)['sday'] > 1) {
      hcst <- CST_Subset(hcst, along = 'sday',
                         indices = (dim(hcst$data)['sday'] + 1) / 2)
    }
    if (dim(data$obs$data)['sday'] > 1) {
      data$obs <- CST_Subset(data$obs, along = 'sday',
                         indices = (dim(data$obs$data)['sday'] + 1) / 2)
    }
    if (dim(data$hcst$data)['sday'] > 1) {
      data$hcst <- CST_Subset(data$hcst, along = 'sday',
                         indices = (dim(data$hcst$data)['sday'] + 1) / 2)
    } 
    if (dim(obs_tr_res)['sday'] > 1) {
      obs_tr_res <- Subset(obs_tr_res, along = 'sday',
                           indices = (dim(obs_tr_res)['sday'] + 1) / 2)
    }
  }
  info(recipe$Run$logger, 
       "#### Calibrated  and Probabilities Done #####")
  if (recipe$Analysis$Workflow$Calibration$save != FALSE) {
    info(recipe$Run$logger, "##### START SAVING CALIBRATED #####")
    # 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

  ## TODO: Add .drop_dims() here
  for (ps in 1:length(categories)) {
    for (perc in 1:(length(categories[[ps]]) + 1)) {
      if (perc == 1) {
        name_elem <- paste0("below_", categories[[ps]][perc]*100)
      } else if (perc == length(categories[[ps]]) + 1) {
        name_elem <- paste0("above_", categories[[ps]][perc-1]*100)
      } else {
        name_elem <- paste0("from_", categories[[ps]][perc-1]*100,
                            "_to_", categories[[ps]][perc]*100)
      }
      probs_hcst <- append(list(Subset(hcst_probs_ev[[ps]],
                                       along = 'bin',
                                       indices = perc,
                                       drop = 'selected')),
                           probs_hcst)
      probs_obs <- append(list(Subset(obs_probs_ev[[ps]],
                                      along = 'bin',
                                      indices = perc,
                                      drop = 'selected')),
                          probs_obs)
      if (!is.null(data$fcst)) {
        probs_fcst <- append(list(Subset(fcst_probs[[ps]],
                                         along = 'bin',
                                         indices = perc,
                                         drop = 'selected')),
                             probs_fcst)
      }
      all_names <- c(all_names, name_elem)
    }
  }
  ## TODO: Apply .drop_dims() directly to hcst_probs_ev etc; and move 
  ## reorganizing and renaming to save_probabilities()
  names(probs_hcst) <- all_names
  probs_hcst <- lapply(probs_hcst, .drop_dims)
  names(probs_obs) <- all_names
  probs_obs <- lapply(probs_obs, .drop_dims)
  if (!is.null(data$fcst)) {
    names(probs_fcst) <- all_names
    probs_fcst <- lapply(probs_fcst, .drop_dims)
  }
  if (recipe$Analysis$Workflow$Probabilities$save %in%
      c('all', 'bins_only')) {
    ## Saving only forecast probabilities for now
    # 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$obs, agg = agg,
    #                    type = "obs")
    save_probabilities(recipe = recipe, probs = probs_fcst,
                       data_cube = data$fcst, agg = agg,
                       type = "fcst")
  }

  return(list(hcst = hcst, obs = data$obs, fcst = data$fcst,
              hcst.full_val = data$hcst, obs.full_val = data$obs,
	      cat_lims = list(obs_lims = lims, hcst_lims = lims_fcst),
              #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))

}


