
source("modules/Loading/R/dates2load.R")
source("modules/Loading/R/get_timeidx.R")
source("modules/Loading/R/check_latlon.R")

load_multimodel <- function(recipe) {
  
  archive <- read_yaml("conf/archive.yml")$esarchive
  ref.name <- recipe$Analysis$Datasets$Reference$name
  if (tolower(recipe$Analysis$Horizon) == 'seasonal') {
    exp.name <- recipe$Analysis$Datasets$System$models
  } else if (tolower(recipe$Analysis$Horizon) == 'decadal') {
    exp.name <- sapply(recipe$Analysis$Datasets$System$models, '[[', 'name')
  } else {stop('Multimodel not implemented for this horizon.')}
  store.freq <- recipe$Analysis$Variables$freq
  variable <- strsplit(recipe$Analysis$Variables$name, ", | |,")[[1]]
  exp_descrip <- archive$System[[exp.name[1]]]
  reference_descrip <- archive$Reference[[ref.name]
  sdates <- dates2load(recipe, recipe$Run$logger)
  createFrom <- str_to_title(recipe$Analysis$Datasets$Multimodel$createFrom)
  
  lats.min <- recipe$Analysis$Region$latmin
  lats.max <- recipe$Analysis$Region$latmax
  lons.min <- recipe$Analysis$Region$lonmin
  lons.max <- recipe$Analysis$Region$lonmax
  circularsort <- check_latlon(lats.min, lats.max, lons.min, lons.max)
  
  if (recipe$Analysis$Variables$freq == "monthly_mean") {
    split_multiselected_dims = TRUE
  } else {
    split_multiselected_dims = FALSE
  }
  
  # Find the saved data directory
  recipe$Run$output_dir <- file.path(recipe$Run$output_dir, "outputs", createFrom)
  if (tolower(recipe$Analysis$Output_format) == "scorecards") {
    hcst_start <- recipe$Analysis$Time$hcst_start
    hcst_end <- recipe$Analysis$Time$hcst_end
    shortdate <- substr(recipe$Analysis$Time$sdate, start = 1, stop = 2)
    filename <- paste0("scorecards_$model$_", ref.name, "_$var$_$file_date$_",
                       hcst_start, "-", hcst_end, "_s", shortdate, ".nc")
  } else {
    filename <- "$var$_$file_date$.nc"
  }

  hcst.path <- file.path(get_dir(recipe = recipe, variable = variable[1]), filename)
  hcst.path <- gsub(variable[1], "$var$", hcst.path)
  hcst.path <- gsub('Multimodel', "$model$", hcst.path)
  fcst.path <- obs.path <- hcst.path
  obs.path <- gsub("_$file_date$", "-obs_$file_date$", obs.path, fixed = T)
  obs.path <- gsub("$model$", gsub('\\.','',exp.name[1]), obs.path, fixed = T)
  
  # Load hindcast 
  #-------------------------------------------------------------------
  hcst <- Start(dat = hcst.path,
                var = variable,
                file_date = sdates$hcst,
                model = gsub('\\.','',exp.name),
                time = 'all',
                latitude = 'all',
                latitude_reorder = Sort(),
                longitude = 'all',
                longitude_reorder = circularsort,
                synonims = list(latitude = c('lat', 'latitude'),
                                longitude = c('lon', 'longitude'),
                                ensemble = c('member', 'ensemble')),
                ensemble = 'all',
                metadata_dims = 'var',
                largest_dims_length = TRUE,
                return_vars = list(latitude = 'dat',
                                   longitude = 'dat',
                                   time = 'file_date'),
                split_multiselected_dims = split_multiselected_dims,
                retrieve = TRUE)
  
  #############################
  #NOTE: NOT TESTED YET
  if (store.freq %in% c("daily_mean", "daily")) {
    # Adjusts dims for daily case, could be removed if startR allows
    # multidim split
    names(dim(hcst))[which(names(dim(hcst)) == 'file_date')] <- "syear"
    default_dims <- c(dat = 1, var = 1, sday = 1, 
                      sweek = 1, syear = 1, time = 1, 
                      latitude = 1, longitude = 1, ensemble = 1)
    default_dims[names(dim(hcst))] <- dim(hcst)
    dim(hcst) <- default_dims
    # Change time attribute dimensions
    default_time_dims <- c(sday = 1, sweek = 1, syear = 1, time = 1)
    names(dim(attr(hcst, "Variables")$common$time))[which(names(
      dim(attr(hcst, "Variables")$common$time)) == 'file_date')] <- "syear"
    default_time_dims[names(dim(attr(hcst, "Variables")$common$time))] <- 
      dim(attr(hcst, "Variables")$common$time)
    dim(attr(hcst, "Variables")$common$time) <- default_time_dims
  }
  ###############################
  
  # Convert hcst to s2dv_cube object
  hcst <- as.s2dv_cube(hcst)
  # Adjust dates for models where the time stamp goes into the next month
  if (recipe$Analysis$Variables$freq == "monthly_mean") {
    hcst$attrs$Dates[] <- hcst$attrs$Dates - seconds(exp_descrip$time_stamp_lag)
  }
  
  # Load forecast
  #-------------------------------------------------------------------
  if (!is.null(recipe$Analysis$Time$fcst_year)) {
    fcst <- Start(dat = fcst.path,
                  var = variable,
                  file_date = sdates$fcst,
                  model = gsub('\\.','',exp.name),
                  time = 'all',
                  latitude = 'all',
                  latitude_reorder = Sort(),
                  longitude = 'all',
                  longitude_reorder = circularsort,
                  synonims = list(latitude = c('lat', 'latitude'),
                                  longitude = c('lon', 'longitude'),
                                  ensemble = c('member', 'ensemble')),
                  ensemble = 'all',
                  metadata_dims = 'var',
                  largest_dims_length = TRUE,
                  return_vars = list(latitude = 'dat',
                                     longitude = 'dat',
                                     time = 'file_date'),
                  split_multiselected_dims = split_multiselected_dims,
                  retrieve = TRUE)
    
    #############################
    #NOTE: NOT TESTED YET
    if (store.freq %in% c("daily_mean", "daily")) {
      # Adjusts dims for daily case, could be removed if startR allows
      # multidim split
      names(dim(fcst))[which(names(dim(fcst)) == 'file_date')] <- "syear"
      default_dims <- c(dat = 1, var = 1, sday = 1,
                        sweek = 1, syear = 1, time = 1,
                        latitude = 1, longitude = 1, ensemble = 1)
      default_dims[names(dim(fcst))] <- dim(fcst)
      dim(fcst) <- default_dims
      # Change time attribute dimensions
      default_time_dims <- c(sday = 1, sweek = 1, syear = 1, time = 1)
      names(dim(attr(fcst, "Variables")$common$time))[which(names(
        dim(attr(fcst, "Variables")$common$time)) == 'file_date')] <- "syear"
      default_time_dims[names(dim(attr(fcst, "Variables")$common$time))] <- 
        dim(attr(fcst, "Variables")$common$time)
      dim(attr(fcst, "Variables")$common$time) <- default_time_dims
    }
    #############################
    
    # Convert fcst to s2dv_cube
    fcst <- as.s2dv_cube(fcst)
    # Adjust dates for models where the time stamp goes into the next month
    if (recipe$Analysis$Variables$freq == "monthly_mean") {
      fcst$attrs$Dates[] <- 
        fcst$attrs$Dates - seconds(exp_descrip$time_stamp_lag)
    }
    
  } else {
    fcst <- NULL
  }
  
  # Load reference
  #-------------------------------------------------------------------
  
  if (store.freq == "monthly_mean") {
    
    obs <- Start(dat =  obs.path,
                 var = variable,
                 file_date = sdates$hcst,
                 time = 'all',
                 latitude = 'all',
                 latitude_reorder = Sort(),
                 longitude = 'all',
                 longitude_reorder = circularsort,
                 synonims = list(latitude = c('lat','latitude'),
                                 longitude = c('lon','longitude')),
                 metadata_dims = 'var',
                 return_vars = list(latitude = 'dat',
                                    longitude = 'dat',
                                    time = 'file_date'),
                 split_multiselected_dims = TRUE,
                 retrieve = TRUE)
    
  } else if (store.freq %in% c("daily_mean", "daily")) {
    
    #############################
    #NOTE: NOT TESTED YET
    
    # Obtain dates and date dimensions from the loaded hcst data to make sure
    # the corresponding observations are loaded correctly.
    dates <- hcst$attrs$Dates
    dim(dates) <- hcst$dims[c("sday", "sweek", "syear", "time")]
    
    # Get year and month for file_date
    dates_file <- sapply(dates, format, '%Y%m')
    dim(dates_file) <- dim(dates)
    # Set hour to 12:00 to ensure correct date retrieval for daily data
    lubridate::hour(dates) <- 12 
    lubridate::minute(dates) <- 00
    # Restore correct dimensions
    dim(dates) <- dim(dates_file)
    
    obs <- Start(dat = obs.path,
                 var = variable,
                 file_date = sort(unique(dates_file)),
                 time = dates,
                 time_var = 'time',
                 time_across = 'file_date',
                 merge_across_dims = TRUE,
                 merge_across_dims_narm = TRUE,
                 latitude = 'all',
                 latitude_reorder = Sort(),
                 longitude = 'all',
                 longitude_reorder = circularsort,
                 synonims = list(latitude = c('lat','latitude'),
                                 longitude = c('lon','longitude')),
                 metadata_dims = 'var',
                 return_vars = list(latitude = 'dat',
                                    longitude = 'dat',
                                    time = 'file_date'),
                 split_multiselected_dims = TRUE,
                 retrieve = TRUE)
    #############################
    
  }
  
  # Adds ensemble dim to obs (for consistency with hcst/fcst)
  default_dims <- c(dat = 1, var = 1, sday = 1, 
                    sweek = 1, syear = 1, time = 1, 
                    latitude = 1, longitude = 1, ensemble = 1)
  default_dims[names(dim(obs))] <- dim(obs)
  dim(obs) <- default_dims
  
  # Convert obs to s2dv_cube
  obs <- as.s2dv_cube(obs)
  
  # Check for consistency between hcst and obs grid
  if (!isTRUE(all.equal(as.vector(hcst$coords$latitude), as.vector(obs$coords$latitude)))) {
    lat_error_msg <- paste("Latitude mismatch between hcst and obs.",
                           "Please check the original grids and the",
                           "regrid parameters in your recipe.")
    error(recipe$Run$logger, lat_error_msg)
    hcst_lat_msg <- paste0("First hcst lat: ", hcst$lat[1],
                           "; Last hcst lat: ", hcst$lat[length(hcst$lat)])
    info(recipe$Run$logger, hcst_lat_msg)
    obs_lat_msg <- paste0("First obs lat: ", obs$lat[1],
                          "; Last obs lat: ", obs$lat[length(obs$lat)])
    info(recipe$Run$logger, obs_lat_msg)
    stop("hcst and obs don't share the same latitudes.")
  }
  if (!isTRUE(all.equal(as.vector(hcst$coords$longitude), as.vector(obs$coords$longitude)))) {
    lon_error_msg <- paste("Longitude mismatch between hcst and obs.",
                           "Please check the original grids and the",
                           "regrid parameters in your recipe.")
    error(recipe$Run$logger, lon_error_msg)
    hcst_lon_msg <- paste0("First hcst lon: ", hcst$lon[1],
                           "; Last hcst lon: ", hcst$lon[length(hcst$lon)])
    info(recipe$Run$logger, hcst_lon_msg)
    obs_lon_msg <- paste0("First obs lon: ", obs$lon[1],
                          "; Last obs lon: ", obs$lon[length(obs$lon)])
    info(recipe$Run$logger, obs_lon_msg)
    stop("hcst and obs don't share the same longitudes.")  
  }
  
  # Print a summary of the loaded data for the user, for each object
  if (recipe$Run$logger$threshold <= 2) {
    data_summary(hcst, recipe)
    data_summary(obs, recipe)
    if (!is.null(fcst)) {
      data_summary(fcst, recipe)
    }
  }
  
  info(recipe$Run$logger,
       "##### DATA LOADING COMPLETED SUCCESSFULLY #####")
  
  return(list(hcst = hcst, fcst = fcst, obs = obs))
}
