## 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)
  if (recipe$Analysis$Datasets$System$name == 'Multimodel'){
    system_name <- paste0('Multimodel-',
                          recipe$Analysis$Datasets$Multimodel$approach)
  } else {
    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 if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
    init_week <- 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
    init_week <- 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)
      if (is.null(attributes(fcst$attrs$time_bounds))) {
        months <- lubridate::month(
                    fcst$attrs$Dates[1, 1, which(start_date == i_syear), ],
                    label = T, abb = F)
      } else {
        if (length(attributes(fcst$attrs$time_bounds)$plotting_attr) > 1) {
          months <- unlist(
            lapply(1:length(fcst$attrs$time_bounds$start), function(i) {
             ftime_ini <- attributes(fcst$attrs$time_bounds)$plotting_attr$ini[i]
             ftime_end <- attributes(fcst$attrs$time_bounds)$plotting_attr$end[i]
             # labels for file name:
             ftime <- paste0(ftime_ini, "-", ftime_end)
             # title names:
             ftime_ini <- init_month + ftime_ini - 1
             ftime_ini <- ifelse(ftime_ini > 12, ftime_ini - 12, ftime_ini)
             ftime_ini <- month.name[ftime_ini]
             ftime_end <- init_month + ftime_end - 1
             ftime_end <- ifelse(ftime_end > 12, ftime_end - 12, ftime_end)
             ftime_end <- month.name[ftime_end]
             toptitle <- paste(ftime_ini, "to", ftime_end)}))
        } else {
          months <- attributes(fcst$attrs$time_bounds)$plotting_attr[[1]]
        }
      }

      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
        if (tolower(recipe$Analysis$Horizon) == "seasonal") {
          titles <- as.vector(months)
        } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
          titles <- as.vector(weeks)
        } else {
          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, ".pdf"),
                   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
      }
      if (tolower(recipe$Analysis$Horizon) == "seasonal") {
        for (i in 1:length(months)) {
          # Get forecast time label
          if (is.null(attributes(fcst$attrs$time_bounds))) {
            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)
          } else {
            forecast_time <- months
          }

          # 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, ".pdf")
          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)
            }
          }
        }
      } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") { 
        for (i in 1:length(weeks)) {
          # Get forecast time label
          forecast_time <- match(weeks[i], week.name) - init_week + 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
          if (tolower(recipe$Analysis$Horizon) == 'seasonal') {
            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"))
          } else {
            toptitle <- paste0(system_name, " / ",
                               str_to_title(var_long_name),
                               "\n", "Most Likely Tercile / ",
                               months[i],  
                               " / Start date: ",
                               i_syear)  
          }        
          # Plot
          if (is.null(attributes(fcst$attrs$time_bounds))) {
            fileout <- paste0(outfile, "_ft", forecast_time, ".pdf")
          } else {
            if (length(attributes(fcst$attrs$time_bounds)$plotting_attr) > 1) {
              fileout <- paste0(outfile, "_ft",
              attributes(fcst$attrs$time_bounds)$plotting_attr$ini[i], "-",
              attributes(fcst$attrs$time_bounds)$plotting_attr$end[i], ".pdf")
            } else {
              fileout <- paste0(outfile, "_ft", months[i], ".pdf")
            }
          }

          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 #####")
}
