###############################################################################
############# Spatial aggregation calculations for scorecards #################
###############################################################################

## This function spatially aggregates the skill metrics and statistics for the 
## scorecards. The aggregated statistical significance is also calculate for the
## metrics when 'score' aggregate is requested. The aggregation regions are 
## defined in the recipe under the scorecards section. This function reads from 
## the atomic recipes.

## Define function
Scorecards_calculations <- function(recipe, data, skill_metrics,
                                    statistics = NULL) {
  
  # Check that skill_metrics is not NULL
  if (is.null(skill_metrics)) {
    stop("Scorecards calculations requested but 'skill_metrics' not provided.")
  } else if (!is.list(skill_metrics)) {
    stop("'skill_metrics' should be a named list with metric arrays")
  }
  # Assign parameters
  ncores <- recipe$Analysis$ncores
  if (is.null(recipe$Analysis$Workflow$Scorecards$signif_alpha)) {
    alpha <- 0.05
  } else {
    alpha <- recipe$Analysis$Workflow$Scorecards$signif_alpha
  }
  
  if (is.null(recipe$Analysis$remove_NAs)) {
    na.rm <- FALSE
  } else {
    na.rm <- recipe$Analysis$remove_NAs 
  }
  
  if (is.null(recipe$Analysis$Workflow$Scorecards$inf_to_na)){
    inf.to.na <- FALSE
  } else {
    inf.to.na <- recipe$Analysis$Workflow$Scorecards$inf_to_na
  }
  
  ## Parameters for saving output data files
  output.path <- paste0(recipe$Run$output_dir, "/plots/Scorecards/")
  dir.create(output.path, recursive = T, showWarnings = F)
  
  system <- recipe$Analysis$Datasets$System$name
  reference <- recipe$Analysis$Datasets$Reference$name
  var <- recipe$Analysis$Variables$name
  start.year <- as.numeric(recipe$Analysis$Time$hcst_start)
  end.year <- as.numeric(recipe$Analysis$Time$hcst_end)
  forecast.months <- recipe$Analysis$Time$ftime_min:recipe$Analysis$Time$ftime_max
  start.months <- substr(recipe$Analysis$Time$sdate, 1,2)
  period <- paste0(start.year, "-", end.year)
  
  ## Parameters for data aggregation
  regions <- recipe$Analysis$Workflow$Scorecards$regions  
  for (i in names(regions)){regions[[i]] <- unlist(regions[[i]])}
  
  metric.aggregation <- recipe$Analysis$Workflow$Scorecards$metric_aggregation
  metrics <- unlist(strsplit(tolower(recipe$Analysis$Workflow$Scorecards$metric),
                             ", | |,"))
  
  ## Define array to filled with data
  aggr_significance <- array(data = NA, 
                             dim = c(time = length(forecast.months), 
                                     region = length(regions),
                                     metric = length(metrics)))
  
  ## For data that is already aggregated by region
  if ("region" %in% names(dim(skill_metrics[[1]]))) { 
    aggr_metrics <- NULL  
    for (met in metrics) {
      skill_metrics[[met]] <- Reorder(skill_metrics[[met]],
                                      c("time", "region", "var"))
      aggr_metrics <- abind(aggr_metrics, skill_metrics[[met]], along=3)
    }
    names(dim(aggr_metrics)) <- c("time", "region", "metric")
  } else { 
    ## Define arrays to filled with data
    aggr_metrics <- array(data = NA, 
                          dim = c(time = length(forecast.months), 
                                  region = length(regions),
                                  metric = length(metrics)))  
    
    lon_dim <- 'longitude'
    lat_dim <- 'latitude'
    time_dim <- 'syear'
    memb_dim <- 'ensemble'
    lon <- as.numeric(data$hcst$coords$longitude)
    lat <- as.numeric(data$hcst$coords$latitude)
    
    ## Skill aggregation
    if (metric.aggregation == 'skill') {
      ## Calculate weighted mean of spatial aggregation
      for (met in metrics) {
        result <- sapply(X = 1:length(regions), 
                         FUN = function(X) {
                           WeightedMean(data = skill_metrics[[met]],
                                        region = regions[[X]],
                                        lon = lon, londim = lon_dim,
                                        lat = lat, latdim = lat_dim,
                                        na.rm = na.rm)
                         }, simplify = 'array')
        
        names(dim(result))[length(dim(result))] <- 'region'
        result <-Subset(result, 'var', 1, drop = 'selected')
        
        if (met =='crpss' && inf.to.na == TRUE) {
          result[result == -Inf] <- NA
        }
        aggr_metrics[,,which(metrics == met)] <- Reorder(data = result, 
                                                         order = c('time', 'region'))
      } ## close on met
    } ## close if on skill
    
    ## Score Aggregation
    if (metric.aggregation == 'score') {
      ## Spatially aggregate data for each metric
      for (met in metrics) {   
        if (met == 'rpss') {
          rps_syear <- sapply(X = 1:length(regions), 
                              FUN = function(X) {
                                WeightedMean(data = skill_metrics$rps_syear,
                                             region = regions[[X]],
                                             lon = lon, londim = lon_dim,
                                             lat = lat, latdim = lat_dim,
                                             na.rm = na.rm)
                              }, simplify = 'array')
          
          rps_clim_syear <- sapply(X = 1:length(regions), 
                                   FUN = function(X) {
                                     WeightedMean(data = skill_metrics$rps_clim_syear,
                                                  region = regions[[X]],
                                                  lon = lon, londim = lon_dim,
                                                  lat = lat, latdim = lat_dim,
                                                  na.rm = na.rm)
                                   }, simplify = 'array')
          
          ## Include name of region dimension
          names(dim(rps_syear))[length(dim(rps_syear))] <- 'region'
          names(dim(rps_clim_syear))[length(dim(rps_clim_syear))] <- 'region'
          
          ## Remove 'var' dimension
          rps_syear <- Subset(rps_syear, 'var', 1, drop = 'selected')
          rps_clim_syear <- Subset(rps_clim_syear, 'var', 1, drop = 'selected')
          
          ## Calculate significance
          sign_rpss <- RandomWalkTest(rps_syear, rps_clim_syear, 
                                      time_dim = time_dim, 
                                      test.type = 'two.sided',
                                      alpha = alpha, pval = FALSE, sign = TRUE, 
                                      ncores = NULL)$sign
          
          ## Average over 'syear' dimension
          rps_syear <- Apply(data = rps_syear, 
                             target_dims = time_dim, 
                             fun = 'mean', ncores = ncores)$output1
          
          rps_clim_syear <- Apply(data = rps_clim_syear, 
                                  target_dims = time_dim,
                                  fun = 'mean',  ncores = ncores)$output1
          
          ## Calculate RPSS from aggregated RPS and RPS_clim
          rpss <- 1 - rps_syear / rps_clim_syear
          
          ## Save metric result in arrays
          aggr_metrics[,,which(metrics == met)] <- Reorder(data = rpss, 
                                                           order = c('time', 'region'))
          aggr_significance[,,which(metrics == met)] <- Reorder(data = sign_rpss,
                                                                order = c('time', 'region'))
          
        } else if (met == 'crpss') {
          crps_syear <- sapply(X = 1:length(regions), 
                               FUN = function(X) {
                                 WeightedMean(data = skill_metrics$crps_syear,
                                              region = regions[[X]],
                                              lon = lon, londim = lon_dim,
                                              lat = lat, latdim = lat_dim,
                                              na.rm = na.rm)
                               }, simplify = 'array')
          crps_clim_syear <- sapply(X = 1:length(regions), 
                                    FUN = function(X) {
                                      WeightedMean(data = skill_metrics$crps_clim_syear,
                                                   region = regions[[X]],
                                                   lon = lon, londim = lon_dim,
                                                   lat = lat, latdim = lat_dim,
                                                   na.rm = na.rm)
                                    }, simplify = 'array')
          
          ## Include name of region dimension
          names(dim(crps_syear))[length(dim(crps_syear))] <- 'region'
          names(dim(crps_clim_syear))[length(dim(crps_clim_syear))] <- 'region'
          
          ## Remove 'var' dimension
          crps_syear <- Subset(crps_syear, 'var', 1, drop = 'selected')
          crps_clim_syear <- Subset(crps_clim_syear, 'var', 1, drop = 'selected')
          
          ## Calculate significance
          sign_crpss <- RandomWalkTest(crps_syear, crps_clim_syear, 
                                       time_dim = time_dim, 
                                       test.type = 'two.sided',
                                       alpha = alpha, pval = FALSE, sign = TRUE, 
                                       ncores = NULL)$sign
          
          ## Average over 'syear' dimension
          crps_syear <- Apply(data = crps_syear, 
                              target_dims = time_dim, 
                              fun = 'mean', ncores = ncores)$output1
          
          crps_clim_syear <- Apply(data = crps_clim_syear, 
                                   target_dims = time_dim,
                                   fun = 'mean', ncores = ncores)$output1
          
          ## Calculate CRPSS from aggregated CRPS and CRPS_clim
          crpss <- 1 - crps_syear / crps_clim_syear
          
          ## Save metric result in arrays
          aggr_metrics[ , , which(metrics == met)] <- Reorder(data = crpss, 
                                                              order = c('time', 'region'))
          aggr_significance[ , , which(metrics == met)] <- Reorder(data = sign_crpss, 
                                                                   order = c('time', 'region'))
        } else if (met == 'enscorr') {
          cov <- sapply(X = 1:length(regions), 
                        FUN = function(X) {
                          WeightedMean(data = statistics$cov,
                                       region = regions[[X]],
                                       lon = lon, londim = lon_dim,
                                       lat = lat, latdim = lat_dim,
                                       na.rm = na.rm)
                        }, simplify = 'array')
          n_eff <- sapply(X = 1:length(regions), 
                          FUN = function(X) {
                            WeightedMean(data = statistics$n_eff,
                                         region = regions[[X]],
                                         lon = lon, londim = lon_dim,
                                         lat = lat, latdim = lat_dim,
                                         na.rm = na.rm)
                          }, simplify = 'array')
          
          ## Calculate variance from standard deviation
          var_hcst <- (statistics$std_hcst)^2
          var_obs <- (statistics$std_obs)^2
          
          var_hcst <- sapply(X = 1:length(regions), 
                             FUN = function(X) {
                               WeightedMean(data = var_hcst,
                                            region = regions[[X]],
                                            lon = lon, londim = lon_dim,
                                            lat = lat, latdim = lat_dim,
                                            na.rm = na.rm)
                             }, simplify = 'array')
          var_obs <- sapply(X = 1:length(regions), 
                            FUN = function(X) {
                              WeightedMean(data = var_obs,
                                           region = regions[[X]],
                                           lon = lon, londim = lon_dim,
                                           lat = lat, latdim = lat_dim,
                                           na.rm = na.rm)
                            }, simplify = 'array')
          
          ## Include name of region dimension
          names(dim(cov))[length(dim(cov))] <- 'region'
          names(dim(var_hcst))[length(dim(var_hcst))] <- 'region'
          names(dim(var_obs))[length(dim(var_obs))] <- 'region'
          names(dim(n_eff))[length(dim(n_eff))] <- 'region'
          
          ## Convert aggregated variance back to standard deviation
          std_hcst <- sqrt(var_hcst)
          std_obs <- sqrt(var_obs)
          
          ## Calculate correlation
          enscorr <- cov / (std_hcst * std_obs)
          
          ## Remove 'var' dimension
          enscorr <- Subset(enscorr, 'var', 1, drop = 'selected')
          n_eff <- Subset(n_eff, 'var', 1, drop = 'selected')
          
          ## Calculate significance of corr
          t_alpha2_n2 <- qt(p = alpha/2, df = n_eff-2, lower.tail = FALSE)
          t <- abs(enscorr) * sqrt(n_eff-2) / sqrt(1-enscorr^2)
          
          sign_corr<- array(data = NA, 
                            dim = c(time = length(forecast.months), 
                                    region = length(regions)))
          
          for (time in 1:dim(sign_corr)[['time']]) {
            for (reg in 1:dim(sign_corr)[['region']]) {
              if ((anyNA(c(t[time, reg], t_alpha2_n2[time, reg])) == FALSE) &&
                  (t[time, reg] >= t_alpha2_n2[time, reg])) {
                sign_corr[time, reg] <- TRUE
              } else {
                sign_corr[time, reg] <- FALSE
              }
            }
          }  
          
          ## Save metric result in arrays
          aggr_metrics[,,which(metrics == met)] <- Reorder(data = enscorr, 
                                                           order = c('time', 'region'))
          aggr_significance[,,which(metrics == met)] <- Reorder(data = sign_corr, 
                                                                order = c('time', 'region'))
          
        } else if (met == 'mean_bias') {
          ## Calculate ensemble mean
          hcst_data_ens <- MeanDims(data$hcst$data, dims = 'ensemble')
          obs_data_ens <- MeanDims(data$obs$data, dims = 'ensemble')
          
          ## Aggregate data over regions
          hcst_data_aggr <- sapply(X = 1:length(regions), 
                                   FUN = function(X) {
                                     WeightedMean(data = hcst_data_ens,
                                                  region = regions[[X]],
                                                  lon = lon, londim = lon_dim,
                                                  lat = lat, latdim = lat_dim,
                                                  na.rm = na.rm)
                                   }, simplify = 'array')
          
          obs_data_aggr <- sapply(X = 1:length(regions), 
                                  FUN = function(X) {
                                    WeightedMean(data = obs_data_ens,
                                                 region = regions[[X]],
                                                 lon = lon, londim = lon_dim,
                                                 lat = lat, latdim = lat_dim,
                                                 na.rm = na.rm)
                                  }, simplify = 'array')
          
          ## Include name of region dimension
          names(dim(hcst_data_aggr))[length(dim(hcst_data_aggr))] <- 'region'
          names(dim(obs_data_aggr))[length(dim(obs_data_aggr))] <- 'region'
          
          ## Remove unnecessary dimension
          hcst_data_aggr <- Subset(hcst_data_aggr, c('dat','var', 'sday','sweek'),
                                   list(1,1,1,1) , drop = 'selected')
          obs_data_aggr <- Subset(obs_data_aggr, c('dat','var', 'sday','sweek'),
                                  list(1,1,1,1) , drop = 'selected')
          
          ## Calculate significance
          pval_mean_bias <- Apply(data = list(x = hcst_data_aggr, y = obs_data_aggr),
                                  target_dims = c('syear'), ncores = ncores,
                                  fun = function(x,y)
                                  {t.test(as.vector(x),as.vector(y))})$p.value
          
          sign_mean_bias <- pval_mean_bias <= alpha
          
          ## Calculate aggregated mean bias metric
          mean_bias <- sapply(X = 1:length(regions), 
                              FUN = function(X) {
                                WeightedMean(data = skill_metrics$mean_bias,
                                             region = regions[[X]],
                                             lon = lon, londim = lon_dim,
                                             lat = lat, latdim = lat_dim,
                                             na.rm = na.rm)
                              }, simplify = 'array')
          
          ## Include name of region dimension
          names(dim(mean_bias))[length(dim(mean_bias))] <- 'region'
          
          ## Remove 'var' dimension
          mean_bias <- Subset(mean_bias, 'var', 1, drop = 'selected')
          
          ## Save metric result in array
          aggr_metrics[ , , which(metrics == met)] <- Reorder(data = mean_bias,
                                                              order = c('time', 'region'))
          aggr_significance[ , , which(metrics == met)] <- Reorder(data = sign_mean_bias,
                                                                   order = c('time', 'region'))
          
        } else if (met == 'enssprerr') {
          ## Calculate metric
          spread <- sapply(X = 1:length(regions), 
                           FUN = function(X) {
                             WeightedMean(data = statistics$spread,
                                          region = regions[[X]],
                                          lon = lon, londim = lon_dim,
                                          lat = lat, latdim = lat_dim,
                                          na.rm = na.rm)
                           }, simplify = 'array')       
          
          error <- sapply(X = 1:length(regions), 
                          FUN = function(X) {
                            WeightedMean(data = skill_metrics$rms,
                                         region = regions[[X]],
                                         lon = lon, londim = lon_dim,
                                         lat = lat, latdim = lat_dim,
                                         na.rm = na.rm)
                          }, simplify = 'array')   
          
          ## Include name of region dimension
          names(dim(spread))[length(dim(spread))] <- 'region'
          names(dim(error))[length(dim(error))] <- 'region'
          
          ## Remove 'var' dimension
          spread <- Subset(spread, 'var', 1, drop = 'selected')
          error <- Subset(error, 'var', 1, drop = 'selected')
          
          enssprerr <- spread / error
          
          # ## Significance calculation
          # 
          # # Effective sample size
          # # exp <- data$hcst$data
          # # ens_exp <- MeanDims(data$hcst$data, dims = 'ensemble')
          # # enospr <- sum(Eno(exp - InsertDim(ens_exp, length(dim(exp)), dim(exp)['ensemble']), "ensemble"))
          # # enodif <- .Eno(ens_exp - ens_obs, na.action = na.pass)
          # 
          # # Removing eno at the moment
          # # F <- (enospr * spread^2 / (enospr - 1)) / (enodif * error^2 / (enodif - 1))
          # F <- spread^2 / error^2
          # # if (!is.na(F) & !is.na(enospr) & !is.na(enodif) & any(enospr > 2) & enodif > 2) {
          #   # p.val <- pf(F, enospr - 1, enodif - 1)
          # 
          #   pval_enssprerr <- Apply(data = list(x = F, y = data$hcst$data),  ### DOES NOT WORK
          #                           target_dims = c('time'),
          #                           fun = function(x, y)
          #                             {pf(x, dim(y)['syear'] -1 ,dim(y)['syear'] -1)})
          # 
          #   pval_enssprerr <- pf(F,dim(data$hcst$data)['syear'] -1 ,dim(data$hcst$data)['syear'] -1)
          #   pval_enssprerr <- 2 * min(pval_enssprerr, 1 - pval_enssprerr)
          # 
          # sign_enssprerr <- pval_enssprerr <= alpha
          
          ## Save metric result in array
          aggr_metrics[ , , which(metrics == met)] <- Reorder(data = enssprerr, 
                                                              order = c('time', 'region'))
          # aggr_significance[,,which(metrics == met)] <- Reorder(data = sign_corr, 
          #                                                       order = c('time', 'region'))
        }
      } ## close loop on metric
    } ## close if on score   
  } ## close if on region
  
  ## Set NAs to False 
  aggr_significance[is.na(aggr_significance)] <- FALSE
  
  ## Include 'var' dimension to be able to save array
  if (!'var' %in% names(dim(aggr_metrics))) {
    aggr_metrics <- InsertDim(aggr_metrics, 1, 1, name = 'var')
    aggr_significance <- InsertDim(aggr_significance, 1, 1, name = 'var')
  }
  
  ## Include attributes 
  attributes(aggr_metrics)$metrics <- metrics
  attributes(aggr_metrics)$forecast.months <- forecast.months
  attributes(aggr_metrics)$regions <- regions
  attributes(aggr_metrics)$system.name <- system
  attributes(aggr_metrics)$reference.name <- reference
  
  aggr_scorecards <- list(aggr_metrics = aggr_metrics,
                          aggr_significance = aggr_significance)
  
  ## Save metric data arrays
  recipe$Run$output_dir <- file.path(recipe$Run$output_dir,
                                     "outputs", "Scorecards/")
  save_metrics(recipe = recipe, metrics = aggr_scorecards,
               data_cube = data$hcst, agg = 'region',
               module = "scorecards")
}
