# recipe: the SUNSET recipe
# fcst: an s2dv_cube object containing the forecast to be plotted. The forecast
#   should have the following dimensions: 
#   'dat', 'var', 'syear', 'time', 'latitude', 'longitude', 'ensemble'
# probabilities: an array containing the probabilities to be plotted. Dimension
#   'bin' is 2 and contains the complementary probabilities. 
# mask: 
# dots: 
# outdir: output directory
# output_conf: list as returned by read_yaml("output_conf.yml")
plot_extreme_probs <- function(recipe, fcst,
                               probabilities,
                               mask = NULL, dots = NULL,
                               outdir, output_conf,
                               brks_min = 5, brks_max = 95,
                               logo = NULL) {
  # Abort if frequency is daily
  if (recipe$Analysis$Variables$freq %in% c("daily", "daily_mean")) {
    stop("Visualization functions not yet implemented for daily data.")
  }
  # Check if arrays contain 'bin' dimension
  if (is.list(probabilities) && !("bin" %in% names(dim(probabilities[[1]])))) {
      stop("The forecast probability categories are stored in lists, while this",
           " function expects arrays with a 'bin' dimension. This function is ",
           "only available for full cross-validation workflows and has not been ",
           " implemented for this case.")
  }

  # Metadata
  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
  }
  # Shapefile 
  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 
  }
  # 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 
  }

  if (!is.null(recipe$Analysis$Workflow$Visualization$projection)) {
    projection <- tolower(recipe$Analysis$Workflow$Visualization$projection)
  } else {
    projection <- "cylindrical_equidistant"
  }
  ## TODO: improve handling of mask
  mask_list <- mask
  # Define indices and filter binary probabilities
  probs_index <- which(vapply(recipe$Analysis$Workflow$Probabilities$percentiles,
                              length, integer(1)) == 1)
  probs_labels <- recipe$Analysis$Workflow$Probabilities$percentiles[probs_index]
  # Remove unneeded dimensions and multiply values by 100 for display
  probabilities <- lapply(probabilities,
                          function(x) {
                            x <- .drop_dims(x)
                            x <- x*100
                          })
  output_files <- list()
  # Loop over variable dimension
  for (var in 1:fcst$dims[['var']]) {
    # Loop over requested binary probabilities
    for (p in 1:length(probs_index)) {
      if (is.list(probabilities)) {
        probs_array <- probabilities[[probs_index[p]]]
      } else {
        probs_array <- probabilities
      }
      # Loop over bin dimension
      for (bin in 1:dim(probs_array)["bin"]) {
        variable <- fcst$attrs$Variable$varName[[var]]
        var_long_name <- .get_variable_longname(data_cube = fcst,
                                                index = var,
                                                format = NULL) ## TODO: change
        # Subset ensemble mean by variable
        var_fcst <- ClimProjDiags::Subset(probs_array,
                                          along = c("bin", "var"),
                                          indices = list(bin = bin, var = var),
                                          drop = 'selected')
        var_fcst <- Reorder(var_fcst, c("syear",
                                        "time",
                                        "longitude",
                                        "latitude"))
        ## TODO: take threshold from Probabilities()
        base_prob <- strsplit(probs_labels[[p]], "/")[[1]]
        base_prob <- as.numeric(base_prob[1]) / 
          as.numeric(base_prob[2])
        if (bin == 1) {
          threshold <- base_prob * 100
          if (variable == 'prlr') {
            # green to yellow
            cols_1 <- c("#41CBC9", "white")
            cols_2 <- c("white", "#FFAB38")
          } else {
            # red to blue
            cols_1 <- c("#FF764D", "white")
            cols_2 <- c("white", "#33BFD1")
          }
        } else {
          threshold <- (1 - base_prob) * 100
          if (variable == 'prlr') {
            # yellow to green
            cols_1 <- c("#FFAB38", "white")
            cols_2 <- c("white", "#41CBC9")
          } else {
            # blue to red
            cols_1 <- c("#33BFD1", "white")
            cols_2 <- c("white", "#FF764D")
          }
        }
        col_fun_1 <- colorRampPalette(cols_1)
        col_fun_2 <- colorRampPalette(cols_2)
        if (all(is.na(var_fcst))) {
          brks <- NULL
        } else {
          brks <- pretty(c(brks_min, brks_max), n = 5, min.n = 5)
          # Add brks_min, brks_max and threshold to brks
          brks <- unique(round(c(brks, brks_min, brks_max, threshold), digits = 2))
          # Add extra break if threshold is below lower break
          if (threshold <= brks_min) {
            brks_min_threshold <- threshold/2
            brks <- c(brks, brks_min_threshold)
          } else {
            brks_min_threshold <- brks_min
          }
          # Add extra break if threshold is above upper break
          if (threshold >= brks_max) {
            brks_max_threshold <- 100 - (100-threshold)/2
            brks <- c(brks, brks_max_threshold)
          } else {
            brks_max_threshold <- brks_max
          }
          # Keep only values within brks_min_threshold and brks_max_threshold
          brks <- sort(unique(brks[brks >= brks_min_threshold & brks <= brks_max_threshold]))
        }
        if (is.null(brks)) {
          cols <- NULL
        } else {
          cols_1 <- col_fun_1(length(brks[brks < round(threshold, 10)]) + 2)
          cols_2 <- col_fun_2(length(brks[brks > round(threshold, 10)]) + 2)
          # cols_1[length(cols_1)] and cols_2[1] are white
          # cols_1[1] and cols_2[length(cols_2)] are for col_inf and col_sup
          cols <- c(cols_1[-c(1, length(cols_1))], cols_2[-c(1, length(cols_2))])
        }
       
        # Loop over forecast years
        for (i_syear in start_date) {
          i_var_fcst <- ClimProjDiags::Subset(var_fcst,
                                              along = c("syear"),
                                              indices = which(start_date == i_syear),
                                              drop = 'selected')
          outfile <- paste0(outdir[[var]], "forecast_probability_",
                            ifelse(bin == 1, "below_p", "above_p"),
                            sprintf("%02.0f", round(base_prob * 100, 2)),
                            "-", i_syear)
          # Mask
          if (!is.null(mask)) {
            if (is.list(mask)) {
              mask_p <- mask[[paste0("rpss-", names(probs_index)[p])]]
            } else {
              mask_p <- mask
            }
            ## TODO: Change name?
            outfile <- paste0(outfile, "_rpssmask")
            var_mask <- ClimProjDiags::Subset(mask_p,
                                              along = c("var"),
                                              indices = list(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)) {
            if (is.list(dots)) {
              dots_p <- dots[[paste0("rpss-", names(probs_index)[p])]]
            } else {
              dots_p <- dots
            }
            outfile <- paste0(outfile, "_rpssdots")
            var_dots <- ClimProjDiags::Subset(dots_p,
                                              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) {
            # Define name of output file and titles
            if (tolower(recipe$Analysis$Horizon) %in% c("seasonal", "subseasonal")) {
              titles <- as.vector(time_labels)
            } else {
              titles <- NULL
            }
            if (tolower(recipe$Analysis$Horizon) == "subseasonal") {
              toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                                 "\n", ifelse(bin == 1, "Below ", "Above "),
                                 round(base_prob*100, 2), "%",
                                 " / ", "Issued on ",
                                 format(ymd(start_date), "%d-%m-%Y"))
            } else {
              toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                                 "\n", ifelse(bin == 1, "Below ", "Above "),
                                 round(base_prob*100, 2), "%",
                                 " / ", "Init.: ", i_syear)
            }
            # Plots
            ## TODO: Should not be forecast ensemble mean...
            output_configuration <- output_conf$multipanel$extreme_probabilities
            base_args <- list(fun = "VizEquiMap",
                              plot_dims = c('longitude', 'latitude'),
                              var = i_var_fcst, lon = longitude,
                              lat = latitude, mask = var_mask, dots = var_dots,
                              filled.continents = FALSE, toptitle = toptitle,
                              title_scale = 0.7, subtitle_scale = 0.8,
                              subtitle_margin_scale = 2, titles = titles,
                              cols = cols, brks = brks,
                              fileout = paste0(outfile, ".pdf"),
                              bar_label_digits = 4, extra_margin = rep(1, 4),
                              bar_label_scale = 1.5, axes_label_scale = 1.1,
                              triangle_ends = c(FALSE, FALSE),
                              shapefile = shapefile) 
            base_args[names(output_configuration)] <- output_configuration
            do.call(PlotLayout, base_args)
            # File conversion
            if (!is.null(recipe$Analysis$Workflow$Visualization$file_format) &&
                tolower(recipe$Analysis$Workflow$Visualization$file_format) != "pdf") {
              fileout <- paste0(outfile, ".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)
            }
            # Add logo
            if (!is.null(logo)) {
              add_logo(file = fileout, logo = logo)
            }
          } else {
            # Define function and parameters depending on projection
            if (projection == 'cylindrical_equidistant') {
              fun <- VizEquiMap
              base_args <- list(var = NULL, dots = NULL, mask = NULL,
                                lon = longitude, lat = latitude,
                                dot_symbol = 20, title_scale = 0.6,
                                font.main = 2, margin_scale = c(1, 5, 5, 5),
                                filled.continents = F, brks = brks, cols = cols,
                                bar_label_digits = 4, bar_label_scale = 1.5,
                                axes_label_scale = 1, units = "%",
                                triangle_ends = c(TRUE, TRUE),
                                col_inf = cols_1[1],
                                col_sup = cols_2[length(cols_2)],
                                shapefile = shapefile)
            } else {
              fun <- PlotRobinson
              common_projections <- c("robinson", "stereographic", "lambert_europe")
              if (projection %in% common_projections) {
                target_proj <- get_proj_code(projection)
              } else {
                target_proj <- projection
              }
              base_args <- list(data = NULL, mask = NULL, dots = NULL,
                                lon = longitude, lat = latitude,
                                lon_dim = 'longitude', lat_dim = 'latitude',
                                target_proj = target_proj, legend = 's2dv',
                                style = 'point', brks = brks, cols = cols,
                                dots_size = 0.2, dots_shape = 47,
                                col_inf = cols_1[1],
                                col_sup = cols_2[length(cols_2)])
            }
            output_configuration <- output_conf[[projection]]$extreme_probabilities
            base_args[names(output_configuration)] <- output_configuration
            # Loop over forecast times
            for (i in 1:length(time_labels)) {  
              # 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", ifelse(bin == 1, "Below ", "Above "),
                                   round(base_prob*100, 2), "%", " / ", time_labels[i],
                                   " ", years[i], " / Start date: ",
                                   format(as.Date(i_syear, format = "%Y%m%d"),
                                          "%d-%m-%Y"))
                forecast_time_caption <- paste0("Forecast month: ", sprintf("%02d", i))
              } else if (tolower(recipe$Analysis$Horizon) == 'subseasonal') {
                toptitle <- paste0(system_name, " / ",
                                   str_to_title(var_long_name),
                                   "\n", ifelse(bin == 1, "Below ", "Above "),
                                   round(base_prob*100, 2), "%", " / ", "Issued on ", 
                                   format(ymd(start_date), "%d-%m-%Y"),
                                   "\n", time_labels[i], years[i])
                forecast_time_caption <- paste0("Forecast week: ", sprintf("%02d", i))
              } else {
                toptitle <- paste0(system_name, " / ",
                                   str_to_title(var_long_name),
                                   "\n", ifelse(bin == 1, "Below ", "Above "),
                                   round(base_prob*100, 2), "%", " / ", time_labels[i],
                                   " ", years[i], " / Start date: ", i_syear)
                forecast_time_caption <- paste0("Forecast month: ", sprintf("%02d", i))
              }
              # Define caption
              ## TODO: Customize technical details
              base_args[['caption']] <-
                paste0("Nominal start date: ", format(ymd(start_date), 
                                                      "%d-%m-%Y"), "\n",
                       forecast_time_caption, "\n",
                       "Reference: ", recipe$Analysis$Datasets$Reference,
                       " [", recipe$Analysis$Time$hcst_start, " - ",
                       recipe$Analysis$Time$hcst_end, "]",
                       "\n", " ")
              # Modify base arguments
              base_args[[1]] <- i_var_fcst[i, , ]
              if (is.null(attributes(fcst$attrs$time_bounds))) {
                fileout <- paste0(outfile, "_ft", sprintf("%02d", i), ".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$mask <- mask_i
              base_args$dots <- dots_i
              # Plot
              do.call(fun,
                      args = c(base_args,
                               list(toptitle = toptitle,
                                    fileout = fileout)))
              output_files <- append(output_files, fileout)
              # File conversion
              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)
              }
              # Add logo
              if (!is.null(logo)) {
                add_logo(file = fileout, logo = logo)
              }
            }
          }
        }
      }
    }
  }
  info(recipe$Run$logger,
       "##### EXTREME PROBABILITIES PLOTS SAVED TO OUTPUT DIRECTORY #####")
}
