## 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_metrics.R")
source("modules/Visualization/R/get_proj_code.R")
source("modules/Visualization/R/get_plot_time_labels.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")

Visualization <- function(recipe,
                          data,
                          skill_metrics = NULL,
                          statistics = NULL,
                          probabilities = NULL,
                          significance = F,
                          output_conf = NULL) {
  # 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

  # Try to set default configuration if not specified by user
  if (is.null(output_conf) && !is.null(recipe$Analysis$Region$name)) {
    output_conf <- read_yaml("modules/Visualization/output_size.yml",
                             eval.exp = TRUE)$region
    if (recipe$Analysis$Region$name %in% names(output_conf)) {
      output_conf <- output_conf[[recipe$Analysis$Region$name]]
    } else {
      warn(recipe$Run$logger,
           paste0("The region name is not found in the 'output_conf' file. ",
                  "The default plot settings will be used."))
      output_conf <- NULL
    }
  # If the user chooses to specify the configuration, warn them.
  } else {
    if (is.list(output_conf)) {
      warning(paste0("Parameter 'output_conf' should be a list matching the ",
                     "parameters of the requested plotting function, i.e. ",
                     "PlotEquiMap, PlotRobinson or PlotLayout. There could be ",
                     "plotting erros if the list is incomplete."))
    } else if (!is.null(output_conf)) {
      warning(paste("Parameter 'output_conf' should be a list.",
                    "Using default configuration."))
      output_conf <- NULL
    }
  }

  # Get plot types and create output directories
  plots <- strsplit(recipe$Analysis$Workflow$Visualization$plots, ", | |,")[[1]]
  ## TODO: Do not modify output dir here
  recipe$Run$output_dir <- paste0(recipe$Run$output_dir, "/plots/")
  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(statistics)) && 
      (is.null(data$fcst))) { 
    error(recipe$Run$logger, 
          paste0("The Visualization module has been called, but there is no ",
                 "fcst in 'data', and 'skill_metrics' and 'statistics' are ",
                 "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
  }
  
  # Warning if significance parameter not included in function call
  if (!missing(significance) && !is.null(recipe$Analysis$Workflow$Visualization$significance)) {
    if (significance != recipe$Analysis$Workflow$Visualization$significance) {
      warning("The significance value provided in the function call does not match ",
              "the value in the recipe. The function call value will be used.")
    }
  }
  
  # If not specified in function call, set significance to recipe value
  if (missing(significance) && !is.null(recipe$Analysis$Workflow$Visualization$significance)) {
    significance <- recipe$Analysis$Workflow$Visualization$significance
  }
  
  # Plot skill metrics
  if ("skill_metrics" %in% plots) {
    if (!is.null(skill_metrics)) {
      if (is.logical(significance)) {
        plot_metrics(recipe = recipe, data_cube = data$hcst, 
                     metrics = skill_metrics, outdir = outdir,
                     significance = significance, output_conf = output_conf)
        info(recipe$Run$logger,
             paste("Skill metrics significance set as", significance))
        } else {
          if (significance %in% c('both', 'dots')) {
            plot_metrics(recipe, data$hcst, skill_metrics, outdir,
                         significance = 'dots', output_conf = output_conf)
            info(recipe$Run$logger, "Skill metrics significance as dots")
            }
          if (significance %in% c('both', 'mask')) {
            plot_metrics(recipe, data$hcst, skill_metrics, outdir,
                         significance = 'mask', output_conf = output_conf)
            info(recipe$Run$logger, "Skill metrics significance as mask")
          }
        }
    } else {
      error(recipe$Run$logger,
            paste0("The skill metric plots have been requested, but the ",
                   "parameter 'skill_metrics' is NULL"))
    }
  }
  
  # Plot statistics
  if ("statistics" %in% plots) {
    if (!is.null(statistics)) {
      plot_metrics(recipe, data$hcst, statistics, outdir,
                   significance, output_conf = output_conf)
    } else {
      error(recipe$Run$logger,
            paste0("The statistics plots have been requested, but the ",
                   "parameter 'statistics' is NULL"))
    }
  }

  # Plot forecast ensemble mean
  if ("forecast_ensemble_mean" %in% plots) {
    if (!is.null(data$fcst)) {
      if (is.null(recipe$Analysis$Workflow$Visualization$mask_ens)) {
        recipe$Analysis$Workflow$Visualization$mask_ens <- FALSE
      }
      if (is.null(recipe$Analysis$Workflow$Visualization$dots)) {
        recipe$Analysis$Workflow$Visualization$dots <- FALSE
      }
      # Plot without mask or dots
      if ((recipe$Analysis$Workflow$Visualization$mask_ens
           %in% c('both', FALSE)) ||
          (recipe$Analysis$Workflow$Visualization$dots
           %in% c('both', FALSE))) {
        plot_ensemble_mean(recipe, data$fcst, outdir, 
                           mask = NULL, dots = NULL,
                           output_conf = output_conf)
      }
      # Plots with masked
      if (recipe$Analysis$Workflow$Visualization$mask_ens %in%
          c('both', TRUE)) {
        if (is.null(skill_metrics)) {
          error(recipe$Run$logger,
                paste0("For the forecast ensemble mean  plot, skill_metrics ",
                       "need to be provided to be masked."))
        } else if (!('enscorr' %in% names(skill_metrics))) {
          error(recipe$Run$logger,
                paste0("For the forecast ensemble mean plot, enscor metric ",
                       "need to be provided to be masked"))
        } else {
          plot_ensemble_mean(recipe, data$fcst,
                             mask = skill_metrics$enscorr,
                             dots = NULL,
                             outdir, output_conf = output_conf)
        }
      }
      # Plots with dotted negative correlated in ens-mean-fcst
      if (recipe$Analysis$Workflow$Visualization$dots %in% c('both', TRUE)) {
        if (is.null(skill_metrics)) {
          error(recipe$Run$logger,
                paste0("For the forecast ensemble mean plot, skill_metrics ",
                       "need to be provided for the dots"))
        } else if (!('enscorr' %in% names(skill_metrics))) {
          error(recipe$Run$logger,
                paste0("For the forecast ensemble mean plot, enscor metric ",
                       "needs to be provided for the dots"))
        } else {
          plot_ensemble_mean(recipe, data$fcst,
                             mask = NULL,
                             dots = skill_metrics$enscorr,
                             outdir, output_conf = output_conf)
        }
      }

    } 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, output_conf = output_conf)
      }
      # 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, output_conf = output_conf)
        }
      }
      # 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, output_conf = output_conf)
        }
      }
    } else {
      error(recipe$Run$logger,
            paste0("For the most likely terciles plot, both the fcst and the ",
                   "probabilities must be provided."))
    }
  }

  # Convert plots to user-chosen format
  if (!is.null(recipe$Analysis$Workflow$Visualization$file_format) &&
      tolower(recipe$Analysis$Workflow$Visualization$file_format) != "pdf") {
    extension <- tolower(recipe$Analysis$Workflow$Visualization$file_format)
    plot_files <- list.files(path = recipe$Run$output_dir, pattern = "\\.pdf$",
                             recursive = TRUE, full.names = TRUE)
    for (file in plot_files) {
      system_command <- paste("convert -density 300", file,
                              "-resize 40% -alpha remove",
                              paste0(tools::file_path_sans_ext(file), ".",
                                     extension))
      system(system_command)
    }
    unlink(plot_files)
    info(recipe$Run$logger,
         paste0("##### PLOT FILES CONVERTED TO ", toupper(extension), " #####"))
  }
}
