###############################################################################
##################### SCORECARDS MODULE FOR SUNSET SUITE ######################
###############################################################################

##### Load source functions #####
source('modules/Scorecards/R/tmp/LoadMetrics.R')
source('modules/Scorecards/R/tmp/WeightedMetrics.R')
source('modules/Scorecards/R/tmp/Utils.R')
source('modules/Scorecards/R/tmp/SCTransform.R')
source('modules/Scorecards/R/tmp/ScorecardsSingle.R')
source('modules/Scorecards/R/tmp/ScorecardsMulti.R')
source('modules/Scorecards/R/tmp/ScorecardsSystemDiff.R')
source('modules/Scorecards/R/tmp/SCPlotScorecard.R') 


## TODO: Change function name to 'Scorecards'?
## Define function
Scorecards <- function(recipe) {
  
  ## Parameters for loading data
  input.path <- "/esarchive/scratch/nmilders/scorecards_data/test/output_test/" #paste0(recipe$Run$output_dir, "/outputs/Skill/")
  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
  calib.method <- tolower(recipe$Analysis$Workflow$Calibration$method)
  
  # NOTE (Eva): This condition needs to be checked, in my case 
  #             (recipe$Analysis$Workflow$Scorecards$start_months = NULL)
  start.months <- 1:12 # I added this line
  # Needs to be corrected:
  # if (recipe$Analysis$Workflow$Scorecards$start_months == 'all') {
  #   start.months <- 1:12
  # } else {
  #   start.months <- as.numeric(strsplit(recipe$Analysis$Workflow$Scorecards$start_months,
  #                                       split = ", | |,")[[1]])
  # }
  
  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.load <- unlist(strsplit(tolower(recipe$Analysis$Workflow$Skill$metric), ", | |,"))
  metrics.visualize <- unlist(strsplit(tolower(recipe$Analysis$Workflow$Scorecards$metric), ", | |,"))

  
  ## Parameters for scorecard table
  inf.to.na <-  recipe$Analysis$Workflow$Scorecards$inf.to.na
  table.label <-  recipe$Analysis$Workflow$Scorecards$table_label
  fileout.label <- recipe$Analysis$Workflow$Scorecards$fileout_label
  legend.white.space <- recipe$Analysis$Workflow$Scorecards$legend_white_space 
  col1.width <- recipe$Analysis$Workflow$Scorecards$col1_width
  col2.width <- recipe$Analysis$Workflow$Scorecards$col2_width
  calculate.diff <- recipe$Analysis$Workflow$Scorecards$calculate_diff
  ncores <- 1 # recipe$Analysis$ncores
  
  ## Load data files
  start.months <- sprintf("%02d", start.months)
  period <- paste0(start.year, "-", end.year)

  ####### SKILL AGGREGATION #######
  if(metric.aggregation == 'skill'){
  
    ## Load data files
    loaded_metrics <- LoadMetrics(system = system,
                                  reference = reference,
                                  var = var,
                                  start.year = start.year,
                                  end.year = end.year,
                                  metrics =  metrics.load,
                                  start.months = start.months,
                                  forecast.months = forecast.months,
                                  inf.to.na = inf.to.na,
                                  input.path = input.path)
    
    ## Spatial Aggregation of metrics
    if('region' %in% names(dim(loaded_metrics[[1]][[1]]))){
      
      ### Convert loaded metrics to array for already aggregated data
      metrics.dim <- attributes(loaded_metrics[[1]][[1]])$metrics 
      forecast.months.dim <- attributes(loaded_metrics[[1]][[1]])$forecast.months 
      start.months.dim <- attributes(loaded_metrics[[1]][[1]])$start.months 
      regions.dim <- regions #list('NAO' = c(lon.min = -80, lon.max = 40, lat.min = 20, lat.max = 80))
      
      aggregated_metrics <- array(dim = c(system = length(loaded_metrics),
                                          reference = length(loaded_metrics[[1]]),
                                          metric = length(metrics.dim), 
                                          time = length(forecast.months.dim), 
                                          sdate = length(start.months.dim), 
                                          region = length(regions.dim)))
      
      
      for (sys in 1:length(names(loaded_metrics))){
        for (ref in 1:length(names(loaded_metrics[[sys]]))){
          aggregated_metrics[sys, ref, , , , ] <- s2dv::Reorder(data = loaded_metrics[[sys]][[ref]], order = c('metric','time','sdate','region'))
        }
      }
      
      ## Add attributes
      attributes(aggregated_metrics)$metrics <- metrics.load
      attributes(aggregated_metrics)$start.months <- attributes(loaded_metrics[[1]][[1]])$start.months
      attributes(aggregated_metrics)$forecast.months <- attributes(loaded_metrics[[1]][[1]])$forecast.months
      attributes(aggregated_metrics)$regions <- regions
      attributes(aggregated_metrics)$system.name <- names(loaded_metrics)
      attributes(aggregated_metrics)$reference.name <- names(loaded_metrics[[1]])
      
      
    } else {
      ## Calculate weighted mean of spatial aggregation
      aggregated_metrics <- WeightedMetrics(loaded_metrics,
                                            regions = regions,
                                            metric.aggregation = metric.aggregation,
                                            ncores = ncores)
    }## close if on region
  } 
  
  ###### SCORE AGGREGATION ######
  if(metric.aggregation == 'score'){
    
    lon_dim_name <- 'longitude'
    lat_dim_name <- 'latitude'
    
    if('rpss' %in% metrics.visualize){
      ## Load data from saved files
      rps_syear <- .loadmetrics(input_path = input.path, system = system,
                                reference = reference, var = var, 
                                period = period, start_months = start.months,
                                forecast_months = forecast.months, 
                                metrics = 'rps_syear', 
                                calib_method = calib.method, syear = TRUE)
      
      rps_clim_syear <- .loadmetrics(input_path = input.path, system = system,
                                      reference = reference, var = var, 
                                      period = period, start_months = start.months,
                                      forecast_months = forecast.months, 
                                      metrics = 'rps_clim_syear', 
                                      calib_method = calib.method, syear = TRUE)
      
      ## Remove dat and var dimensions
      rps_syear <- Subset(rps_syear, along = c('dat', 'var'), indices = list(1,1) , drop = 'selected')
      rps_clim_syear <- Subset(rps_clim_syear, along = c('dat', 'var'), indices = list(1,1) , drop = 'selected')
      
    ## Calculate skill score value to show in scorecard 
      
      ## Spatially aggregate rps data
      rps_syear_spatial_aggr <- sapply(X = 1:length(regions), 
                                       FUN = function(X) {
                                         WeightedMean(data = rps_syear,
                                                      region = regions[[X]],
                                                      lon = as.vector(attributes(rps_syear)$Variables$dat1$longitude),
                                                      lat = as.vector(attributes(rps_syear)$Variables$dat1$latitude),
                                                      londim = lon_dim_name,
                                                      latdim = lat_dim_name,
                                                      na.rm = F)
                                       }, simplify = 'array')
      
      ## Include name of region dimension
      names(dim(rps_syear_spatial_aggr))[length(dim(rps_syear_spatial_aggr))] <- 'region'
      
      ## Spatially aggregate rps_clim data
      rps_clim_syear_spatial_aggr <- sapply(X = 1:length(regions), 
                                       FUN = function(X) {
                                         WeightedMean(data = rps_clim_syear,
                                                      region = regions[[X]],
                                                      lon = as.vector(attributes(rps_clim_syear)$Variables$dat1$longitude),
                                                      lat = as.vector(attributes(rps_clim_syear)$Variables$dat1$latitude),
                                                      londim = lon_dim_name,
                                                      latdim = lat_dim_name,
                                                      na.rm = F)
                                       }, simplify = 'array')
      
      ## Include name of region dimension
      names(dim(rps_clim_syear_spatial_aggr))[length(dim(rps_clim_syear_spatial_aggr))] <- 'region'

      ## Temporally aggregate data
      rps_temp_aggr <- Apply(data = rps_syear_spatial_aggr, 
                             target_dims = 'syear', 
                             fun = 'mean', ncores = ncores)$output1
      
      rps_clim_temp_aggr <- Apply(data = rps_clim_syear_spatial_aggr, 
                                  target_dims = 'syear',
                                  fun = 'mean',  ncores = ncores)$output1
      
      
    ## Calculate significance
      sign_rpss <- RandomWalkTest(rps_syear_spatial_aggr, rps_clim_syear_spatial_aggr, 
                                  time_dim = 'syear', test.type = 'two.sided',
                                  alpha = 0.05, pval = FALSE, sign = TRUE, 
                                  ncores = NULL)$sign
    
    } ## close if on rpss
  
    if('crpss' %in% metrics.visualize){
      ## Load data from saved files
      crps_syear <- .loadmetrics(input_path = input.path, system = system,
                                reference = reference, var = var, 
                                period = period, start_months = start.months,
                                forecast_months = forecast.months, 
                                metrics = 'crps_syear', 
                                calib_method = calib.method, syear = TRUE)
      
      crps_clim_syear <- .loadmetrics(input_path = input.path, system = system,
                                     reference = reference, var = var, 
                                     period = period, start_months = start.months,
                                     forecast_months = forecast.months, 
                                     metrics = 'crps_clim_syear', 
                                     calib_method = calib.method, syear = TRUE)
      
      ## Remove dat and var dimensions
      crps_syear <- Subset(crps_syear, along = c('dat', 'var'), indices = list(1,1) , drop = 'selected')
      crps_clim_syear <- Subset(crps_clim_syear, along = c('dat', 'var'), indices = list(1,1) , drop = 'selected')
      
      ## Calculate skill score value to show in scorecard 
      
      ## Spatially aggregate crps data
      crps_syear_spatial_aggr <- sapply(X = 1:length(regions), 
                                       FUN = function(X) {
                                         WeightedMean(data = crps_syear,
                                                      region = regions[[X]],
                                                      lon = as.vector(attributes(crps_syear)$Variables$dat1$longitude),
                                                      lat = as.vector(attributes(crps_syear)$Variables$dat1$latitude),
                                                      londim = lon_dim_name,
                                                      latdim = lat_dim_name,
                                                      na.rm = F)
                                       }, simplify = 'array')
      
      ## Include name of region dimension
      names(dim(crps_syear_spatial_aggr))[length(dim(crps_syear_spatial_aggr))] <- 'region'
      
      ## Spatially aggregate crps_clim data
      crps_clim_syear_spatial_aggr <- sapply(X = 1:length(regions), 
                                            FUN = function(X) {
                                              WeightedMean(data = crps_clim_syear,
                                                           region = regions[[X]],
                                                           lon = as.vector(attributes(crps_clim_syear)$Variables$dat1$longitude),
                                                           lat = as.vector(attributes(crps_clim_syear)$Variables$dat1$latitude),
                                                           londim = lon_dim_name,
                                                           latdim = lat_dim_name,
                                                           na.rm = F)
                                            }, simplify = 'array')
      
      ## Include name of region dimension
      names(dim(crps_clim_syear_spatial_aggr))[length(dim(crps_clim_syear_spatial_aggr))] <- 'region'
      
      ## Temporally aggregate data
      crps_temp_aggr <- Apply(data = crps_syear_spatial_aggr, 
                             target_dims = 'syear', 
                             fun = 'mean', ncores = ncores)$output1
      
      crps_clim_temp_aggr <- Apply(data = crps_clim_syear_spatial_aggr, 
                                  target_dims = 'syear',
                                  fun = 'mean', ncores = ncores)$output1
      
      
      ## Calculate significance
      sign_crpss <- RandomWalkTest(crps_syear_spatial_aggr, crps_clim_syear_spatial_aggr, 
                                  time_dim = 'syear', test.type = 'two.sided',
                                  alpha = 0.05, pval = FALSE, sign = TRUE, 
                                  ncores = NULL)$sign
      
    } ## close if on crpss
  } ## close if on score   

  ## Create simple scorecard tables 
  ## (one system only)
  ## Metrics input must be in the same order as function SC_spatial_aggregation
  scorecard_single <- ScorecardsSingle(data = aggregated_metrics,
                                       system = system,
                                       reference = reference,
                                       var = var,
                                       start.year = start.year,
                                       end.year = end.year,
                                       start.months = start.months,
                                       forecast.months = forecast.months,
                                       region.names = names(regions),
                                       metrics = metrics.visualize,
                                       table.label =  table.label,
                                       fileout.label = fileout.label,
                                       legend.white.space = legend.white.space, 
                                       col1.width = col1.width,
                                       col2.width = col2.width,
                                       output.path = output.path)
  
  ## Create multi system/reference scorecard tables
  ## (multiple systems with one reference or one system with multiple references)
  ## Metrics input must be in the same order as function SC_spatial_aggregation
  if(length(system) >  1 || length(reference) >  1){
    scorecard_multi <- ScorecardsMulti(data = aggregated_metrics,
                                       system = system,
                                       reference = reference,
                                       var = var,
                                       start.year = start.year,
                                       end.year = end.year,
                                       start.months = start.months,
                                       forecast.months = forecast.months,
                                       region.names = attributes(regions)$names,
                                       metrics = metrics.visualize,
                                       table.label =  table.label,
                                       fileout.label = fileout.label,
                                       output.path = output.path)
  } ## close if
  
  
  if(calculate.diff == TRUE){
    if(length(system) ==  2 || length(reference) ==  2){
      scorecard_diff <- ScorecardsSystemDiff(data = aggregated_metrics,
                                             system = system,
                                             reference = reference,
                                             var = var,
                                             start.year = start.year,
                                             end.year = end.year,
                                             start.months = start.months,
                                             forecast.months = forecast.months,
                                             region.names = attributes(regions)$names,
                                             metrics = metrics.visualize,
                                             table.label = table.label,
                                             fileout.label =  fileout.label,
                                             legend.white.space = legend.white.space, 
                                             col1.width = col1.width,
                                             col2.width = col2.width,
                                             output.path = output.path)
    } else {stop ("Difference scorecard can only be computed with two systems or two references.")}
  } ## close if on calculate.diff
  
}

