# 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, retrieve = retrieve) {
  ## 
  archive <- read_yaml(paste0("conf/archive_decadal.yml"))[[recipe$Run$filesystem]]
  
  # Print Start() info or not
  DEBUG <- TRUE
  # 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 <- recipe$Analysis$Variables$name #'tas'
  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 == "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)
  }
  
  #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')
  } else {
    table <- 'day'
  }
  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
  multi_path <- tmp$multi_path
  
  #TODO: to make this case work; enhance Start() if it's possible
  if (multi_path & length(variable) > 1) {
    stop("The recipe requests multiple variables and start dates from both dpccA-hindcast and dcppB-forecast. This case is not available for now.")
  }
  
  Start_default_arg_list <- list(
    dat = path_list,
    var = variable,
    syear = paste0(sdates_hcst),
    chunk = 'all',
    chunk_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,
    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'),
    #    path_glob_permissive = 2, # for version
    synonims = list(longitude = c('lon', 'longitude'),
                    latitude = c('lat', 'latitude')),
    return_vars = list(latitude = NULL, longitude = NULL,
                       time = c('syear', 'chunk')),
    silent = !DEBUG,
    retrieve = retrieve)
  
  if (length(variable) > 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'))
  }
  
  if (!multi_path) {
    Start_hcst_arg_list <- Start_default_arg_list
    hcst <- do.call(Start, Start_hcst_arg_list)
    hcst_dims <- attr(hcst, "Dimensions")
  } else {
    Start_hcst_arg_list <- Start_default_arg_list
    Start_hcst_arg_list[['syear']] <- NULL
    Start_hcst_arg_list[['chunk_depends']] <- NULL
    remove_ind <- which(Start_hcst_arg_list[['return_vars']][['time']] == 'syear')
    Start_hcst_arg_list[['return_vars']][['time']] <- Start_hcst_arg_list[['return_vars']][['time']][-remove_ind]
    
    hcst <- do.call(Start, Start_hcst_arg_list)
    # Reshape and reorder dimensions
    ## dat should be 1, syear should be length of dat; reorder dimensions
    dim(hcst) <- c(dat = 1, syear = as.numeric(dim(hcst))[1], dim(hcst)[2:6])
    hcst <- s2dv::Reorder(hcst, c('dat', 'var', 'syear', 'time', 'latitude', 'longitude', 'ensemble'))
    
    # Manipulate time attr because Start() cannot read it correctly
    wrong_time_attr <- attr(hcst, 'Variables')$common$time # dim: [time], the first syear only
    tmp <- array(dim = c(dim(hcst)[c('syear', 'time')]))
    tmp[1, ] <- wrong_time_attr
    yr_diff <- (sdates_hcst - sdates_hcst[1])[-1] #diff(sdates_hcst)
    for (i_syear in 1:length(yr_diff)) {
      tmp[(i_syear + 1), ] <- wrong_time_attr + lubridate::years(yr_diff[i_syear])
    }
    attr(hcst, 'Variables')$common$time <- as.POSIXct(tmp, origin = '1970-01-01', tz = 'UTC')
    
  }
  
  tmp_time_attr <- attr(hcst, 'Variables')$common$time
 
  ## TODO: for retrieve = FALSE, do this in pre-processing
  # change syear to c(sday, sweek, syear)
  # dim(hcst) should be [dat, var, sday, sweek, syear, time, latitude, longitude, ensemble]
  if (retrieve) {
    dim(hcst) <- c(dim(hcst)[1:2], sday = 1, sweek = 1, dim(hcst)[3:7])
  } else {

  ## TODO: Why tf does this raise an error?
  # if (!identical(dim(tmp_time_attr), hcst_dims[c('syear', 'time')])) {
  #   error(recipe$Run$logger, retrieve = retrieve,
  #         "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))
  # Get hcst dates
  dates <- attr(hcst, 'Variables')$common$time
  
  #TODO: as.s2dv_cube() needs to be improved to recognize "variable" is under $dat1
  if (multi_path) {
    attributes(hcst)$Variables$common[[variable]] <- attributes(hcst)$Variables$dat1[[variable]]
  }
  
  # Change class from startR_array to s2dv_cube
  if (retrieve) {
    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
    multi_path <- tmp$multi_path
    
    #TODO: to make this case work; enhance Start() if it's possible
    if (multi_path & length(variable) > 1) {
      stop("The recipe requests multiple variables and start dates from both dpccA-hindcast and dcppB-forecast. This case is not available for now.")
    }
    
    # monthly & daily
    if (!multi_path) {
      #NOTE: the adjustment for two cases (multiple files per sdate or not) has been made in hcst
      Start_fcst_arg_list <- Start_default_arg_list
      Start_fcst_arg_list[['dat']] <- path_list
      Start_fcst_arg_list[['syear']] <- paste0(sdates_fcst)
      fcst <- do.call(Start, Start_fcst_arg_list)
      
      
    } else { # multi_path
      
      #TODO: time attribute is not correct. Improve Start().
      Start_fcst_arg_list <- Start_default_arg_list
      Start_fcst_arg_list[['dat']] <- path_list
      Start_fcst_arg_list[['syear']] <- NULL
      Start_fcst_arg_list[['chunk_depends']] <- NULL
      remove_ind <- which(Start_fcst_arg_list[['return_vars']][['time']] == 'syear')
      Start_fcst_arg_list[['return_vars']][['time']] <- Start_fcst_arg_list[['return_vars']][['time']][-remove_ind]
      fcst <- do.call(Start, Start_fcst_arg_list)
      
      # Reshape and reorder dimensions
      ## dat should be 1, syear should be length of dat; reorder dimensions
      ## dim(fcst) should be [dat, var, syear, time, latitude, longitude, ensemble]
      if (retrieve) {
        dim(fcst) <- c(dat = 1, syear = as.numeric(dim(fcst))[1], dim(fcst)[2:6])
        fcst <- s2dv::Reorder(fcst, c('dat', 'var', 'syear', 'time',
                                      'latitude', 'longitude', 'ensemble'))
      }
      
      # Manipulate time attr because Start() cannot read it correctly
      wrong_time_attr <- attr(fcst, 'Variables')$common$time # dim: [time], the first syear only
      tmp <- array(dim = c(dim(fcst)[c('syear', 'time')]))
      tmp[1, ] <- wrong_time_attr
      yr_diff <- (sdates_fcst - sdates_fcst[1])[-1] #diff(sdates_fcst)
      for (i_syear in 1:length(yr_diff)) {
        tmp[(i_syear + 1), ] <- wrong_time_attr + lubridate::years(yr_diff[i_syear]) 
      }
      attr(fcst, 'Variables')$common$time <- as.POSIXct(tmp, origin = '1970-01-01', tz = 'UTC')
      
    }
    
    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]
    if (retrieve) {
      dim(fcst) <- c(dim(fcst)[1:2], sday = 1, sweek = 1, dim(fcst)[3:7])
      if (!identical(dim(tmp_time_attr), dim(fcst)[c('syear', 'time')])) {
        error(recipe$Run$logger, retrieve = retrieve,
              "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))
    
    #TODO: as.s2dv_cube() needs to be improved to recognize "variable" is under $dat1
    if (multi_path) {
      attributes(fcst)$Variables$common[[variable]] <- attributes(fcst)$Variables$dat1[[variable]]
    }
    
    # Change class from startR_array to s2dv_cube
    if (retrieve) {
      fcst <- as.s2dv_cube(fcst)
      # Only syear could be different
      if (!identical(dim(hcst$data)[-5], dim(fcst$data)[-5])) {
        error(recipe$Run$logger, retrieve = retrieve,
              "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, archive$Reference[[ref.name]]$src,
                        store.freq, "$var$$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 = retrieve)
    
  } 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 = retrieve)
  }
  
  
  #dim(attr(obs, 'Variables')$common$time)
  # sday sweek syear  time 
  #    1     1     2    14 
  
  obs_dates <- attr(obs, 'Variables')$common$time

  if (retrieve) {
    # 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, retrieve = retrieve,
            "obs and hcst dimensions do not match.")
      stop()
    }
    # Add ensemble dim to obs
    dim(obs) <- c(dim(obs), ensemble = 1)
    # Transform to s2dv_cube
    obs <- as.s2dv_cube(obs)
  }
  
  #-------------------------------------------
  # Step 4. Verify the consistance between data
  #-------------------------------------------
  if (retrieve) {
    # dimension
    if (any(!names(dim(obs$data)) %in% names(dim(hcst$data)))) {
      error(recipe$Run$logger, retrieve = retrieve,
            "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, retrieve = retrieve,
              "hcst and obs don't share the same dimension length.")
        stop()
      }
    }
    
    # time attribute
    if (!identical(format(dates, '%Y%m'), format(obs_dates, '%Y%m'))) {
      error(recipe$Run$logger, retrieve = retrieve,
            "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, retrieve = retrieve, 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, retrieve = retrieve, 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, retrieve = retrieve, 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, retrieve = retrieve, 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, retrieve = retrieve, 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, retrieve = retrieve, 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, retrieve = retrieve,
              "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, retrieve = retrieve,
                "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, retrieve = retrieve,
  #          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, retrieve = retrieve,
       "##### DATA LOADING COMPLETED SUCCESSFULLY #####")
  
  .log_memory_usage(recipe$Run$logger, retrieve = retrieve, when = "After loading")
  return(list(hcst = hcst, fcst = fcst, obs = obs))
}
