library(stringr)
library(lubridate)

plot_metrics <- function(recipe, data_cube, metrics,
                          outdir, significance = F, output_conf) {
  # recipe: Auto-S2S recipe
  # archive: Auto-S2S archive
  # data_cube: s2dv_cube object with the corresponding hindcast data
  # metrics: list of named metric arrays with named dimensions
  # 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 metrics is not list
  if (!is.list(metrics) || is.null(names(metrics))) {
    stop("The element 'metrics' must be a list of named arrays.")
  }
  
  latitude <- data_cube$coords$latitude
  longitude <- data_cube$coords$longitude
  archive <- get_archive(recipe)
  if (recipe$Analysis$Datasets$System$name == 'Multimodel'){
    system_name <- paste0('Multimodel-',
                          recipe$Analysis$Datasets$Multimodel$approach)
  } else {
    system_name <- archive$System[[recipe$Analysis$Datasets$System$name]]$name
  }
  hcst_period <- paste0(recipe$Analysis$Time$hcst_start, "-",
                        recipe$Analysis$Time$hcst_end)
  start_date <- recipe$Analysis$Time$sdate
  if (tolower(recipe$Analysis$Horizon) == "seasonal") {
    init_month <- as.numeric(substr(recipe$Analysis$Time$sdate,
                                    start = 1, stop = 2))
    months <- lubridate::month(Subset(data_cube$attrs$Dates,
                                    "syear", indices = 1),
                             label = T, abb = F,locale = "en_GB")
    month_label <- tolower(month.name[init_month])
    month_abbreviation <- month.abb[init_month]
  } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
    init_week <- as.character(recipe$Analysis$Time$sdate)
    wday_ini <- wday(ymd(init_week), week_start = 1)
    ftime_min <- recipe$Analysis$Time$ftime_min
    ftime_max <- recipe$Analysis$Time$ftime_max
    week_valid_ini <- ymd(init_week) + (1 - wday(ymd(init_week), week_start = 1)) +
      (ftime_min:ftime_max)*7
    week_valid_end <- week_valid_ini + 6
    # Forecast times:
    weeks <- paste0(0, ftime_min:ftime_max)
    # This week_label appears on the name of the file. It's just the start date.
    week_label <- recipe$Analysis$Time$sdate
  } else { # Decadal 
    init_month <- 1 
    init_week <- 1
    months <- lubridate::month(Subset(data_cube$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")
  statistics <- c("cov", "std_hcst", "std_obs", "var_hcst", "var_obs", "n_eff")

  # Loop over variables and assign colorbar and plot parameters to each metric
  for (var in 1:data_cube$dims[['var']]) {
    var_name <- data_cube$attrs$Variable$varName[[var]] ## For statistics
    var_metric <- lapply(metrics, function(x) {
                        ClimProjDiags::Subset(x, along = 'var',
                                              indices = var,
                                              drop = 'selected')})
    for (name in c(skill_scores, scores, statistics, "mean_bias", "enssprerr")) {
      if (name %in% names(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])
          metric <- var_metric[[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"
          metric <- var_metric[[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"
          metric <- var_metric[[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) {
          metric <- var_metric[[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") {
          metric <- var_metric[[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") {
          metric <- var_metric[[name]]
          display_name <- "Mean Bias"
          max_value <- max(abs(quantile(metric, 0.02, na.rm = T)),
                           abs(quantile(metric, 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 <- data_cube$attrs$Variable$metadata[[var_name]]$units
        } else if (name %in% "cov") {
          metric <- var_metric[[name]]
          display_name <- "Covariance"
          max_value <- max(abs(quantile(metric, 0.02, na.rm = T)),
                           abs(quantile(metric, 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 <- paste0(data_cube$attrs$Variable$metadata[[var_name]]$units, "²")
        } else if (name %in% "std_hcst") {
          metric <- var_metric[[name]]
          display_name <- "Hindcast Standard Deviation"
          max_value <- max(abs(quantile(metric, 0.02, na.rm = T)),
                           abs(quantile(metric, 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 <- data_cube$attrs$Variable$metadata[[var_name]]$units
        } else if (name %in% "std_obs") {
          metric <- var_metric[[name]]
          display_name <- "Observation Standard Deviation"
          max_value <- max(abs(quantile(metric, 0.02, na.rm = T)),
                           abs(quantile(metric, 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 <- data_cube$attrs$Variable$metadata[[var_name]]$units
        } else if (name %in% "n_eff") {
          metric <- var_metric[[name]]
          display_name <- "Effective Sample Size"
          max_value <- max(metric)
          brks <- max_value * seq(0, 1, by = 0.1)
          colorbar <- clim.colors(length(brks) + 1, diverging_palette)
          cols <- colorbar[2:(length(colorbar) - 1)]
          col_inf <- NULL
          col_sup <- NULL
        }
        
        # Reorder dimensions
        metric <- Reorder(metric, c("time", "longitude", "latitude"))
        # If the significance has been requested and the variable has it,
        # retrieve it and reorder its dimensions.
        if (significance != FALSE) { # Both, dots, mask or TRUE	
          significance_name <- paste0(name, "_significance")
          if ((significance_name %in% names(metrics))) {
            metric_significance <- var_metric[[significance_name]]
            metric_significance <- Reorder(metric_significance, c("time", 
                                                                  "longitude",
                                                                  "latitude"))
            # Split significance into list of lists, along the time dimension
            # This allows for plotting the significance dots correctly.
            metric_significance <- ClimProjDiags::ArrayToList(metric_significance,
                                                              dim = "time",
                                                              level = "sublist",
                                                              names = "dots")
          } else {
            metric_significance <- NULL
          }
        } else {
          metric_significance <- NULL
        }
        
        # Define output file name and titles
        if (tolower(recipe$Analysis$Horizon) == "seasonal") {
          outfile <- paste0(outdir[var], name, "-", month_label)
        } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
          outfile <- paste0(outdir[var], name, "-", week_label)
        } else {
          outfile <- paste0(outdir[var], name)
        }
        
        # Get variable name and long name
        var_name <- data_cube$attrs$Variable$varName[[var]]
        var_long_name <- data_cube$attrs$Variable$metadata[[var_name]]$long_name
        
        # Multi-panel or single-panel plots
        if (recipe$Analysis$Workflow$Visualization$multi_panel) {
          # Define titles
          if (tolower(recipe$Analysis$Horizon) %in% c("seasonal", "decadal")) {
            toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                             "\n", display_name, " / ", hcst_period)
            ## time_bounds in data_cube to know if Time_aggregation was applied
            if (!is.null(attributes(data_cube$attrs$time_bounds))) {
              info(recipe$Run$logger, "Using plotting attrs from time_bounds.")
              if (length(attributes(data_cube$attrs$time_bounds)$plotting_attr) > 1) {
                titles <- unlist(
                        lapply(1:length(attributes(data_cube$attrs$time_bounds)$plotting_attr$ini_ftime),
                               function(x) {
                               paste("Forecast time",
                                     attributes(data_cube$attrs$time_bounds)$plotting_attr$ini_ftime[x],
                                     "to",
                                     attributes(data_cube$attrs$time_bounds)$plotting_attr$end_ftime[x])}))
              } else {
                titles <- attributes(data_cube$attrs$time_bounds)$plotting_attr[[1]]
              }
            } else {
              titles <- as.vector(months)
            }
          } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
            toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                           "\n", display_name, " / ", "Issued on ",
                           format(ymd(start_date), "%d-%m-%Y"), 
                           " / ", hcst_period)
            titles <- paste("Valid from", format(week_valid_ini, "%d-%m"),
                            "to", format(week_valid_end, "%d-%m"))
	  } else {
            toptitle <- "Unknown"
            titles <- "Unknown"
          }
          ## TODO: Combine PlotLayout with PlotRobinson?
          output_configuration <- output_conf$multipanel$skill_metrics
          base_args <- list(fun = "PlotEquiMap",
                            plot_dims = c('longitude', 'latitude'),
                            var = asplit(metric, MARGIN = 1),
                            lon = longitude, lat = latitude,
                            special_args = metric_significance,
                            dot_symbol = 20, toptitle = toptitle,
                            title_scale = 0.6, titles = titles,
                            filled.continents = FALSE, brks = brks,
                            cols = cols, col_inf = col_inf, col_sup = col_sup,
                            fileout = paste0(outfile, ".pdf"),
                            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, height = 11)
          base_args[names(output_configuration)] <- output_configuration
          do.call(PlotLayout, base_args)
        } else {
          # Define function and parameters depending on projection
          if (projection == 'cylindrical_equidistant') {
            fun <- PlotEquiMap
            base_args <- list(var = NULL, dots = NULL,
                              lon = longitude, lat = latitude,
                              dot_symbol = 20, dot_size = 1,
                              title_scale = 0.6, margin_scale = c(1, 5, 5, 5),
                              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)
          } 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', dots = NULL, brks = brks, 
                              cols = cols,
                              col_inf = col_inf, col_sup = col_sup,
                              units = units,
                              dots_size = 0.2, dots_shape = 47)
          }
          output_configuration <- output_conf[[projection]]$skill_metrics
          base_args[names(output_configuration)] <- output_configuration
          # Loop over forecast times
          for (i in 1:dim(metric)[['time']]) {
            # Get forecast time label
            if (recipe$Analysis$Horizon == "seasonal") {
              # Case without time aggregation:
              if (is.null(attributes(data_cube$attrs$time_bounds))) {
                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)
              } else {
                if (length(attributes(data_cube$attrs$time_bounds)$plotting_attr) > 1) {
                  forecast_time_ini <-attributes(data_cube$attrs$time_bounds)$plotting_attr$ini[i]
                  forecast_time_end <- attributes(data_cube$attrs$time_bounds)$plotting_attr$end[i]
                  # labels for file name:
                  forecast_time <- paste0(forecast_time_ini, "-", forecast_time_end)
                  # title names:
                  ## TODO: There's probably a better way to do this, like using
                  ## time_bounds$start and time_bounds$end directly.
                  forecast_time_ini <- init_month + forecast_time_ini - 1
                  forecat_time_ini <- ifelse(forecast_time_ini > 12, forecast_time_ini - 12, forecast_time_ini) 
                  forecast_time_ini <- month.name[forecast_time_ini]
                  forecast_time_end <- init_month + forecast_time_end - 1
                  forecast_time_end <- ifelse(forecast_time_end > 12, forecast_time_end - 12, forecast_time_end)
                  forecast_time_end <- month.name[forecast_time_end]
                  toptitle <- paste(system_name, "/",
                                    str_to_title(var_long_name),
                                    "\n", display_name, "/", forecast_time_ini, "to", 
                                    forecast_time_end, "/",
                                    hcst_period)
                } else {
                  forecast_time <- attributes(data_cube$attrs$time_bounds)$plotting_attr[[1]][i]
                  toptitle <- paste(system_name, "/",
                                    str_to_title(var_long_name),
                                    "\n", display_name, "/", 
                                    forecast_time, "/",
                                    hcst_period)
                }
              }
              nominal_startdate_caption <- paste0("1st of ", str_to_title(month_label))
              forecast_time_caption <- paste0("Forecast month: ", forecast_time)
            } else if (tolower(recipe$Analysis$Horizon == "subseasonal")) {
              # Get forecast time label
              forecast_time <- weeks[i]
              toptitle <- paste(system_name, "/",
                                str_to_title(var_long_name),
                                "\n", display_name, " / ", "Issued on ",
                                format(ymd(start_date), "%d-%m-%Y"),
                                " / ", hcst_period, "\n",
                                paste("Valid from", format(week_valid_ini[i], "%d-%m"),
                                      "to", format(week_valid_end[i], "%d-%m"), "of",
                                      year(ymd(start_date))))
              
              nominal_startdate_caption <- ymd(start_date)
              forecast_time_caption <- paste0("Forecast week: ", sprintf("%02d", i))
            } else if (recipe$Analysis$Horizon == "decadal") {
              # Case without time aggregation:
              if (is.null(attributes(data_cube$attrs$time_bounds))) {
                forecast_time <- match(months[i], month.name) + floor((i - 1)/12)*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)
              } else {
                if (length(attributes(data_cube$attrs$time_bounds)$plotting_attr) > 1) {
                  forecast_time_ini <- attributes(data_cube$attrs$time_bounds)$plotting_attr$ini[i]
                  forecast_time_end <- attributes(data_cube$attrs$time_bounds)$plotting_attr$end[i]
                  # labels for file name:
                  forecast_time <- paste0(forecast_time_ini, "-", forecast_time_end)
                  # title names:
                  forecast_time_ini <- init_month + forecast_time_ini - 1
                  forecast_time_ini <- ifelse(forecast_time_ini > 12, forecast_time_ini - 12, forecast_time_ini) 
                  forecast_time_ini <- month.name[forecast_time_ini]
                  forecast_time_end <- init_month + forecast_time_end - 1
                  forecast_time_end <- ifelse(forecast_time_end > 12, forecast_time_end - 12, forecast_time_end)
                  forecast_time_end <- month.name[forecast_time_end]
                  toptitle <- paste(system_name, "/",
                                    str_to_title(var_long_name),
                                    "\n", display_name, "/", forecast_time_ini, "to", 
                                    forecast_time_end, "/",
                                    hcst_period)
                } else {
                  forecast_time <- attributes(data_cube$attrs$time_bounds)$plotting_attr[[1]][i]
                  toptitle <- paste(system_name, "/",
                                    str_to_title(var_long_name),
                                    "\n", display_name, "/", 
                                    forecast_time, "/",
                                    hcst_period)
                }
              }
              nominal_startdate_caption <- start_date
              forecast_time_caption <- paste0("Forecast month: ", forecast_time)
            } else {
              warning("Unknown time horizon?")
            }
            # Modify base arguments
            base_args[[1]] <- metric[i, , ]
            if (!is.null(metric_significance)) {
              sign_file_label <- NULL
              if (is.logical(significance)) {
                if (significance) {
                  base_args[[2]] <- metric_significance[[i]][[1]]
                  sign_file_label <- '_mask'
                }
              } else {
                if (significance == 'dots') {
                  if (projection != 'cylindrical_equidistant') {
                    base_args[[10]] <- metric_significance[[i]][[1]]
                  } else {
                    # The position of arguments in base_args requires this cond
                    # so PlotEquiMap plots dots when requested
                    base_args[[2]] <- metric_significance[[i]][[1]]
                  }
                  sign_file_label <- '_dots'
                } else if (significance == 'mask') {
                  base_args[[2]] <- metric_significance[[i]][[1]]
                  sign_file_label <- '_mask'
                }
              }
              significance_caption <- "alpha = 0.05"
            } else {
              significance_caption <- NULL
              sign_file_label <- NULL
            }
            if (identical(fun, PlotRobinson)) {
              ## TODO: Customize alpha and other technical details depending on the metric
              base_args[['caption']] <- 
                paste0("Nominal start date: ", nominal_startdate_caption, "\n",
                       forecast_time_caption, "\n",
                       "Reference: ", recipe$Analysis$Datasets$Reference, "\n",
                       "Units: ", data_cube$attrs$Variable$metadata[[var_name]]$units, "\n",
                       significance_caption)
            }
            fileout <- paste0(outfile, "_ft", forecast_time,
                              sign_file_label, ".pdf")
            # Plot
            info(recipe$Run$logger,
                 paste("Plotting", display_name))

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