# Function to write autosubmit configuration from an SUNSET recipe. The function
# reads the corresponding AS configuration file templates and fills them with
# the information needed to run the experiment. The modified configuration
# files are saved in the `conf/` folder of the Autosubmit experimet.
# 
# recipe: the SUNSET recipe
# nchunks: the number of 'chunks' to be processed by Autosubmit, as returned
#   by divide_recipe().
# chunk_to_recipe: list with the correspondence between the chunk number and 
#   the name of the atomic recipe, as returned by divide_recipe().
# split_to_recipe: list with the correspondence between the split number and
#   the name of the multi-model atomic recipe, as returned by divide_recipe().

write_autosubmit_conf <- function(recipe, nchunks,
                                  chunk_to_recipe,
                                  split_to_recipe) {
  # Experiment ID
  expid <- recipe$Run$auto_conf$expid
  # Directory with the experiment templates
  template_dir <- paste0("autosubmit/conf_", recipe$Run$filesystem, "/")
  # Read autosubmit info for the specific filesystem (e.g. esarchive)
  auto_specs <- read_yaml("conf/autosubmit.yml")[[recipe$Run$filesystem]]
  # Output directory
  dest_dir <- file.path(auto_specs$experiment_dir, expid, "conf")
  proj_dir <- file.path(auto_specs$experiment_dir, expid, "proj", "auto-s2s")
  tmp_dir <- auto_specs$tmp_dir
  # HPC details
  platform <- recipe$Run$auto_conf$platform
  hpc_user <- recipe$Run$auto_conf$hpc_user
  # Create project directory if it does not exist yet so that chunk_to_recipe
  # and split_to_recipe files can be created
  if (!dir.exists(proj_dir)) {
    dir.create(proj_dir, recursive = TRUE)
  }
  # Modify the configuration files according to the info in the recipe
  for (file in list.files(template_dir)) {
    conf_type <- strsplit(file, split = "[.]")[[1]][1]
    extension <- strsplit(file, split = "[.]")[[1]][2]
    dest_file <- paste0(conf_type, "_", expid, ".", extension)
    # Read configuration file
    conf <- read.config(file = paste0(template_dir, file))
    if (conf_type == "autosubmit") {
      # Section 1: autosubmit.conf
      ## expid, email notifications and address
      conf$config$EXPID <- expid
      conf$config$AUTOSUBMIT_VERSION <- auto_specs$auto_version
      if (recipe$Run$auto_conf$email_notifications) {
        conf$mail$NOTIFICATIONS <- "True"
      } else {
        conf$mail$NOTIFICATIONS <- "False"
      }
      conf$mail$TO <- recipe$Run$auto_conf$email_address
    } else if (conf_type == "expdef") {
      # Section 2: expdef
      ## expid, numchunks, project_type?, project_destination = auto-s2s
      conf$DEFAULT$EXPID <- expid
      conf$experiment$DATELIST <- format(Sys.Date(), "%Y%m%d")
      conf$experiment$NUMCHUNKS <- nchunks
      conf$local$PROJECT_PATH <- recipe$Run$code_dir
    } else if (conf_type == "jobs") {
      # Section 3: jobs
      ## wallclock, notify_on, platform?, processors, 
      # Create bash file to associate chunk number to recipe name
      chunk_file <- paste0(proj_dir, "/", "chunk_to_recipe")
      ## TODO: Specify dir in chunk_to_recipe?
      .create_bash_file(fileout = chunk_file,
                        dictionary = chunk_to_recipe,
                        variable = "CHUNK")
      # Define job parameters
      notify_on <- ""
      if (recipe$Run$auto_conf$notify_completed) {
        notify_on <- "COMPLETED"
      }
      if (recipe$Run$auto_conf$notify_failed) {
        notify_on <- paste(notify_on, "FAILED")
      }
      # Define last job before file transfer
      last_job <- "verification"
      # Verification job
      conf$JOBS$verification$WALLCLOCK <- recipe$Run$auto_conf$wallclock
      conf$JOBS$verification$NOTIFY_ON <- notify_on
      conf$JOBS$verification$PROCESSORS <- recipe$Run$auto_conf$processors_per_job # ncores?
      conf$JOBS$verification$CUSTOM_DIRECTIVES <- recipe$Run$auto_conf$custom_directives
      conf$JOBS$verification$PLATFORM <- platform
      # Recipe transfer job: only include if a tmp_dir is defined
      if (!is.null(tmp_dir)) {
        conf$JOBS$verification$DEPENDENCIES <- "transfer_recipes"
        conf$JOBS$transfer_recipes$NOTIFY_ON <- notify_on
      } else {
        conf$JOBS$transfer_recipes <- NULL
      }
      # Multimodel job: Only include if section exists in the recipe
      # is set to execute = 'True' or 'both'
      if (!is.null(recipe$Analysis$Datasets$Multimodel) && 
          tolower(recipe$Analysis$Datasets$Multimodel$execute) == "false") {
        conf$JOBS$multimodel <- NULL
        conf$JOBS$scorecards$DEPENDENCIES <- "verification"
      } else {
        conf$JOBS$multimodel$PLATFORM <- platform
        conf$JOBS$scorecards$DEPENDENCIES <- "multimodel"
        # Create bash file to associate split number to recipe name
        split_file <- paste0(proj_dir, "split_to_recipe")
        .create_bash_file(fileout = split_file,
                          dictionary = split_to_recipe,
                          variable = "SPLIT")
        # Define multimodel dependencies in the format required by AS config
        mm_dependencies <- lapply(split_to_recipe, grep, chunk_to_recipe)
        mm_dependencies <- lapply(mm_dependencies, paste, collapse = ",")
        names(mm_dependencies) <- paste(1:length(mm_dependencies))
        for (split in names(mm_dependencies)) {
          conf$JOBS$multimodel$DEPENDENCIES$verification$SPLITS_FROM[[split]]$CHUNKS_TO <- 
            mm_dependencies[[split]]
        }
        # 'Splits' parameter should be the number of mulimodel jobs
        conf$JOBS$multimodel$SPLITS <- length(mm_dependencies)
        # Define the rest of the parameters
        conf$JOBS$multimodel$NOTIFY_ON <- notify_on
        conf$JOBS$multimodel$PROCESSORS <- recipe$Run$auto_conf$processors_multimodel
        conf$JOBS$multimodel$CUSTOM_DIRECTIVES <- recipe$Run$auto_conf$custom_directives_multimodel
        conf$JOBS$multimodel$WALLCLOCK <- recipe$Run$auto_conf$wallclock_multimodel
        last_job <- "multimodel"
      }
      # Scorecards job: include only if section exists in the recipe and 
      # is set to 'execute: True'
      if (!("Scorecards" %in% names(recipe$Analysis$Workflow)) || 
	  (!recipe$Analysis$Workflow$Scorecards$execute)) {
        conf$JOBS$scorecards <- NULL
      } else {
        conf$JOBS$scorecards$NOTIFY_ON <- notify_on
        conf$JOBS$scorecards$PLATFORM <- platform
        last_job <- "scorecards"
      }
      # Results transfer job: only include if a tmp_dir is defined
      if (!is.null(tmp_dir)) {
        conf$JOBS$transfer_results$DEPENDENCIES <- last_job
        conf$JOBS$transfer_results$NOTIFY_ON <- notify_on
      } else {
        conf$JOBS$transfer_results <- NULL
      }
    } else if (conf_type == "platforms") {
      ## TODO: Allow user to choose platform
      # Section 4: platform configuration
      ## nord3v2 configuration... platform name? user, processors_per_node
      conf$Platforms[[platform]]$USER <- hpc_user
      unused_machines <- names(conf$Platforms)[!names(conf$Platforms) %in% c("transfer", platform)]
      conf$Platforms[unused_machines] <- NULL
      if (!is.null(tmp_dir)) {
        conf$Platforms$transfer$USER <- hpc_user
      } else {
        conf$Platforms$transfer <- NULL
      }
    } else if (conf_type == "proj") {
      # Section 5: proj
      ## modules? Info that goes on script, e.g. directory
      conf$common$OUTDIR <- recipe$Run$output_dir
      conf$common$SCRIPT <- recipe$Run$auto_conf$script
      conf$common$RECIPE <- paste0(recipe$name, ".yml")
      if (!is.null(tmp_dir)) {
        conf$common$TMPDIR <- gsub("\\$.+?\\$", hpc_user, tmp_dir)
      } else {
        conf$common$TMPDIR <- recipe$Run$output_dir
      }
    }
    # Write config file inside autosubmit dir
    write.config(conf, file.path(dest_dir, dest_file),
                 write.type = auto_specs$conf_format)
    Sys.chmod(file.path(dest_dir, dest_file), mode = "755", use_umask = F)
  }
  info(recipe$Run$logger,
       paste("##### AUTOSUBMIT CONFIGURATION WRITTEN FOR", expid, "#####"))
  info(recipe$Run$logger,
       paste0("You can check your experiment configuration at: ",
              dest_dir))
  # Print instructions/commands for user
  if (recipe$Run$Terminal) {
    ## TODO: Change SSH message for other environments (outside BSC)
    info(recipe$Run$logger,
         paste("Please SSH into bscesautosubmit01 or bscesautosubmit02 and run",
               "the following commands:"))
    info(recipe$Run$logger,
         paste("module load", auto_specs$module_version))
    info(recipe$Run$logger,
         paste("autosubmit create", expid))
    info(recipe$Run$logger,
         paste("autosubmit refresh", expid))
    info(recipe$Run$logger,
         paste("nohup autosubmit run", expid, "& disown"))
  } else {
    print(paste("Please SSH into bscesautosubmit03 or bscesautosubmit04 and run",
                "the following commands:"))
    print(paste("module load", auto_specs$module_version))
    print(paste("autosubmit create", expid))
    print(paste("autosubmit refresh", expid))
    print(paste("nohup autosubmit run", expid, "& disown"))
  }
}

.create_bash_file <- function(fileout, dictionary, variable) {
  file_connection <- file(fileout)
  script_lines <- paste0("case $", variable, " in")
  for (item in 1:length(dictionary)) {
    script_command <- paste0("  ", item, ") recipe='", dictionary[[item]], "' ;;")
    script_lines <- c(script_lines, script_command)
  }
  script_lines <- c(script_lines, "esac")
  writeLines(script_lines, file_connection)
  close(file_connection)
}
