## 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 = NULL,
                                      dots = NULL, 
                                      outdir,
                                      output_conf,
                                      logo = NULL,
                                      return_elements = FALSE) {
  
  ## 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.")
  }
 
  if (!is.null(recipe$Analysis$Workflow$Visualization$shapefile)) {
    shapefile <- read_sf(recipe$Analysis$Workflow$Visualization$shapefile)
    shapefile <- spTransform(as(shapefile, "Spatial"),
                             CRSobj = CRS("+proj=longlat"))
    shapefile <- SpatialPolygons2map(shapefile)
  } else {
    shapefile <- NULL
  }

  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
  }
  # Initialization dates
    if (tolower(recipe$Analysis$Horizon) == "seasonal") {
    start_date <- paste0(recipe$Analysis$Time$fcst_year,
                         recipe$Analysis$Time$sdate)
    init_date <- as.numeric(substr(recipe$Analysis$Time$sdate,
                                    start = 1, stop = 2))
  } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
    start_date <- recipe$Analysis$Time$sdate
    init_date <- start_date ## start_date and init_date are same for subseasonal
  } else {
    ## TODO: Sort out decadal initial month (is it always January?)
    start_date <- paste0(recipe$Analysis$Time$fcst_year,
                         recipe$Analysis$Time$sdate)
    init_date <- 1
  }
  # Retrieve and rearrange probability bins for the forecast
  if ("bin" %in% names(dim(probabilities[[1]]))) {
    ## TODO: Adjust array dimensions if needed. What if there is more than one 
    ## set of categories?
    terciles_index <- which(lapply(recipe$Analysis$Workflow$Probabilities$percentiles,
                                   function(x) {
                                     all(match(x, c("1/3", "2/3")))
                                   }) == TRUE)
    probabilities <- probabilities[[terciles_index]]
    probs_fcst <- .drop_dims(probabilities)
  } else {
    if (is.null(probabilities$prob_b33) ||
        is.null(probabilities$prob_33_to_66) ||
        is.null(probabilities$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$prob_b33,
                        probabilities$prob_33_to_66,
                        probabilities$prob_a66,
                        along = 0)
    names(dim(probs_fcst)) <- c("bin",
                                names(dim(probabilities$prob_b33)))
  }

  for (var in 1:fcst$dims[['var']]) {
    ## NOTE: Variables starting with var_* represent the variable counter.
    variable <- fcst$attrs$Variable$varName[[var]]
    custom_dictionary <- recipe$Analysis$Visualization$custom_dictionary
    var_long_name <- .get_variable_longname(data_cube = fcst,
                                            index = var,
                                            format = custom_dictionary)
    # 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) {
      ## NOTE: Variables the start with i_* represent the forecast year counter.
      # 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 <- file.path(outdir[[var]],
                           paste0("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 
      } else {
        var_mask <- NULL
      }
      # 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
      } else {
        var_dots <- NULL
      }

      # Plot title labels look different depending on horizon and aggregation
      labels <- .get_plot_time_labels(dates = fcst$attrs$Dates,
                                      horizon = recipe$Analysis$Horizon,
                                      init_date = init_date,
                                      start_date = start_date,
                                      i_syear = i_syear,
                                      time_bounds = fcst$attrs$time_bounds)
      time_labels <- labels$time_labels
      years <- labels$years
      
      if (recipe$Analysis$Workflow$Visualization$multi_panel) {
        ## TODO: Ensure this works for daily and sub-daily cases
        titles <- as.vector(time_labels)
        # Define top title: system name, variable, forecast date
        if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
          toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                             "\n", "Most Likely Tercile / ", "Issued on ",
                             format(ymd(start_date), "%d-%m-%Y"))
        } else {
          toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                             "\n", "Most Likely Tercile / ", "Initialization: ",
                             i_syear)
        }
        # Plots
        ## NOTE: PlotLayout() and PlotMostLikelyQuantileMap() are still being worked
        ## on. This option does not work with mask or dots for now. 
        suppressWarnings(
        PlotLayout(PlotMostLikelyQuantileMap, c('bin', 'longitude', 'latitude'),
                   cat_dim = 'bin',
                   i_var_probs, longitude, latitude,
                   mask = var_mask,
                   dots = var_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),
                   shapefile = shapefile) # , width = 11, height = 8)
        )
      } else {
        output_configuration <- output_conf$cylindrical_equidistant$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 = 0.9,
                          legend_scale = 0.7,
                          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(1, 5, 5, 5),
                          # plot_margin = c(5.1, 4.1, 4.1, 2.1),
                          return_leg = T,
                          triangle_ends = c(F, T), width = 10, height = 8,
                          shapefile = shapefile)
        base_args[names(output_configuration)] <- output_configuration
        for (i in 1:length(time_labels)) {
          ## variables ending in *_i represent each forecast time
          # Get forecast time label
          if (tolower(recipe$Analysis$Horizon) %in% c("seasonal", "decadal")) {
            if (is.null(attributes(fcst$attrs$time_bounds))) {
              forecast_time <- match(time_labels[i], month.name) - init_date + 1
              # Months should range from 1 to 12
              if (forecast_time < 1) {
                forecast_time <- forecast_time + 12
              }
              forecast_time <- sprintf("%02d", forecast_time)
            } else {
              ## TODO: Check this
              # forecast_time <- sprintf("%02d", time_labels[i])
              forecast_time <- time_labels[i]
            }
          } else if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
            forecast_time <- sprintf("%02d", i)
          }
          # 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 / ",
                               time_labels[i], " ", years[i],
                               " / Start date: ",
                               format(as.Date(i_syear, format = "%Y%m%d"),
                                      "%d-%m-%Y"))
          } else if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
            toptitle <- paste0(system_name, " / ",
                               str_to_title(var_long_name),
                               "\n", "Most Likely Tercile / ",
                               "Issued on ",
                               format(ymd(start_date), "%d-%m-%Y"),
                               "\n", time_labels[i], years[i])
          } else {
            toptitle <- paste0(system_name, " / ",
                               str_to_title(var_long_name),
                               "\n", "Most Likely Tercile / ",
                               time_labels[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", time_labels[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()
          if (!is.null(recipe$Analysis$Workflow$Visualization$file_format) &&
              tolower(recipe$Analysis$Workflow$Visualization$file_format) != "pdf") {
            extension <- tolower(recipe$Analysis$Workflow$Visualization$file_format)
            system_command <- paste("convert -density 300", fileout,
                                    "-resize 40% -alpha remove",
                                     paste0(tools::file_path_sans_ext(fileout), ".",
                                            extension))
            system(system_command)
            unlink(fileout)
            fileout <- paste0(tools::file_path_sans_ext(fileout), ".",
                              extension)
          }
          if (!is.null(logo)) {
            add_logo(file = fileout, logo = logo)
          }
        }
      }
    }
  }
  if (return_elements) {
    elements <- list(
      latitude = latitude,
      longitude = longitude,
      system_name = system_name,
      start_date = start_date,
      init_date = init_date,
      base_args_brks = if (exists("base_args")) base_args$brks else NULL,
      cols = if (exists("cols")) cols else NULL,
      toptitle = if (exists("toptitle")) toptitle else NULL,
      labels = if (exists("labels")) labels else NULL,
      time_labels = if (exists("time_labels")) time_labels else NULL,
      forecast_time = if (exists("forecast_time")) forecast_time else NULL,
      years = if (exists("years")) years else NULL,
      fileout = if (exists("fileout")) fileout else NULL
    )
    return(elements)
  }
  info(recipe$Run$logger,
       "##### MOST LIKELY TERCILE PLOTS SAVED TO OUTPUT DIRECTORY #####")
}
