## 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")
source("modules/Visualization/R/get_variable_longname.R")
source("modules/Visualization/R/plot_most_likely_terciles_map.R")
source("modules/Visualization/R/plot_forecast_map.R")
source("modules/Visualization/R/plot_extreme_probs.R")
## TODO: Remove after esviz release
source("modules/Visualization/R/tmp/VizEquiMap.R")
source("modules/Visualization/R/tmp/VizRobinson.R")
source("modules/Visualization/R/tmp/PlotRobinson.R")
source("modules/Visualization/R/tmp/zzz.R")
source("modules/Visualization/R/tmp/ColorBarContinuous.R")
source("modules/Visualization/R/tmp/ClimPalette.R")

Visualization <- function(recipe,
                          data,
                          skill_metrics = NULL,
                          statistics = NULL,
                          probabilities = NULL,
                          significance = F,
                          output_conf = NULL,
                          logo = 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 <- file.path(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
  }

  # If not specified in function call, set logo to recipe value
  if (missing(logo) && !is.null(recipe$Analysis$Workflow$Visualization$logo)) {
    logo <- recipe$Analysis$Workflow$Visualization$logo
  }
  
  # 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,
                     logo = logo)
        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,
                   logo = logo)
    } else {
      error(recipe$Run$logger,
            paste0("The statistics plots have been requested, but the ",
                   "parameter 'statistics' is NULL."))
    }
  }

  # Plot forecast ensemble mean
  if ("forecast_map" %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
      }
      if (is.null(recipe$Analysis$Workflow$Visualization$forecast_method)) {
        recipe$Analysis$Workflow$Visualization$forecast_method <- "median"
      } 
      # Loop over the methods requested
      for (method in strsplit(recipe$Analysis$Workflow$Visualization$forecast_method, ", | |,")[[1]]) {
        # 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_forecast_map(recipe, data$fcst,
                            mask = NULL, dots = NULL,
                            outdir = outdir, method = method,
                            output_conf = output_conf,
                            logo = logo)
        }
        # 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 ",
                         "needs 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 ",
                         "needs to be provided to be masked."))
          } else {
            plot_forecast_map(recipe, data$fcst,
                              mask = skill_metrics$enscorr,
                              dots = NULL,
                              outdir = outdir, method = method,
                              output_conf = output_conf,
                              logo = logo)
          }
        }
        # 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 ",
                         "needs 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_forecast_map(recipe, data$fcst,
                              mask = NULL,
                              dots = skill_metrics$enscorr,
                              outdir = outdir, method = method,
                              output_conf = output_conf,
                              logo = logo)
          }
        }
      } ## End loop over methods
    } else {
      error(recipe$Run$logger,
            paste0("The forecast ensemble mean plot has been requested, but ",
                   "there is no fcst element in 'data'."))
    }
  }
  # Probability plots
  ## TODO: Remove in the future
  if ("probs_fcst" %in% names(probabilities)) {
    names(probabilities)[which(names(probabilities) == "probs_fcst")] <- "fcst"    
  }
  # 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,
                                  probs$fcst,
                                  mask = NULL,
                                  dots = NULL,
                                  outdir, output_conf = output_conf,
                                  logo = logo)
      }
      # 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 ",
                       "needs 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 ",
                       "needs to be provided to be masked."))
        } else {
          plot_most_likely_terciles(recipe, data$fcst,
                                    probabilities$fcst,
                                    mask = skill_metrics$rpss,
                                    dots = NULL,
                                    outdir, output_conf = output_conf,
                                    logo = logo)
        }
      }
      # 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 ",
                       "needs 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$fcst,
                                    mask = NULL,
                                    dots = skill_metrics$rpss,
                                    outdir, output_conf = output_conf,
                                    logo = logo)
        }
      }
    } else {
      error(recipe$Run$logger,
            paste0("For the most likely terciles plot, both the fcst and the ",
                   "probabilities must be provided."))
    }
  }
 
  # Plot extreme probabilities
  if ("extreme_probabilities" %in% plots) {
    if ((!is.null(probabilities)) && (!is.null(data$fcst))) {
      # Assign default parameters
      if (is.null(recipe$Analysis$Workflow$Visualization$brks_min)) {
        recipe$Analysis$Workflow$Visualization$brks_min <- 5
      }
      if (is.null(recipe$Analysis$Workflow$Visualization$brks_max)) {
        recipe$Analysis$Workflow$Visualization$brks_max <- 95
      } 
      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 mask/dots
      if ((recipe$Analysis$Workflow$Visualization$mask_terciles 
           %in% c('both', FALSE)) || 
          (recipe$Analysis$Workflow$Visualization$dots
           %in% c('both', FALSE))) {
        plot_extreme_probs(recipe = recipe, fcst = data$fcst,
                           probabilities = probabilities$fcst,
                           mask = NULL,
                           dots = NULL,
                           outdir, output_conf = output_conf,
                           brks_min = recipe$Analysis$Workflow$Visualization$brks_min,
                           brks_max = recipe$Analysis$Workflow$Visualization$brks_max,
                           logo = logo)
      }
      # Plots with mask
      if (recipe$Analysis$Workflow$Visualization$mask_terciles %in% 
          c('both', TRUE)) {
        if (is.null(skill_metrics)) {
          error(recipe$Run$logger,
                paste0("For the extreme probabilities plot, skill_metrics ",
                       "needs to be provided to be masked."))
        } else if (!('rpss' %in% names(skill_metrics))) {
          error(recipe$Run$logger,
                paste0("For the extreme probabilities plot, rpss metric ",
                       "need to be provided to be masked."))
        } else {
          plot_extreme_probs(recipe = recipe, fcst = data$fcst,
                             probabilities = probabilities$fcst,
                             mask = skill_metrics$rpss,
                             dots = NULL,
                             outdir, output_conf = output_conf,
                             brks_min = recipe$Analysis$Workflow$Visualization$brks_min,
                             brks_max = recipe$Analysis$Workflow$Visualization$brks_max,
                             logo = logo)
        }
      }
      # Plots with dots
      if (recipe$Analysis$Workflow$Visualization$dots %in% c('both', TRUE)) {
        if (is.null(skill_metrics)) {
          error(recipe$Run$logger,
                paste0("For the extreme probabilities plot, skill_metrics ",
                       "needs to be provided for the dots."))
        } else if (!('rpss' %in% names(skill_metrics))) {
          error(recipe$Run$logger,
                paste0("For the extreme probabilities plot, rpss metric ",
                       "needs to be provided for the dots."))
        } else {
          plot_extreme_probs(recipe = recipe, fcst = data$fcst,
                             probabilities = probabilities$fcst,
                             mask = NULL,
                             dots = skill_metrics$rpss,
                             outdir, output_conf = output_conf,
                             brks_min = recipe$Analysis$Workflow$Visualization$brks_min,
                             brks_max = recipe$Analysis$Workflow$Visualization$brks_max,
                             logo = logo)
        }
      }
      
    } else {
      error(recipe$Run$logger,
            paste0("For the extreme probabilities plot, both the fcst and the ",
                   "probabilities must be provided."))
    }
  }
}

