# 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'
# mask: 
# dots: 
# outdir: output directory
# output_conf: list as returned by read_yaml("output_conf.yml")
# method: Computation to be applied to the forecast. Options:
#   'median' (default), 'iqr', 'mean'.
plot_forecast_map <- function(recipe, fcst, mask = NULL, dots = NULL, 
                              outdir, output_conf,
                              method = 'median',
                              logo = NULL) {
  ## 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.")
  }
  # Read 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
  }
  # Coordinates and names
  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 
  }

  if (!is.null(recipe$Analysis$Workflow$Visualization$projection)) {
    projection <- tolower(recipe$Analysis$Workflow$Visualization$projection)
  } else {
    projection <- "cylindrical_equidistant"
  }
  # Compute ensemble mean or other
  if (!is.null(method)) {
    method <- tolower(method)
  }
  if (method == 'mean') {
    forecast_product <- s2dv::MeanDims(fcst$data, 'ensemble')
  } else if (method == 'iqr') {
    ## Change to uppercase to display correctly in plot title?
    method <- "IQR"
    forecast_product <- Apply(fcst$data, target_dim = 'ensemble',
                              fun = function(x) {
                              IQR(x)})$output1
  } else {
    ## return median by default?
    forecast_product <- Apply(fcst$data, target_dim = 'ensemble',
                              fun = function(x) {
                              median(x)})$output1
  }
  # Loop over variable dimension
  for (var in 1:fcst$dims[['var']]) {
    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)
    units <- fcst$attrs$Variable$metadata[[variable]]$units
    # Subset ensemble mean by variable
    var_fcst <- ClimProjDiags::Subset(forecast_product,
                                      along = c("dat", "var",
                                                "sday", "sweek"),
                                      indices = list(1, var, 1, 1),
                                      drop = 'selected')
    var_fcst <- Reorder(var_fcst, c("syear",
                                    "time",
                                    "longitude",
                                    "latitude"))
    ## TODO: Redefine column colors, possibly depending on variable
    if (variable == 'prlr') {
      palette = "BrBG"
      rev = F
    } else {
      palette = "RdBu"
      rev = T
    }
    # Define brks, centered around zero in the case of anomalies
    if (grepl("anomaly", var_long_name)) {
      variable <- paste(variable, "anomaly")
      max_value <- max(abs(var_fcst), na.rm = TRUE)
      ugly_intervals <- seq(-max_value, max_value, max_value/20)
      brks <- pretty(ugly_intervals, n = 12, min.n = 8)
      cols <- grDevices::hcl.colors(length(brks) - 1, palette, rev = rev)
    } else {
      if (variable == 'prlr') {
        cols <- c("#FFAB38", "white", "#41CBC9")
        col_fun <- colorRampPalette(cols)
      } else {
        cols <- c("#33BFD1", "white", "#FF764D")
        col_fun <- colorRampPalette(cols)
      }
      if (all(is.na(var_fcst))) {
        brks <- NULL
      } else {
        max_value <- max(abs(var_fcst), na.rm = T)
        brks <- pretty(c(-1 * max_value, max_value), n = 15, min.n = 8)
        ## old:
        ## brks <- pretty(range(var_fcst, na.rm = T), n = 15, min.n = 8)
      }
      if (is.null(brks)) {
        cols <- NULL
      } else {
        cols <- col_fun(length(brks) - 1)
      }
    }
    
    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_ensemble_", method,
                        "-", i_syear)
      # Mask
      if (!is.null(mask)) {
        outfile <- paste0(outfile, "_enscormask")
        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, "_enscordots")
        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) { 
        # 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", "Forecast Ensemble", method, " / ", "Issued on ",
                             format(ymd(start_date), "%d-%m-%Y"))
        } else {
          toptitle <- paste0(system_name, " / ", str_to_title(var_long_name),
                             "\n", "Forecast Ensemble ", method, "/ ", "Init.: ", i_syear)
        }
        # Plots
        output_configuration <- output_conf$multipanel$forecast_ensemble_mean
        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,
                          units = units, cols = cols, brks = brks,
                          fileout = paste0(outfile, ".pdf"),
                          bar_extra_margin = rep(0.9, 4),
                          bar_label_digits = 4, extra_margin = rep(1, 4),
                          bar_label_scale = 1.5, axes_label_scale = 1.1,
                          shapefile = shapefile)
        base_args[names(output_configuration)] <- output_configuration
        do.call(PlotLayout, base_args)
      } 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.7,
                            font.main = 2, margin_scale = c(3, 5, 5, 5),
                            bar_extra_margin = rep(0.9, 4),
                            filled.continents = F, brks = brks, cols = cols,
                            bar_label_digits = 4, bar_label_scale = 1.5,
                            axes_label_scale = 1, units = units,
                            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, drawleg = 'bar',
                            style = 'point', brks = brks, cols = cols,
                            dots_size = 0.2, dots_shape = 47)
        }
        output_configuration <- output_conf[[projection]]$forecast_ensemble_mean
        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", "Ensemble ", method, " / ",
                               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", "Ensemble ", method, " / ",
                               "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", "Ensemble ", method, " / ",
                               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, "\n",
                   "Units: ", units)
          # 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)))
          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)
          }
        }        
      }
    }
  }
  info(recipe$Run$logger,
       "##### FORECAST MAPS SAVED TO OUTPUT DIRECTORY #####")
}

