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

plot_data <- function(recipe,
		      data,
		      skill_metrics = NULL,
		      probabilities = NULL,
		      archive = NULL,
		      significance = F) {
  # Try to produce and save several basic plots.
  # recipe: the auto-s2s recipe as read by read_yaml()
  # archive: the auto-s2s archive 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

  outdir <- paste0(get_dir(recipe), "/plots/")
  dir.create(outdir, 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()
  }

  if (is.null(archive)) {
    if (tolower(recipe$Analysis$Horizon) == "seasonal") {
      archive <- 
	read_yaml(paste0("conf/archive.yml"))[[recipe$Run$filesystem]]
    } else if (tolower(recipe$Analysis$Horizon) == "decadal") {
      archive <- 
	read_yaml(paste0("conf/archive_decadal.yml"))[[recipe$Run$filesystem]]
    }
  }

  # Plot skill metrics
  if (!is.null(skill_metrics)) {
    plot_skill_metrics(recipe, archive, data$hcst, skill_metrics, outdir,
		       significance)
  }

  # Plot forecast ensemble mean
  if (!is.null(data$fcst)) {
    plot_ensemble_mean(recipe, archive, data$fcst, outdir)
  }

  # Plot Most Likely Terciles
  if ((!is.null(probabilities)) && (!is.null(data$fcst))) {
    plot_most_likely_terciles(recipe, archive, data$fcst,
			      probabilities, outdir)
  }
}

plot_skill_metrics <- function(recipe, archive, data_cube, skill_metrics,
			       outdir, significance = F) {
  # recipe: Auto-S2S recipe
  # archive: Auto-S2S archive
  # data_cube: 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

  ## TODO: OPTION for CERISE: Using PuOr
  # Abort if frequency is daily
  if (recipe$Analysis$Variables$freq == "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.")
  }
  
  latitude <- data_cube$coords$lat
  longitude <- data_cube$coords$lon
  system_name <- archive$System[[recipe$Analysis$Datasets$System$name]]$name
  hcst_period <- paste0(recipe$Analysis$Time$hcst_start, "-",
			recipe$Analysis$Time$hcst_end)
  init_month <- lubridate::month(as.numeric(substr(recipe$Analysis$Time$sdate,
					           start = 1, stop = 2)),
				 label = T, abb = T)
  # Define color palette and number of breaks according to output format
  ## TODO: Make separate function
  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")
  scores <- c("rps", "frps", "crps", "frps_specs")
  # Assign colorbar to each metric type
  ## TODO: Triangle ends
  for (name in c(skill_scores, scores, "mean_bias", "enssprerr")) {
    if (name %in% names(skill_metrics)) {
      # 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")) {
	display_name <- toupper(strsplit(name, "_")[[1]][1])
        skill <- skill_metrics[[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 <- skill_metrics[[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 <- skill_metrics[[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 <- skill_metrics[[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 <- skill_metrics[[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 == "mean_bias") {
	skill <- skill_metrics[[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)]
      }
      options(bitmapType = "cairo")
      # 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 <- skill_metrics[[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
      outfile <- paste0(outdir, name, ".png")
      toptitle <- paste(display_name, "-", data_cube$attrs$Variable$varName,
        		 "-", system_name, "-", init_month, hcst_period)
      months <- unique(lubridate::month(data_cube$attrs$Dates,
        			              label = T, abb = F))
      titles <- as.vector(months)
      # Plot
      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 = outfile,
                 bar_label_digits = 3,
                 bar_extra_margin = rep(0.9, 4),
                 bar_label_scale = 1.5,
                 axes_label_scale = 1.3)
      )
    }
  }
  info(recipe$Run$logger,
       "##### SKILL METRIC PLOTS SAVED TO OUTPUT DIRECTORY #####")
}

plot_ensemble_mean <- function(recipe, archive, fcst, outdir) {

  ## TODO: Add 'anomaly' to plot title
  # Abort if frequency is daily
  if (recipe$Analysis$Variables$freq == "daily_mean") {
    stop("Visualization functions not yet implemented for daily data.")
  }

  latitude <- fcst$coords$lat
  longitude <- fcst$coords$lon
  system_name <- archive$System[[recipe$Analysis$Datasets$System$name]]$name
  variable <- recipe$Analysis$Variables$name
  units <- attr(fcst$Variable, "variable")$units
  start_date <- paste0(recipe$Analysis$Time$fcst_year,
		       recipe$Analysis$Time$sdate)
  # Compute ensemble mean
  ensemble_mean <- s2dv::MeanDims(fcst$data, 'ensemble')
  # Drop extra dims, add time dim if missing:
  ensemble_mean <- drop(ensemble_mean)

  if (!("time" %in% names(dim(ensemble_mean)))) {
    dim(ensemble_mean) <- c("time" = 1, dim(ensemble_mean))
  }
  if (!'syear' %in% names(dim(ensemble_mean))) {
    ensemble_mean <- Reorder(ensemble_mean, c("time",
					      "longitude",
					      "latitude"))
  } else {
    ensemble_mean <- Reorder(ensemble_mean, c("syear",
					      "time",
					      "longitude",
					      "latitude"))
  }
  ## TODO: Redefine column colors, possibly depending on variable
  if (variable == 'prlr') {
    palette = "BrBG"
    rev = F
  } else {
    palette = "RdBu"
    rev = T
  }
  # Define brks, centered on in the case of anomalies
  ##
  if (grepl("anomaly",
	    fcst$attrs$Variable$metadata[[variable]]$long_name)) {
    variable <- paste(variable, "anomaly")
    max_value <- max(abs(ensemble_mean))
    ugly_intervals <- seq(-max_value, max_value, max_value/20)
    brks <- pretty(ugly_intervals, n = 12, min.n = 8)
  } else {
    brks <- pretty(range(ensemble_mean, na.rm = T), n = 15, min.n = 8)
  }
  cols <- grDevices::hcl.colors(length(brks) - 1, palette, rev = rev)
  options(bitmapType = "cairo")

  for (i_syear in start_date) {
    # Define name of output file and titles
    if (length(start_date) == 1) {
      i_ensemble_mean <- ensemble_mean
      outfile <- paste0(outdir, "forecast_ensemble_mean.png")
    } else {
      i_ensemble_mean <- ensemble_mean[which(start_date == i_syear), , , ]
      outfile <- paste0(outdir, "forecast_ensemble_mean_", i_syear, ".png")
    }
    toptitle <- paste("Forecast Ensemble Mean -", variable, "-", system_name,
  	  	      "- Initialization:", i_syear)
    months <- lubridate::month(fcst$attrs$Dates[1, 1, which(start_date == i_syear), ],
                               label = T, abb = F)
    titles <- as.vector(months)
    # Plots
    PlotLayout(PlotEquiMap, c('longitude', 'latitude'),
   	       i_ensemble_mean, longitude, latitude,
               filled.continents = F,
	       toptitle = toptitle,
	       title_scale = 0.6,
	       titles = titles,
	       units = units,
	       cols = cols,
	       brks = brks,
	       fileout = outfile,
	       bar_label_digits = 4,
	       bar_extra_margin = rep(0.7, 4),
	       bar_label_scale = 1.5,
	       axes_label_scale = 1.3)
  }
  info(recipe$Run$logger,
       "##### FCST ENSEMBLE MEAN PLOT SAVED TO OUTPUT DIRECTORY #####")
}

plot_most_likely_terciles <- function(recipe, archive,
				      fcst,
				      probabilities, 
				      outdir) {

  ## TODO: Add 'anomaly' to plot title
  # Abort if frequency is daily
  if (recipe$Analysis$Variables$freq == "daily_mean") {
    stop("Visualization functions not yet implemented for daily data.")
  }

  latitude <- fcst$coords$lat
  longitude <- fcst$coords$lon
  system_name <- archive$System[[recipe$Analysis$Datasets$System$name]]$name
  variable <- recipe$Analysis$Variables$name
  start_date <- paste0(recipe$Analysis$Time$fcst_year,
		       recipe$Analysis$Time$sdate)

  # Retrieve and rearrange probability bins for the forecast
  if (is.null(probabilities$probs_fcst$prob_b33) ||
      is.null(probabilities$probs_fcst$prob_33_to_66) ||
      is.null(probabilities$probs_fcst$prob_a66)) {
    stop("The forecast tercile probability bins are not present inside ",
	 "'probabilities', the most likely tercile map cannot be plotted.")
  }

  probs_fcst <- abind(probabilities$probs_fcst$prob_b33,
		      probabilities$probs_fcst$prob_33_to_66,
		      probabilities$probs_fcst$prob_a66,
		      along = 0)
  names(dim(probs_fcst)) <- c("bin",
			      names(dim(probabilities$probs_fcst$prob_b33)))

  ## TODO: Improve this section
  # Drop extra dims, add time dim if missing:
  probs_fcst <- drop(probs_fcst)
  if (!("time" %in% names(dim(probs_fcst)))) {
    dim(probs_fcst) <- c("time" = 1, dim(probs_fcst))
  }
  if (!'syear' %in% names(dim(probs_fcst))) {
    probs_fcst <- Reorder(probs_fcst, c("time", "bin", "longitude", "latitude"))
  } else {
    probs_fcst <- Reorder(probs_fcst,
			  c("syear", "time", "bin", "longitude", "latitude"))
  }

  for (i_syear in start_date) {
    # Define name of output file and titles
    if (length(start_date) == 1) {
      i_probs_fcst <- probs_fcst
      outfile <- paste0(outdir, "forecast_most_likely_tercile.png")
    } else {
      i_probs_fcst <- probs_fcst[which(start_date == i_syear), , , , ]
      outfile <- paste0(outdir, "forecast_most_likely_tercile_", i_syear, ".png")
    }
    toptitle <- paste("Most Likely Tercile -", variable, "-", system_name, "-",
 		      "Initialization:", i_syear)
    months <- lubridate::month(fcst$attrs$Dates[1, 1, which(start_date == i_syear), ],
	  	               label = T, abb = F)
    ## TODO: Ensure this works for daily and sub-daily cases
    titles <- as.vector(months)

    # Plots
    ## NOTE: PlotLayout() and PlotMostLikelyQuantileMap() are still being worked
    ## on. 
    suppressWarnings(
    PlotLayout(PlotMostLikelyQuantileMap, c('bin', 'longitude', 'latitude'),
  	       cat_dim = 'bin',
  	       i_probs_fcst, longitude, latitude,
	       coast_width = 1.5,
	       title_scale = 0.6,
	       legend_scale = 0.8, #cex_bar_titles = 0.6,
	       toptitle = toptitle,
	       titles = titles,
	       fileout = outfile,
	       bar_label_digits = 2,
	       bar_scale = rep(0.7, 4),
	       bar_label_scale = 1.2,
	       axes_label_scale = 1.3,
	       triangle_ends = c(F, F), width = 11, height = 8)
    )
  }

  info(recipe$Run$logger,
       "##### MOST LIKELY TERCILE PLOT SAVED TO OUTPUT DIRECTORY #####")
}
