source("modules/Visualization/R/tmp/PlotMostLikelyQuantileMap.R")
source("modules/Visualization/R/tmp/PlotCombinedMap.R")
source("modules/Visualization/R/tmp/GradientCatsColorBar.R")


## TODO: Change name
plot_most_likely_terciles <- function(recipe,
                                      fcst,
                                      probabilities,
                                      mask,
                                      dots, 
                                      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
  archive <- get_archive(recipe)
  system_name <- archive$System[[recipe$Analysis$Datasets$System$name]]$name
  start_date <- paste0(recipe$Analysis$Time$fcst_year,
                       recipe$Analysis$Time$sdate)
  init_month <- as.numeric(substr(recipe$Analysis$Time$sdate,
                                  start = 1, stop = 2))

  # 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:
  for (var in 1:fcst$dims[['var']]) {
    variable <- fcst$attrs$Variable$varName[[var]]
    var_long_name <- fcst$attrs$Variable$metadata[[variable]]$long_name
    # Choose colors depending on the variable
    if (variable %in% c('prlr')) { ## add others
      cols <- list(c("#FFC473", #FFAB38", 
                     "darkorange1"),
                   c("#b5b5b5", "black"),
                   c("#A0E5E4",#41CBC9", 
                     "deepskyblue3"))
    } else {
      cols <- list(c("#A0E5E4",#33BFD1", 
                     "deepskyblue3"),
                   c("#b5b5b5", "black"),
                   c("#FFB19A", #FF764D", 
                     "indianred3"))
    }
    var_probs <- ClimProjDiags::Subset(probs_fcst,
                                       along = c("var"),
                                       indices = var,
                                       drop = 'selected')
    var_probs <- Reorder(var_probs,
                         c("syear", "time", "bin", "longitude", "latitude"))
    for (i_syear in start_date) {
      # Define name of output file and titles
      i_var_probs <- var_probs[which(start_date == i_syear), , , , ]
      outfile <- paste0(outdir[[var]], "forecast_most_likely_tercile-", i_syear)
      # Mask
      if (!is.null(mask)) {
        outfile <- paste0(outfile, "_rpssmask")
        var_mask <- ClimProjDiags::Subset(mask,
                                          along = c("var"),
                                          indices = var,
                                          drop = 'selected')
        dim_mask <- dim(var_mask)
        var_mask <- as.numeric(var_mask <= 0)
        dim(var_mask) <- dim_mask 
      }
      # Dots
      if (!is.null(dots)) {
        outfile <- paste0(outfile, "_rpssdots")
        var_dots <- ClimProjDiags::Subset(dots,
                                          along = c("var"),
                                          indices = var,
                                          drop = 'selected')
        dim_dots <- dim(var_dots)
        var_dots <- as.numeric(var_dots <= 0)
        dim(var_dots) <- dim_dots
      }
      toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                         "\n", "Most Likely Tercile / Initialization: ",
                         i_syear)
      months <- lubridate::month(fcst$attrs$Dates[1, 1, which(start_date == i_syear), ],
                                 label = T, abb = F)
      years <- lubridate::year(fcst$attrs$Dates[1, 1, which(start_date == i_syear), ])
      if (recipe$Analysis$Workflow$Visualization$multi_panel) {
        ## 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_var_probs, longitude, latitude,
                   mask = mask,
                   dots = dots,
                   coast_width = 1.5,
                   title_scale = 0.6,
                   title_margin_scale = 0.7,
                   subtitle_scale = 1,
                   legend_scale = 0.8, #cex_bar_titles = 0.6,
                   toptitle = toptitle,
                   titles = titles,
                   fileout = paste0(outfile, ".png"),
                   bar_label_digits = 2,
                   bar_scale = rep(0.7, 4),
                   extra_margin = rep(1, 4),
                   bar_label_scale = 1.2,
                   axes_label_scale = 1.1,
                   triangle_ends = c(F, F)) # , width = 11, height = 8)
        )
      } else {
        for (i in 1:length(months)) {
          # Get forecast time label
          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)
          # Get mask subset
          if (!is.null(mask)) {
            mask_i <- Subset(var_mask, along = 'time', indices = i, drop = TRUE)
          } else {
            mask_i <- NULL
          }
          # Get dots subset
          if (!is.null(dots)) {
            dots_i <- Subset(var_dots, along = 'time', indices = i, drop = TRUE)
          } else {
            dots_i <- NULL
          }
          # Define plot title
          toptitle <- paste0(system_name, " / ",
                             str_to_title(var_long_name),
                             "\n", "Most Likely Tercile / ",
                             months[i], " ", years[i],
                             " / Start date: ", 
                             format(as.Date(i_syear, format="%Y%m%d"), 
                                            "%d-%m-%Y"))
          # Plot
          fileout <- paste0(outfile, "_ft", forecast_time, ".png")
          PlotMostLikelyQuantileMap(cat_dim = 'bin',
                                    probs = i_var_probs[i, , , ],
                                    lon = longitude, lat = latitude,
                                    coast_width = 1.5,
                                    mask = mask_i,
                                    dots = dots_i,
                                    col_mask = 'antiquewhite',
                                    cols = cols,
                                    title_scale = 1,
                                    legend_scale = 0.8,
                                    cex_bar_titles = 0.9,
                                    toptitle = toptitle,
                                    fileout = fileout,
                                    bar_label_digits = 2,
                                    bar_label_scale = 0.7,
                                    axes_label_scale = 1.1,
                                    triangle_ends = c(F, F) , width = 10, height = 8)
        }
      }
    }
  }
  info(recipe$Run$logger,
       "##### MOST LIKELY TERCILE PLOTS SAVED TO OUTPUT DIRECTORY #####")
}
