## Functions required for normal cat and triangles end.
## See https://earth.bsc.es/gitlab/external/cstools/-/issues/125
source("modules/Visualization/R/tmp/PlotMostLikelyQuantileMap.R")
source("modules/Visualization/R/tmp/PlotCombinedMap.R")
source("modules/Visualization/R/tmp/ColorBar.R")
source("modules/Visualization/R/tmp/clim.palette.R")
source("modules/Visualization/R/tmp/Utils.R")
source("modules/Visualization/R/tmp/PlotEquiMap.R")
source("modules/Visualization/R/tmp/ColorBar_onebox.R")
source("modules/Visualization/R/tmp/GradientCatsColorBar.R")

## TODO: Change name
plot_most_likely_terciles <- function(recipe,
                                      fcst,
                                      probabilities,
                                      mask,
                                      dots, 
                                      outdir,
                                      output_conf) {

  ## TODO: Add 'anomaly' to plot title
  # Abort if frequency is daily
  if (recipe$Analysis$Variables$freq %in% c("daily", "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)
  if (tolower(recipe$Analysis$Horizon) == "seasonal") {
    init_month <- as.numeric(substr(recipe$Analysis$Time$sdate,
                                    start = 1, stop = 2))
  } else {
    ## TODO: Sort out decadal initial month (is it always January?)
    init_month <- 1
  }

  # 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"), 
                   c("grey"),
                   c("#A0E5E4","#41CBC9")) 
      col_sup <- list("darkorange1", "grey", "deepskyblue3")
    } else {
      cols <- list(c("#A0E5E4","#33BFD1"),
                   c("grey"),
                   c("#FFB19A", "#FF764D")) 
      col_sup <- list("deepskyblue3", "grey", "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 <- ClimProjDiags::Subset(var_probs,
                                           along = c("syear"),
                                           indices = which(start_date == i_syear),
                                           drop = 'selected')
      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 {
        output_configuration <- output_conf$PlotEquiMap$most_likely_terciles
        base_args <- list(cat_dim = 'bin',
                          probs = NULL,
                          lon = longitude, lat = latitude,
                          coast_width = 1.5,
                          mask = NULL,
                          dots = NULL,
                          dot_symbol = 4,
                          dot_size = 1,
                          col_mask = 'antiquewhite',
                          cols = cols,
                          col_sup = col_sup,
                          title_scale = 1,
                          legend_scale = 0.8,
                          cex_bar_titles = 0.9,
                          bar_label_digits = 2,
                          bar_label_scale = 0.7,
                          bar_limits = list(c(40, 85), c(40, 85), c(40, 85)),
                          brks = list(4,2,4),
                          axes_label_scale = 1.1,
                          plot_margin = c(5.1, 4.1, 4.1, 2.1),
                          return_leg = T,
                          triangle_ends = c(F, T) , width = 10, height = 8)
        base_args[names(output_configuration)] <- output_configuration
        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")
          base_args$probs <- i_var_probs[i, , , ]
          base_args$mask <- mask_i
          base_args$dots <- dots_i
          cb_info <- do.call(PlotMostLikelyQuantileMap,
                  args = c(base_args,
                           list(toptitle = toptitle,
                                fileout = fileout)))
          # Add color bars with 1 range for normal category:
          for (i_bar in 1:cb_info$nmap) {
            tmp <- cb_info
            tmp$brks <- tmp$brks[[i_bar]]
            tmp$cols <- tmp$cols[[i_bar]]
            tmp$bar_limits <- tmp$bar_limits[[i_bar]]
            tmp$col_sup <- tmp$col_sup[[i_bar]]
            tmp$title <- tmp$bar_titles[i_bar]
            tmp$bar_titles <- NULL
            tmp$nmap <- NULL
            tmp$var_limits <- NULL
            if (length(cb_info$brks[[i_bar]]) > 2) {
            # plot colorbar as normal
              do.call(ColorBar, tmp)
            } else {
            # plot a square
              tmp$brks <- 4 
              tmp$draw_ticks <- F
              tmp$box_label <- "> 40"
              tmp$triangle_ends <- c(F, F) 
              tmp$draw_separators <- FALSE
              do.call(ColorBar_onebox, tmp)
            }
          }
          dev.off()
        }
      }
    }
  }
  info(recipe$Run$logger,
       "##### MOST LIKELY TERCILE PLOTS SAVED TO OUTPUT DIRECTORY #####")
}
