library(stringr)

plot_skill_metrics <- function(recipe, skill_metrics,
                               outdir, significance = F, output_conf) {
  # recipe: Auto-S2S recipe
  # archive: Auto-S2S archive
  # cube_info: s2dv_cube object with the corresponding hindcast data
  # skill_metrics: list of named skill metrics arrays
  # outdir: output directory
  # significance: T/F, whether to display the significance dots in the plots

  # Abort if frequency is daily
  if (recipe$Analysis$Variables$freq %in% c("daily", "daily_mean")) {
    error(recipe$Run$logger, "Visualization functions not yet implemented
          for daily data.")
    stop()
  }
  # Abort if skill_metrics is not list
  if (!is.list(skill_metrics) || is.null(names(skill_metrics))) {
    stop("The element 'skill_metrics' must be a list of named arrays.")
  }

  cube_info <- skill_metrics[["metadata"]]
  skill_metrics[["metadata"]] <- NULL

  latitude <- cube_info$coords$lat
  longitude <- cube_info$coords$lon
  archive <- get_archive(recipe)
  system_name <- archive$System[[recipe$Analysis$Datasets$System$name]]$name
  hcst_period <- paste0(recipe$Analysis$Time$hcst_start, "-",
                        recipe$Analysis$Time$hcst_end)
  if (tolower(recipe$Analysis$Horizon) == "seasonal") {
    init_month <- as.numeric(substr(recipe$Analysis$Time$sdate,
                                    start = 1, stop = 2))
  } else {
    ## TODO: Sort out decadal initial month (is it always January?)
    init_month <- 1
  }
  month_label <- tolower(month.name[init_month])
  month_abbreviation <- month.abb[init_month]
  # Get months
  months <- lubridate::month(Subset(cube_info$attrs$Dates,
                                    "syear", indices = 1),
                             label = T, abb = F,locale = "en_GB")
  if (!is.null(recipe$Analysis$Workflow$Visualization$projection)) {
    projection <- tolower(recipe$Analysis$Workflow$Visualization$projection)
  } else {
    projection <- "cylindrical_equidistant"
  }

  # Define color palette and number of breaks according to output format
  if (tolower(recipe$Analysis$Output_format) %in% c("scorecards", "cerise")) {
    diverging_palette <- "purpleorange"
    sequential_palette <- "Oranges"
  } else {
    diverging_palette <- "bluered"
    sequential_palette <- "Reds"
  }
  # Group different metrics by type
  skill_scores <- c("rpss", "bss90", "bss10", "frpss", "crpss", "mean_bias_ss",
                    "enscorr", "rpss_specs", "bss90_specs", "bss10_specs",
                    "enscorr_specs", "rmsss", "msss")
  scores <- c("rps", "frps", "crps", "frps_specs", "mse")
  # Loop over variables and assign colorbar and plot parameters to each metric
  for (var in 1:cube_info$dims[['var']]) {
    var_skill <- lapply(skill_metrics, function(x) {
                        ClimProjDiags::Subset(x, along = 'var',
                                              indices = var,
                                              drop = 'selected')})
    for (name in c(skill_scores, scores, "mean_bias", "enssprerr")) {
      if (name %in% names(skill_metrics)) {
        units <- NULL
        # Define plot characteristics and metric name to display in plot
        if (name %in% c("rpss", "bss90", "bss10", "frpss", "crpss",
                        "rpss_specs", "bss90_specs", "bss10_specs",
                        "rmsss", "msss")) {
          display_name <- toupper(strsplit(name, "_")[[1]][1])
          skill <- var_skill[[name]]
          brks <- seq(-1, 1, by = 0.2)
          colorbar <- clim.colors(length(brks) + 1, diverging_palette)
          cols <- colorbar[2:(length(colorbar) - 1)]
          col_inf <- colorbar[1]
          col_sup <- NULL
        } else if (name == "mean_bias_ss") {
          display_name <- "Mean Bias Skill Score"
          skill <- var_skill[[name]]
          brks <- seq(-1, 1, by = 0.2)
          colorbar <- clim.colors(length(brks) + 1, diverging_palette)
          cols <- colorbar[2:(length(colorbar) - 1)]
          col_inf <- colorbar[1]
          col_sup <- NULL
        } else if (name %in% c("enscorr", "enscorr_specs")) {
          display_name <- "Ensemble Mean Correlation"
          skill <- var_skill[[name]]
          brks <- seq(-1, 1, by = 0.2)
          cols <- clim.colors(length(brks) - 1, diverging_palette)
          col_inf <- NULL
          col_sup <- NULL
        } else if (name %in% scores) {
          skill <- var_skill[[name]]
          display_name <- toupper(strsplit(name, "_")[[1]][1])
          brks <- seq(0, 1, by = 0.1)
          colorbar <- grDevices::hcl.colors(length(brks), sequential_palette)
          cols <- colorbar[1:(length(colorbar) - 1)]
          col_inf <- NULL
          col_sup <- colorbar[length(colorbar)]
        } else if (name == "enssprerr") {
          skill <- var_skill[[name]]
          display_name <- "Spread-to-Error Ratio"
          brks <- c(0, 0.6, 0.7, 0.8, 0.9, 1, 1.2, 1.4, 1.6, 1.8, 2)
          colorbar <- clim.colors(length(brks), diverging_palette)
          cols <- colorbar[1:length(colorbar) - 1]
          col_inf <- NULL
          col_sup <- colorbar[length(colorbar)]
        } else if (name %in% "mean_bias") {
          skill <- var_skill[[name]]
          display_name <- "Mean Bias"
          max_value <- max(abs(quantile(skill, 0.02, na.rm = T)),
                           abs(quantile(skill, 0.98, na.rm = T)))
          brks <- max_value * seq(-1, 1, by = 0.2)
          colorbar <- clim.colors(length(brks) + 1, diverging_palette)
          cols <- colorbar[2:(length(colorbar) - 1)]
          col_inf <- colorbar[1]
          col_sup <- colorbar[length(colorbar)]
          units <- cube_info$attrs$Variable$metadata[[var_name]]$units
        }
        # Reorder dimensions
        skill <- Reorder(skill, c("time", "longitude", "latitude"))
        # If the significance has been requested and the variable has it,
        # retrieve it and reorder its dimensions. 
        significance_name <- paste0(name, "_significance")
        if ((significance) && (significance_name %in% names(skill_metrics))) {
          skill_significance <- var_skill[[significance_name]]
          skill_significance <- Reorder(skill_significance, c("time", 
                                                              "longitude",
                                                              "latitude"))
          # Split skill significance into list of lists, along the time dimension
          # This allows for plotting the significance dots correctly.
          skill_significance <- ClimProjDiags::ArrayToList(skill_significance,
                                                           dim = "time",
                                                           level = "sublist",
                                                           names = "dots")
        } else {
          skill_significance <- NULL
        }
        # Define output file name and titles
        if (tolower(recipe$Analysis$Horizon) == "seasonal") {
          outfile <- paste0(outdir[var], name, "-", month_label)
        } else {
          outfile <- paste0(outdir[var], name)
        }
        # Get variable name and long name
        var_name <- cube_info$attrs$Variable$varName[[var]]
        var_long_name <- cube_info$attrs$Variable$metadata[[var_name]]$long_name
        # Multi-panel or single-panel plots
        if (recipe$Analysis$Workflow$Visualization$multi_panel) {
          # Define titles
          toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                             "\n", display_name, " / ", hcst_period)
          titles <- as.vector(months)
          ## TODO: Combine PlotLayout with PlotRobinson?
          suppressWarnings(
          PlotLayout(PlotEquiMap, c('longitude', 'latitude'),
                     asplit(skill, MARGIN=1), # Splitting array into a list
                     longitude, latitude,
                     special_args = skill_significance, 
                     dot_symbol = 20,
                     toptitle = toptitle,
                     title_scale = 0.6,
                     titles = titles,
                     filled.continents = F,
                     brks = brks,
                     cols = cols,
                     col_inf = col_inf,
                     col_sup = col_sup,
                     fileout = paste0(outfile, ".png"),
                     bar_label_digits = 3,
                     bar_extra_margin = rep(0.9, 4),
                     extra_margin = rep(1, 4),
                     bar_label_scale = 1.5,
                     axes_label_scale = 1.3,
                     width = 11,#default i
                     height = 11)
          )
        } else {
          # Define function and parameters depending on projection
          if (projection == 'cylindrical_equidistant') {
            fun <- PlotEquiMap
            output_configuration <- output_conf$PlotEquiMap$skill_metric
            base_args <- list(var = NULL, dots = NULL,
                              lon = longitude, lat = latitude,
                              dot_symbol = 20, dot_size = 1,
                              title_scale = 0.6,
                              filled.continents = F, brks = brks, cols = cols, 
                              col_inf = col_inf, col_sup = col_sup, 
                              units = units, font.main = 2,
                              bar_label_digits = 3, bar_label_scale = 1.5,
                              axes_label_scale = 1, width = 8, height = 5)
              base_args[names(output_configuration)] <- output_configuration
          } else {
            fun <- PlotRobinson
            common_projections <- c("robinson", "stereographic", "lambert_europe")
            if (projection %in% common_projections) {
              target_proj <- get_proj_code(projection)
            } else {
              target_proj <- projection
            }
            ## TODO: handle output_conf
            base_args <- list(data = NULL, mask = NULL,
                              lon = longitude, lat = latitude,
                              lon_dim = 'longitude', lat_dim = 'latitude',
                              target_proj = target_proj, legend = 's2dv',
                              style = 'point', brks = brks, cols = cols,
                              col_inf = col_inf, col_sup = col_sup)
          }
          # Loop over forecast times
          for (i in 1:dim(skill)[['time']]) {
            # Get forecast time label
            forecast_time <- match(months[i], month.name) - init_month + 1
                        
            if (forecast_time < 1) {
              forecast_time <- forecast_time + 12
            }
            forecast_time <- sprintf("%02d", forecast_time)
            # Define plot title
            toptitle <- paste(system_name, "/",
                              str_to_title(var_long_name),
                              "\n", display_name, "/", months[i], "/",
                              hcst_period)
            # Modify base arguments
            base_args[[1]] <- skill[i, , ]
            if (!is.null(skill_significance)) {
              base_args[[2]] <- skill_significance[[i]][[1]]
              significance_caption <- "alpha = 0.05"
            } else {
              significance_caption <- NULL
            }
            if (identical(fun, PlotRobinson)) {
              ## TODO: Customize alpha and other technical details depending on the metric
              base_args[['caption']] <- 
                paste0("Nominal start date: ",
                       "1st of ", str_to_title(month_label), "\n",
                       "Forecast month: ", forecast_time, "\n",
                       "Reference: ", recipe$Analysis$Datasets$Reference, "\n",
                       significance_caption)
            }
            fileout <- paste0(outfile, "_ft", forecast_time, ".png")
            # Plot

            do.call(fun,
                    args = c(base_args,
                             list(toptitle = toptitle,
                                  fileout = fileout)))
          }
        }
      }
    }
  }
  info(recipe$Run$logger,
       "##### SKILL METRIC PLOTS SAVED TO OUTPUT DIRECTORY #####")
}
