source("tools/add_logo.R")
plot_ensemble_mean <- 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.")
  }
  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') {
    ensemble_mean <- s2dv::MeanDims(fcst$data, 'ensemble')
  } else if (method == 'iqr') {
    ensemble_mean <- Apply(fcst$data, target_dim = 'ensemble',
                          fun = function(x) {
                          IQR(x)})$output1
  } else {
    ensemble_mean <- 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]]
    var_long_name <- fcst$attrs$Variable$metadata[[variable]]$long_name
    units <- fcst$attrs$Variable$metadata[[variable]]$units
    # Subset ensemble mean by variable
    var_ens_mean <- ClimProjDiags::Subset(ensemble_mean,
                                          along = c("dat", "var",
                                                    "sday", "sweek"),
                                          indices = list(1, var, 1, 1),
                                          drop = 'selected')
    var_ens_mean <- Reorder(var_ens_mean, 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_ens_mean), 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_ens_mean))) {
        brks <- NULL
      } else {
        brks <- pretty(range(var_ens_mean, 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_ens_mean <- ClimProjDiags::Subset(var_ens_mean,
                                              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
      }
      # 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
      }

      # 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 = "PlotEquiMap",
                          plot_dims = c('longitude', 'latitude'),
                          var = i_var_ens_mean, lon = longitude,
                          lat = latitude, mask = mask, dots = 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_label_digits = 4, extra_margin = rep(1, 4),
                          bar_label_scale = 1.5, axes_label_scale = 1.1)
        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 <- PlotEquiMap
          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 = units)
        } 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)
        }
        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
          if (identical(fun, PlotRobinson)) {
            ## 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_ens_mean[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)))
	              # Convert plots to user-chosen format
          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 ENSEMBLE MEAN PLOTS SAVED TO OUTPUT DIRECTORY #####")
}

