# Load required libraries/funs
source("modules/Loading/R/get_regrid_params.R")
source("modules/Loading/R/dates2load.R")
source("modules/Loading/R/get_timeidx.R")
source("modules/Loading/R/check_latlon.R")
source("modules/Loading/R/compare_exp_obs_grids.R")

load_seasonal <- function(recipe) {
  
  # -------------------------------------------
  # Set params -----------------------------------------
  
  hcst.inityear <- recipe$Analysis$Time$hcst_start
  hcst.endyear <- recipe$Analysis$Time$hcst_end
  lats.min <- recipe$Analysis$Region$latmin
  lats.max <- recipe$Analysis$Region$latmax
  lons.min <- recipe$Analysis$Region$lonmin
  lons.max <- recipe$Analysis$Region$lonmax
  ref.name <- recipe$Analysis$Datasets$Reference$name
  exp.name <- recipe$Analysis$Datasets$System$name
  
  variable <- strsplit(recipe$Analysis$Variables$name, ", | |,")[[1]]
  store.freq <- recipe$Analysis$Variables$freq
  
  # get sdates array
  ## LOGGER: Change dates2load to extract logger from recipe?
  sdates <- dates2load(recipe, recipe$Run$logger)
  
  idxs <- NULL
  idxs$hcst <- get_timeidx(sdates$hcst,
                           recipe$Analysis$Time$ftime_min,
                           recipe$Analysis$Time$ftime_max,
                           time_freq=store.freq)

  if (!(is.null(sdates$fcst))) {
    idxs$fcst <- get_timeidx(sdates$fcst,
                             recipe$Analysis$Time$ftime_min,
                             recipe$Analysis$Time$ftime_max,
                             time_freq=store.freq)
  }

  ## TODO: Examine this verifications part, verify if it's necessary
  # stream <- verifications$stream
  # sdates <- verifications$fcst.sdate
  
  ## TODO: define fcst.name
  ##fcst.name <- recipe$Analysis$Datasets$System[[sys]]$name
  
  # get datasets dict:
  archive <- get_archive(recipe)
  exp_descrip <- archive$System[[exp.name]]
  
  freq.hcst <- unlist(exp_descrip[[store.freq]][variable[1]])
  reference_descrip <- archive$Reference[[ref.name]]
  freq.obs <- unlist(reference_descrip[[store.freq]][variable[1]])
  obs.dir <- reference_descrip$src
  fcst.dir <- exp_descrip$src
  hcst.dir <- exp_descrip$src
  fcst.nmember <- exp_descrip$nmember$fcst
  hcst.nmember <- exp_descrip$nmember$hcst
  
  ## TODO: it is necessary?
  ##if ("accum" %in% names(reference_descrip)) {
  ##  accum <- unlist(reference_descrip$accum[store.freq][[1]])
  ##} else {
  ##  accum <- FALSE
  ##}
  
  if (store.freq %in% c("daily", "daily_mean")) {
    frequency <- "daily_mean"
  } else if (store.freq == "monthly_mean") {
    frequency <- "monthly_mean"
  }
  var_dir_obs <- reference_descrip[[frequency]][variable]
  var_dir_exp <- exp_descrip[[frequency]][variable]
  
  # -----------
  obs.path <- paste0(archive$src_ref, obs.dir, "$var_dir$",
                     "$var$_$file_date$.nc")
  
  hcst.path <- paste0(archive$src_sys, hcst.dir, "$var_dir$",
                      "$var$_$file_date$.nc")
  
  fcst.path <- paste0(archive$src_sys, hcst.dir, "$var_dir$",
                      "$var$_$file_date$.nc")
  
  # Define regrid parameters:
  #-------------------------------------------------------------------
  regrid_params <- get_regrid_params(recipe, archive)
  
  # Longitude circular sort and latitude check
  #-------------------------------------------------------------------
  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
  }

  # Load hindcast 
  #-------------------------------------------------------------------
  
  hcst <- Start(dat = hcst.path,
                var = variable,
                var_dir = var_dir_exp,
                file_date = sdates$hcst,
                time = idxs$hcst,
                var_dir_depends = 'var',
                latitude = values(list(lats.min, lats.max)),
                latitude_reorder = Sort(),
                longitude = values(list(lons.min, lons.max)),
                longitude_reorder = circularsort,
                transform = regrid_params$fcst.transform,
                transform_params = list(grid = regrid_params$fcst.gridtype,
                                        method = regrid_params$fcst.gridmethod),
                transform_vars = c('latitude', 'longitude'),
                synonims = list(latitude = c('lat', 'latitude'),
                                longitude = c('lon', 'longitude'),
                                ensemble = c('member', 'ensemble', 'lev')),
                ensemble = indices(1:hcst.nmember),
                metadata_dims = 'var', # change to just 'var'?
                return_vars = list(latitude = 'dat',
                                   longitude = 'dat',
                                   time = 'file_date'),
                split_multiselected_dims = split_multiselected_dims,
                retrieve = TRUE)
  
  # Remove var_dir dimension
  if ("var_dir" %in% names(dim(hcst))) {
    hcst <- Subset(hcst, along = "var_dir", indices = 1, drop = "selected")
  }
  
  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
  ## TODO: Give correct dimensions to $Dates
  ## (sday, sweek, syear instead of file_date)
  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)) {
    # the call uses file_date instead of fcst_syear so that it can work
    # with the daily case and the current version of startR not allowing
    # multiple dims split
    
    fcst <- Start(dat = fcst.path,
                  var = variable,
                  var_dir = var_dir_exp,
                  var_dir_depends = 'var',
                  file_date = sdates$fcst,
                  time = idxs$fcst,
                  latitude = values(list(lats.min, lats.max)),
                  latitude_reorder = Sort(),
                  longitude = values(list(lons.min, lons.max)),
                  longitude_reorder = circularsort,
                  transform = regrid_params$fcst.transform,
                  transform_params = list(grid = regrid_params$fcst.gridtype,
                                          method = regrid_params$fcst.gridmethod),
                  transform_vars = c('latitude', 'longitude'),
                  synonims = list(latitude = c('lat', 'latitude'),
                                  longitude = c('lon', 'longitude'),
                                  ensemble = c('member', 'ensemble', 'lev')),
                  ensemble = indices(1:fcst.nmember),
                  metadata_dims = 'var',
                  return_vars = list(latitude = 'dat',
                                     longitude = 'dat',
                                     time = 'file_date'),
                  split_multiselected_dims = split_multiselected_dims,
                  retrieve = TRUE)
    
    if ("var_dir" %in% names(dim(fcst))) {
      fcst <- Subset(fcst, along = "var_dir", indices = 1, drop = "selected")
    }
    
    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
  #-------------------------------------------------------------------
  
  # 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")]
  
  # Separate Start() call for monthly vs daily data
  if (store.freq == "monthly_mean") {
    
    dates_file <- format(as.Date(dates, '%Y%m%d'), "%Y%m")
    dim(dates_file) <- dim(dates)
    
    obs <- Start(dat = obs.path,
                 var = variable,
                 var_dir = var_dir_obs,
                 var_dir_depends = 'var',
                 file_date = dates_file,
                 latitude = values(list(lats.min, lats.max)),
                 latitude_reorder = Sort(),
                 longitude = values(list(lons.min, lons.max)),
                 longitude_reorder = circularsort,
                 transform = regrid_params$obs.transform,
                 transform_params = list(grid = regrid_params$obs.gridtype,
                                         method = regrid_params$obs.gridmethod),
                 transform_vars = c('latitude', 'longitude'),
                 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")) {
    
    # 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,
                 var_dir = var_dir_obs,
                 var_dir_depends = 'var',
                 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 = values(list(lats.min, lats.max)),
                 latitude_reorder = Sort(),
                 longitude = values(list(lons.min, lons.max)),
                 longitude_reorder = circularsort,
                 transform = regrid_params$obs.transform,
                 transform_params = list(grid = regrid_params$obs.gridtype,
                                         method = regrid_params$obs.gridmethod),
                 transform_vars = c('latitude', 'longitude'),
                 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)
  }
  
  
  # Remove var_dir dimension
  if ("var_dir" %in% names(dim(obs))) {
    obs <- Subset(obs, along = "var_dir", indices = 1, drop = "selected")
  }
  # 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 (!(recipe$Analysis$Regrid$type == 'none')) {
    compare_exp_obs_grids(hcst, obs)
  }
  
  # Remove negative values in accumulative variables
  dictionary <- read_yaml("conf/variable-dictionary.yml")
  for (var_idx in 1:length(variable)) {
    var_name <- variable[var_idx]
    if (dictionary$vars[[var_name]]$accum) {
      info(recipe$Run$logger,
           paste0("Accumulated variable ", var_name,
                  ": setting negative values to zero."))
      # obs$data[, var_idx, , , , , , , ] <- pmax(Subset(obs$data,
      #                                                along = "var",
      #                                                indices = var_idx, F), 0)
      obs$data[, var_idx, , , , , , , ][obs$data[, var_idx, , , , , , , ] < 0] <- 0
      hcst$data[, var_idx, , , , , , , ][hcst$data[, var_idx, , , , , , , ] < 0] <- 0
      if (!is.null(fcst)) {
        fcst$data[, var_idx, , , , , , , ][fcst$data[, var_idx, , , , , , , ] < 0] <- 0
      }
    }
  }
  info(recipe$Run$logger,
       "##### DATA LOADING COMPLETED SUCCESSFULLY #####")
  
  ############################################################################
  #
  #     CHECKS ON MISSING FILES
  #
  ############################################################################
  
  #obs.NA_dates.ind <- Apply(obs, 
  #                          fun=(function(x){ all(is.na(x))}), 
  #                          target_dims=c('time', 'latitude', 'longitude'))[[1]]
  #obs.NA_dates <- dates_file[obs.NA_dates.ind]
  #obs.NA_dates <- obs.NA_dates[order(obs.NA_dates)]
  #obs.NA_files <- paste0(obs.dir, store.freq,"/",variable,"_",
  #                       freq.obs,"obs.grid","/",variable,"_",obs.NA_dates,".nc")
  #
  #if (any(is.na(hcst))){
  #  fatal(recipe$Run$logger,
  #    paste(" ERROR: MISSING HCST VALUES FOUND DURING LOADING # ",
  #          " ################################################# ",
  #          " ######           MISSING FILES               #### ",
  #          " ################################################# ",
  #          "hcst files:",
  #          hcst.NA_files,
  #          " ################################################# ",
  #          " ################################################# ",
  #          sep="\n"))
  #  quit(status = 1)
  #}
  #
  #if (any(is.na(obs)) && !identical(obs.NA_dates,character(0))){
  #  fatal(recipe$logger,
  #    paste(" ERROR: MISSING OBS VALUES FOUND DURING  LOADING # ",
  #          " ################################################# ",
  #          " ######           MISSING FILES               #### ",
  #          " ################################################# ",
  #          "obs files:",
  #          obs.NA_files,
  #          " ################################################# ",
  #          " ################################################# ",
  #          sep="\n"))
  #  quit(status=1)
  #}
  #
  #info(recipe$logger,
  #     "######### DATA LOADING COMPLETED SUCCESFULLY ##############")
  
  ############################################################################
  ############################################################################
  .log_memory_usage(recipe$Run$logger, when = "After loading")  
  return(list(hcst = hcst, fcst = fcst, obs = obs))
  
}
