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

#====================================================================

# recipe_file <- "recipes/atomic_recipes/recipe_decadal.yml"
# recipe_file <- "recipes/atomic_recipes/recipe_decadal_daily.yml"

load_decadal <- function(recipe) {
  ## 
  archive <- get_archive(recipe)
  
  # Print Start() info or not
  DEBUG <- TRUE
  
  ## TODO: this should come from the main script 
  # Create output folder and log:
  
  #-------------------------
  # Read from recipe:
  #-------------------------
  exp.name <- recipe$Analysis$Datasets$System$name # 'HadGEM3'
  ref.name <- recipe$Analysis$Datasets$Reference$name # 'era5'
  member <- strsplit(recipe$Analysis$Datasets$System$member, ', | |,')[[1]] #c("r1i1p1f2", "r2i1p1f2")
  variable <- strsplit(recipe$Analysis$Variables$name, ", | |,")[[1]]
  store.freq <- recipe$Analysis$Variables$freq #monthly_mean
  lats.min <- as.numeric(recipe$Analysis$Region$latmin) #0
  lats.max <- as.numeric(recipe$Analysis$Region$latmax) #10
  lons.min <- as.numeric(recipe$Analysis$Region$lonmin) #0
  lons.max <- as.numeric(recipe$Analysis$Region$lonmax) #10
  
  # change to: sdates <- dates2load(recipe, logger)
  sdates_hcst <- as.numeric(recipe$Analysis$Time$hcst_start):as.numeric(recipe$Analysis$Time$hcst_end) #1960:2015
  sdates_fcst <- recipe$Analysis$Time$fcst
  
  if (store.freq == "monthly_mean") {
    time_ind <- (as.numeric(recipe$Analysis$Time$ftime_min):as.numeric(recipe$Analysis$Time$ftime_max))    
  } else if (store.freq %in% c("daily", "daily_mean")) {
    time_ind <- get_daily_time_ind(ftimemin = as.numeric(recipe$Analysis$Time$ftime_min),
                                   ftimemax = as.numeric(recipe$Analysis$Time$ftime_max),
                                   initial_month = archive$System[[exp.name]]$initial_month,
                                   sdates = sdates_hcst,
                                   calendar = archive$System[[exp.name]]$calendar)
    store.freq <- "daily_mean"
  }
  
  #NOTE: May be used in the future
  #  season <- recipe$Analysis$Time$season
  
  #-------------------------
  # Read from archive:
  #-------------------------
  if (store.freq == "monthly_mean") {
    table <- archive$System[[exp.name]][[store.freq]]$table[variable]  #list(tas = 'Amon')
    grid <- archive$System[[exp.name]][[store.freq]]$grid[variable]  #list(tas = 'gr')
  } else {
    table <- 'day'
    # For grid, get first element as they are all the same?
    grid <- archive$System[[exp.name]][[store.freq]]$grid[variable][[1]]
  }
  # grid <- archive$System[[exp.name]][[store.freq]]$grid[variable]  #list(tas = 'gr')
  version <- archive$System[[exp.name]][[store.freq]]$version[variable]  #list(tas = 'v20210910')
  if (identical(member, 'all')) {
    member <- strsplit(archive$System[[exp.name]]$member, ',')[[1]]
  }
  
  #-------------------------
  # derived from above:
  #-------------------------
  # Check lat and lon and decide CircularSort
  circularsort <- check_latlon(latmin = lats.min, latmax = lats.max, lonmin = lons.min, lonmax = lons.max)
  
  # generate transform params for system and ref
  regrid_params <- get_regrid_params(recipe, archive)
  
  # Only if the time length in each chunk may differ that we need largest_dims_length to be TRUE. Otherwise, set FALSE to increase efficiency.
  need_largest_dims_length <- ifelse(exp.name %in% c('HadGEM3-GC31-MM', 'EC-Earth3-i2'), TRUE, FALSE)
  
  #-------------------------------------------
  # Step 1: Load the hcst
  #-------------------------------------------
  #monthly and daily
  tmp <- get_dcpp_path(archive = archive, exp.name = exp.name, table = table, grid = grid,
                       version = version, sdates = sdates_hcst)
  path_list <- tmp$path_list
  dcpp_list <- tmp$dcpp_list
  
  Start_default_arg_list <- list(
    dat = path_list,
    var = variable,
    syear = paste0(sdates_hcst),
    chunk = 'all',
    chunk_depends = 'syear',
    dcpp = dcpp_list,
    dcpp_depends = 'syear',
    time = indices(time_ind),
    time_across = 'chunk',
    merge_across_dims = TRUE,
    largest_dims_length = need_largest_dims_length,
    latitude = values(list(lats.min, lats.max)),
    latitude_reorder = Sort(decreasing = TRUE),
    longitude = values(list(lons.min, lons.max)),
    longitude_reorder = circularsort,
    ensemble = member,
    path_glob_permissive = 2, # for version
    transform = regrid_params$fcst.transform,
    transform_extra_cells = 2,
    transform_params = list(grid = regrid_params$fcst.gridtype,
                            method = regrid_params$fcst.gridmethod),
    transform_vars = c('latitude', 'longitude'),
    synonims = list(longitude = c('lon', 'longitude'),
                    latitude = c('lat', 'latitude')),
    return_vars = list(latitude = NULL, longitude = NULL,
                       time = c('syear', 'chunk')),
    silent = !DEBUG,
    retrieve = T)
  
  if (length(table) > 1) {
    Start_default_arg_list <- c(Start_default_arg_list, 
                                list(table = table, grid = grid, version = version,
                                     table_depends = 'var', grid_depends = 'var', version_depends = 'var',
                                     metadata_dims = 'var'))
    Start_default_arg_list[["path_glob_permissive"]] <- FALSE
  }
 
  Start_hcst_arg_list <- Start_default_arg_list
  hcst <- do.call(Start, Start_hcst_arg_list)
    
  tmp_time_attr <- attr(hcst, 'Variables')$common$time
 
  ## TODO: Remove this part? 
  # change syear to c(sday, sweek, syear)
  # dim(hcst) should be [dat, var, sday, sweek, syear, time, latitude, longitude, ensemble]
  dim(hcst) <- c(dim(hcst)[1:2], sday = 1, sweek = 1, dim(hcst)[3:length(dim(hcst))])
  if (!identical(dim(tmp_time_attr), dim(hcst)[c('syear', 'time')])) {
    error(recipe$Run$logger,
          "hcst has problem in matching data and time attr dimension.")
    stop()
  }
  dim(attr(hcst, 'Variables')$common$time) <- c(sday = 1, sweek = 1, dim(tmp_time_attr))
  
  # Remove 'dcpp' and other extra dimensions:
  if (length(table) > 1) {
    hcst <- Subset(hcst, along = c("dcpp", "table", "grid", "version"),
                   indices = list(1, 1, 1, 1), drop = "selected")
  } else {
    hcst <- Subset(hcst, along = "dcpp", indices = 1, drop = "selected")
  }
  # Change class from startR_array to s2dv_cube
  suppressWarnings(
    hcst <- as.s2dv_cube(hcst)
  )
  
  #-------------------------------------------
  # Step 2: Load the fcst
  #-------------------------------------------
  if (!is.null(recipe$Analysis$Time$fcst)) {
    
    tmp <- get_dcpp_path(archive = archive, exp.name = exp.name, table = table, grid = grid, 
                         version = version, sdates = sdates_fcst)
    path_list <- tmp$path_list
    dcpp_list <- tmp$dcpp_list
    
    # monthly & daily
    Start_fcst_arg_list <- Start_default_arg_list
    Start_fcst_arg_list[['dat']] <- path_list
    Start_fcst_arg_list[['syear']] <- paste0(sdates_fcst)
    Start_fcst_arg_list[['dcpp']] <- dcpp_list
    fcst <- do.call(Start, Start_fcst_arg_list)
      
    tmp_time_attr <- attr(fcst, 'Variables')$common$time
    
    # change syear to c(sday, sweek, syear)
    # dim(fcst) should be [dat, var, sday, sweek, syear, time, latitude, longitude, ensemble]
    dim(fcst) <- c(dim(fcst)[1:2], sday = 1, sweek = 1, dim(fcst)[3:length(dim(fcst))])
    if (!identical(dim(tmp_time_attr), dim(fcst)[c('syear', 'time')])) {
      error(recipe$Run$logger,
            "fcst has problem in matching data and time attr dimension.")
      stop()
    }
    dim(attr(fcst, 'Variables')$common$time) <- c(sday = 1, sweek = 1, dim(tmp_time_attr))

    # Remove 'dcpp' and any other extra dimensions:
    if (length(table) > 1) {
      fcst <- Subset(fcst, along = c("dcpp", "table", "grid", "version"),
                     indices = list(1, 1, 1, 1), drop = "selected")
    } else {
      fcst <- Subset(fcst, along = "dcpp", indices = 1, drop = "selected")
    }
    
    # Change class from startR_array to s2dv_cube
    suppressWarnings(
      fcst <- as.s2dv_cube(fcst)
    )
    
    # Only syear could be different
    syear_dim <- which(names(dim(hcst$data)) == 'syear')
    if (!identical(dim(hcst$data)[-syear_dim], dim(fcst$data)[-syear_dim])) {
      error(recipe$Run$logger,
            "hcst and fcst do not share the same dimension structure.")
      stop()
    }
    
  } else {
    fcst <- NULL
  }
  
  #-------------------------------------------
  # Step 3. Load the reference
  #-------------------------------------------
  obs.path <- file.path(archive$src_ref, archive$Reference[[ref.name]]$src,
                        "$var_dir$", "$var$_$file_date$.nc")
  var_dir_obs <- archive$Reference[[ref.name]][[store.freq]][variable] # list(tas = "_f1h-r1440x721cds", tos = "_f1h-r1440x721cds")
  
  #  obs.path <- file.path(archive$src, archive$Reference[[ref.name]]$src, store.freq,
  #                        paste0(variable, archive$Reference[[ref.name]][[store.freq]][[variable]]))
  #  obs.files <- paste0('$var$_$file_date$.nc')
  
  # Get from startR_cube
  #  dates <- attr(hcst, 'Variables')$common$time
  # Get from s2dv_cube
  dates <- hcst$attrs$Dates  
  dates_file <- sapply(dates, format, '%Y%m')
  dim(dates_file) <- dim(dates)
  
  if (store.freq == "daily_mean") {
    #////////////////
    # Method 1: use hcst time attr as obs time selector
    #////////////////
    
    # 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 = unique(format(dates, '%Y%m')),
                 time = dates, # [sday, sweek, syear, time]
                 time_across = 'file_date',
                 merge_across_dims = TRUE,
                 split_multiselected_dims = TRUE,
                 latitude = values(list(lats.min, lats.max)),
                 latitude_reorder = Sort(decreasing = TRUE), 
                 longitude = values(list(lons.min, lons.max)),
                 longitude_reorder = circularsort,
                 transform = regrid_params$obs.transform,
                 transform_extra_cells = 2,
                 transform_params = list(grid = regrid_params$obs.gridtype, #nc file
                                         method = regrid_params$obs.gridmethod),
                 transform_vars = c('latitude', 'longitude'),
                 synonims = list(latitude = c('lat','latitude'), 
                                 longitude = c('lon','longitude')),
                 return_vars = list(latitude = NULL, longitude = NULL,
                                    time = 'file_date'),
                 silent = !DEBUG,
                 retrieve = TRUE)
    
  } else if (store.freq == "monthly_mean") {
    #////////////////
    # Method 2: reshape hcst time attr's date into an array with time dim then as obs date selector
    #////////////////
   
    obs <- Start(dat = obs.path,
                 var = variable,
                 var_dir = var_dir_obs,
                 var_dir_depends = 'var',
                 file_date = dates_file, #dates_arr, # [sday, sweek, syear, time]
                 split_multiselected_dims = TRUE,
                 latitude = values(list(lats.min, lats.max)),
                 latitude_reorder = Sort(decreasing = TRUE),
                 longitude = values(list(lons.min, lons.max)),
                 longitude_reorder = circularsort,
                 transform = regrid_params$obs.transform,
                 transform_extra_cells = 2,
                 transform_params = list(grid = regrid_params$obs.gridtype, #nc file
                                         method = regrid_params$obs.gridmethod),
                 transform_vars = c('latitude', 'longitude'),
                 synonims = list(latitude = c('lat','latitude'),
                                 longitude = c('lon','longitude')),
                 return_vars = list(latitude = NULL, longitude = NULL,
                                    time = 'file_date'),
                 metadata_dims = 'var',
                 silent = !DEBUG,
                 retrieve = TRUE)
  }
  
  # Remove var_dir dimension
  obs <- Subset(obs, along = "var_dir", indices = 1, drop = "selected")
  
  # Only ensemble dim could be different
  if (!identical(dim(obs), dim(hcst$data)[-9])) {
    error(recipe$Run$logger,
          "obs and hcst dimensions do not match.")
    stop()
  }
  # Add ensemble dim to obs
  dim(obs) <- c(dim(obs), ensemble = 1)
  
  # Change class from startR_array to s2dv_cube
  suppressWarnings(
    obs <- as.s2dv_cube(obs)
  )
  
  #-------------------------------------------
  # Step 4. Verify the consistance between data
  #-------------------------------------------
  # dimension
  if (any(!names(dim(obs$data)) %in% names(dim(hcst$data)))) {
    error(recipe$Run$logger,
          "hcst and obs don't share the same dimension names.")
    stop()
  } else {
    ens_ind <- which(names(dim(obs$data)) == 'ensemble')
    match_ind <- match(names(dim(obs$data))[-ens_ind], names(dim(hcst$data)))
    if (!all(dim(hcst$data)[match_ind] == dim(obs$data)[-ens_ind])) {
      error(recipe$Run$logger,
            "hcst and obs don't share the same dimension length.")
      stop()
    }
  }
  
  # time attribute
  if (!identical(format(hcst$attrs$Dates, '%Y%m'),
                 format(obs$attrs$Dates, '%Y%m'))) {
    error(recipe$Run$logger,
          "hcst and obs don't share the same time.")
    stop()
  }
  
  # lat and lon attributes
  if (!(recipe$Analysis$Regrid$type == 'none')) {
    if (!isTRUE(all.equal(as.vector(hcst$lat), as.vector(obs$lat)))) {
      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$lon), as.vector(obs$lon)))) {
      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.")
    }
  }
  
  # Check fcst
  if (!is.null(fcst)) {
    # dimension
    if (any(!names(dim(fcst$data)) %in% names(dim(hcst$data)))) {
      error(recipe$Run$logger,
            "hcst and fcst don't share the same dimension names.")
      stop()
    } else {
      ens_ind <- which(names(dim(fcst$data)) %in% c('ensemble', 'syear'))
      match_ind <- match(names(dim(fcst$data))[-ens_ind], names(dim(hcst$data)))
      if (!all(dim(hcst$data)[match_ind] == dim(fcst$data)[-ens_ind])) {
        error(recipe$Run$logger,
              "hcst and fcst don't share the same dimension length.")
        stop()
      }
    }
    
    # lat and lon attributes
    if (!(recipe$Analysis$Regrid$type == 'none')) {
      compare_exp_obs_grids(hcst, obs)
    }
  }
  
  #-------------------------------------------
  # Step 5. Tune data
  #-------------------------------------------
  # # 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
  #     }
  #   }
  
  #-------------------------------------------
  # Step 6. Print summary
  #-------------------------------------------
  
  info(recipe$Run$logger,
       "##### DATA LOADING COMPLETED SUCCESSFULLY #####")
  
  .log_memory_usage(recipe$Run$logger, when = "After loading")
  return(list(hcst = hcst, fcst = fcst, obs = obs))
}
