## TODO: Add the possibility to read the data directly from netCDF
## TODO: Adapt to multi-model case
## TODO: Adapt to multi-variable case
## TODO: Add param 'raw'?
## TODO: Decadal plot names

source("modules/Visualization/R/plot_skill_metrics.R")
source("modules/Visualization/R/get_proj_code.R")
## TODO: Remove after the next s2dv release
source("modules/Visualization/R/tmp/PlotRobinson.R")
source("modules/Visualization/R/plot_most_likely_terciles_map.R")
source("modules/Visualization/R/plot_ensemble_mean.R")

plot_data <- function(recipe,
                      data,
                      skill_metrics = NULL,
                      probabilities = NULL,
                      significance = F) {
  # Try to produce and save several basic plots.
  # recipe: the auto-s2s recipe as read by read_yaml()
  # data: list containing the hcst, obs and (optional) fcst s2dv_cube objects
  # calibrated_data: list containing the calibrated hcst and (optional) fcst
  # s2dv_cube objects
  # skill_metrics: list of arrays containing the computed skill metrics
  # significance: Bool. Whether to include significance dots where applicable

  plots <- strsplit(recipe$Analysis$Workflow$Visualization$plots, ", | |,")[[1]]
  recipe$Run$output_dir <- paste0(recipe$Run$output_dir, "/plots/")
  ## TODO: Sort this out
  outdir <- get_dir(recipe = recipe,
                    variable = data$hcst$attrs$Variable$varName)
  for (directory in outdir) {
   dir.create(directory, showWarnings = FALSE, recursive = TRUE)
  }

  if ((is.null(skill_metrics)) && (is.null(data$fcst))) { 
    error(recipe$Run$logger, "The Visualization module has been called,
          but there is no fcst in 'data', and 'skill_metrics' is NULL
          so there is no data that can be plotted.")
    stop()
  }
  # Set default single-panel plots if not specified
  if (is.null(recipe$Analysis$Workflow$Visualization$multi_panel)) {
    recipe$Analysis$Workflow$Visualization$multi_panel <- FALSE
  }
  # Plot skill metrics
  if ("skill_metrics" %in% plots) {
    if (!is.null(skill_metrics)) {
      plot_skill_metrics(recipe, data$hcst, skill_metrics, outdir,
                         significance)
    } else {
      error(recipe$Run$logger,
            paste0("The skill metric plots have been requested, but the ",
                   "parameter 'skill_metrics' is NULL"))
    }
  }

  # Plot forecast ensemble mean
  if ("forecast_ensemble_mean" %in% plots) {
    if (!is.null(data$fcst)) {
      plot_ensemble_mean(recipe, data$fcst, outdir)
    } else {
      error(recipe$Run$logger,
            paste0("The forecast ensemble mean plot has been requested, but ",
                   "there is no fcst element in 'data'"))
    }
  }

  # Plot Most Likely Terciles
  if ("most_likely_terciles" %in% plots) {
    if ((!is.null(probabilities)) && (!is.null(data$fcst))) {
      # Assign default parameters
      if (is.null(recipe$Analysis$Workflow$Visualization$mask_terciles)) {
        recipe$Analysis$Workflow$Visualization$mask_terciles <- FALSE
      }
      if (is.null(recipe$Analysis$Workflow$Visualization$dots)) {
        recipe$Analysis$Workflow$Visualization$dots <- FALSE
      }
      # Plots without masked terciles/dots
      if ((recipe$Analysis$Workflow$Visualization$mask_terciles 
           %in% c('both', FALSE)) || 
          (recipe$Analysis$Workflow$Visualization$dots
           %in% c('both', FALSE))) {
        plot_most_likely_terciles(recipe, data$fcst,
                                  probabilities,
                                  mask = NULL,
                                  dots = NULL,
                                  outdir)
      }
      # Plots with masked terciles
      if (recipe$Analysis$Workflow$Visualization$mask_terciles %in% 
          c('both', TRUE)) {
        if (is.null(skill_metrics)) {
          error(recipe$Run$logger,
                paste0("For the most likely terciles plot, skill_metrics ",
                       "need to be provided to be masked."))
        } else if (!('rpss' %in% names(skill_metrics))) {
          error(recipe$Run$logger,
                paste0("For the most likely terciles plot, rpss metric ",
                       "need to be provided to be masked"))
        } else {
          plot_most_likely_terciles(recipe, data$fcst,
                                    probabilities,
                                    mask = skill_metrics$rpss,
                                    dots = NULL,
                                    outdir)
        }
      }
      # Plots with dotted terciles
      if (recipe$Analysis$Workflow$Visualization$dots %in% c('both', TRUE)) {
        if (is.null(skill_metrics)) {
          error(recipe$Run$logger,
                paste0("For the most likely terciles plot, skill_metrics ",
                       "need to be provided for the dots"))
        } else if (!('rpss' %in% names(skill_metrics))) {
          error(recipe$Run$logger,
                paste0("For the most likely terciles plot, rpss metric ",
                       "needs to be provided for the dots"))
        } else {
          plot_most_likely_terciles(recipe, data$fcst,
                                    probabilities,
                                    mask = NULL,
                                    dots = skill_metrics$rpss,
                                    outdir)
        }
      }
    } else {
      error(recipe$Run$logger,
            paste0("For the most likely terciles plot, both the fcst and the ",
                   "probabilities must be provided."))
    }
  }
}
