#################################################
# Load GRIB files from MARS
#################################################

source('modules/Loading/R/GRIB/GrbLoad.R')
source('tools/libs.R')

load_GRIB <- function(recipe) {

  # Set params
  #-------------------------------------------------------------------

  # get recipe info
  hcst.inityear <- recipe$Analysis$Time$hcst_start
  hcst.endyear <- recipe$Analysis$Time$hcst_end
  hcst.sdate <- recipe$Analysis$Time$sdate
  hcst.ftime <-  as.numeric(recipe$Analysis$Time$ftime_min):as.numeric(recipe$Analysis$Time$ftime_max)
  fcst.year <- recipe$Analysis$Time$fcst_year
  lats.min <- recipe$Analysis$Region$latmin  # can only be -90
  lats.max <- recipe$Analysis$Region$latmax  # can only be 90
  lons.min <- recipe$Analysis$Region$lonmin  # can only be 0
  lons.max <- recipe$Analysis$Region$lonmax  # can only be 360
  ref.name <- recipe$Analysis$Datasets$Reference$name
  exp.name <- recipe$Analysis$Datasets$System$name
  variable <- recipe$Analysis$Variables$name  
  store.freq <- recipe$Analysis$Variables$freq

  regrid.method <- recipe$Analysis$Regrid$method
  regrid.type <- recipe$Analysis$Regrid$type

  # get MARS datasets dict:
  archive <- read_yaml("conf/archive.yml")[[recipe$Run$filesystem]]
  exp_descrip <- archive$System[[exp.name]]
  freq.hcst <- unlist(exp_descrip[[store.freq]][variable])
  reference_descrip <- archive$Reference[[ref.name]]
  freq.obs <- unlist(reference_descrip[[store.freq]][variable])
  obs.dir <- reference_descrip$src
  fcst.dir <- exp_descrip$src
  hcst.dir <- exp_descrip$src
  #NOTE: We can use this info in GrbLoad() to substitute param 'has.memb'
  fcst.nmember <- exp_descrip$nmember$fcst
  hcst.nmember <- exp_descrip$nmember$hcst

  info(recipe$Run$logger,
       "========== PARAMETERS RETRIEVED. ==========")

  # Load hindcast 
  #-------------------------------------------------------------------

## original file dir
#exp_path <- "/esarchive/exp/ecmwf/system5_m1/original_files/fcmean_od_sfc_msmm_ecmf/"
## soft link to original file dir
#exp_path <- "/esarchive/scratch/aho/tmp/GRIB/GRIB_system5_tas/"  #files are not correct
# The correct files
#exp_path <- "/esarchive/scratch/aho/tmp/GRIB/GRIB_system5_tas_CORRECTED/"

  hcst.path <- paste0(archive$src, hcst.dir, freq.hcst)
  hcst.year <- paste0(as.numeric(hcst.inityear):as.numeric(hcst.endyear))
  hcst.files <- paste0(hcst.path, variable, '_', hcst.year, hcst.sdate, '.grb')

  if (!regrid.type %in% c('none', 'to_system')) {
    if (regrid.type == 'to_reference') {
      regrid_list <- c(method = regrid.method, type = reference_descrip$reference_grid)
    } else { # e.g., "r360x181"
      regrid_list <- list(method = regrid.method, type = regrid.type)
    }
  } else {
    regrid_list <- NULL
  }

  .log_memory_usage(recipe$Run$logger, when = "Before loading the data")
  hcst <- GrbLoad(dat = as.list(hcst.files), time_step = hcst.ftime, has.memb = hcst.nmember,
                    syear_time_dim = NULL, regrid = regrid_list)
  gc()

  info(recipe$Run$logger,
       "========== HCST LOADED. ==========")

  # Load forecast
  #-------------------------------------------------------------------
  if (!is.null(fcst.year)) {
    fcst.path <- paste0(archive$src, hcst.dir)
    fcst.files <- paste0(fcst.path, variable, '_', fcst.year, hcst.sdate, '.grb')
    fcst <- GrbLoad(dat = as.list(fcst.files), time_step = hcst.ftime, has.memb = fcst.nmember,
                    syear_time_dim = NULL, regrid = regrid_list)
    gc()
  } else {
    fcst <- NULL
  }

  info(recipe$Run$logger,
       "========== FCST LOADED. ==========")

  # Load reference
  #-------------------------------------------------------------------
#obs_path <- "/esarchive/scratch/aho/tmp/GRIB/GRIB_era5_tas/"
  obs.path <- paste0(archive$src, obs.dir, freq.obs, "/")
  # Use hcst time attr to load obs
  hcst_times <- attr(hcst, 'time')
  hcst_times_strings <- format(hcst_times, '%Y%m')

  obs.files <- paste0(obs.path, variable, '_', hcst_times_strings, '.grb')

  if (!regrid.type %in% c('none')) {
    if (regrid.type == 'to_reference') {
      regrid_list <- list(method = regrid.method, type = reference_descrip$reference_grid)
    } else if (regrid.type == 'to_system') {
      regrid_list <- list(method = regrid.method, type = exp_descrip$reference_grid)
    } else { # e.g., "r360x181"
      regrid_list <- list(method = regrid.method, type = regrid.type)
    }
  } else {
    regrid_list <- NULL
  }

  #NOTE: only 1 time step in each obs file
  obs <- GrbLoad(dat = as.list(obs.files), time_step = 1, has.memb = NULL, 
                 syear_time_dim = dim(hcst_times), regrid = regrid_list)
  gc()

  .log_memory_usage(recipe$Run$logger, when = "After loading the data")
  info(recipe$Run$logger,
       "========== OBS LOADED. ==========")


#################################################################################

#dim(hcst)
#    syear      time  latitude longitude  ensemble 
#        4         3       640      1280        51 

##BEFORE TRANSFER TO S2DV_CUBE
#str(hcst)
# num [1:4, 1:3, 1:640, 1:1280, 1:51] 252 252 252 252 251 ...
# - attr(*, "edition")= num 1
# - attr(*, "shortName")= chr "2t"
# - attr(*, "longitude")= num [1:1280] 0 0.281 0.563 0.844 1.125 ...
# - attr(*, "latitude")= num [1:640] 89.8 89.5 89.2 88.9 88.7 ...
# - attr(*, "time")= POSIXct[1:12], format: "2000-12-01" "2001-12-01" ...

#dim(attr(hcst, 'time'))
#syear  time 
#    4     3 

##BEFORE TRANSFER TO S2DV_CUBE
#str(obs)
# num [1:4, 1:3, 1:640, 1:1280] 251 251 251 251 251 ...
# - attr(*, "edition")= num 1
# - attr(*, "shortName")= chr "2t"
# - attr(*, "longitude")= num [1:1280] 0 0.281 0.562 0.844 1.125 ...
# - attr(*, "latitude")= num [1:640] 89.8 89.5 89.2 88.9 88.7 ...
# - attr(*, "time")= POSIXct[1:12], format: "2000-12-01" "2001-12-01" ...

#################################################################################

  info(recipe$Run$logger,
       "========== REGRID DONE. ==========")


  # Turn into s2dv_cube
  #-------------------------------------------------------------------
  # hcst
  metadata_list <- vector("list", length = 1)
  names(metadata_list) <- variable
  metadata_list[[variable]] <- list(long_name = attr(hcst, 'name'),
                                    units = attr(hcst, 'units'))
  load_parameters_list <- list(dat1 = list(file_date = list(paste0(hcst.year, hcst.sdate))))

  hcst <- s2dv_cube(data = array(hcst, dim = dim(hcst)),
                    coords = list(dat = 'dat1',
                                  var = variable,
                                  sday = 1,
                                  sweek = 1,
                                  syear = hcst.year,
                                  time = hcst.ftime,
                                  latitude = attr(hcst, 'latitude'),
                                  longitude = attr(hcst, 'longitude'),
                                  ensemble = 1:hcst.nmember),
                    varName = attr(hcst, 'shortName'),
                    metadata = metadata_list,
                    Dates = attributes(hcst)$time,
                    source_files = hcst.files,
                    load_parameters = load_parameters_list,
                    # extra attrs
                    gribEdition = attr(hcst, 'edition'))

  # fcst
  if (!is.null(fcst)) {
    metadata_list <- vector("list", length = 1)
    names(metadata_list) <- variable
    metadata_list[[variable]] <- list(long_name = attr(fcst, 'name'),
                                    units = attr(fcst, 'units'))
    load_parameters_list <- list(dat1 = list(file_date = list(paste0(fcst.year, hcst.sdate))))

    fcst <- s2dv_cube(data = array(fcst, dim = dim(fcst)),
                      coords = list(dat = 'dat1',
                                    var = variable,
                                    sday = 1,
                                    sweek = 1,
                                    syear = fcst.year,
                                    time = hcst.ftime,
                                    latitude = attr(fcst, 'latitude'),
                                    longitude = attr(fcst, 'longitude'),
                                    ensemble = 1:fcst.nmember),
                      varName = attr(fcst, 'shortName'),
                      metadata = metadata_list,
                      Dates = attributes(fcst)$time,
                      source_files = fcst.files,
                      load_parameters = load_parameters_list,
                      gribEdition = attr(fcst, 'edition'))
  }

  # obs
  metadata_list <- vector("list", length = 1)
  names(metadata_list) <- variable
  metadata_list[[variable]] <- list(long_name = attr(obs, 'name'),
                                    units = attr(obs, 'units'))
  load_parameters_list <- list(dat1 = list(file_date = list(hcst_times_strings)))

  obs <-  s2dv_cube(data = array(obs, dim = dim(obs)),
                    coords = list(dat = 'dat1',
                                  var = variable,
                                  sday = 1,
                                  sweek = 1,
                                  #NOTE: Can we directly use hcst info?
                                  syear = hcst.year,
                                  time = hcst.ftime,
                                  latitude = attr(obs, 'latitude'), 
                                  longitude = attr(obs, 'longitude'),
				  ensemble = 1),
                    varName = attr(obs, 'shortName'),
                    metadata = metadata_list,
                    Dates = attributes(obs)$time,
                    source_files = obs.files,
                    load_parameters = load_parameters_list,
                    gribEdition = attr(obs, 'edition'))


#str(hcst)
#List of 4
# $ data  : num [1, 1, 1, 1, 1:2, 1:2, 1:640, 1:1280, 1:51] 252 253 248 251 251 ...
#  ..- attr(*, "edition")= num 1
#  ..- attr(*, "shortName")= chr "2t"
#  ..- attr(*, "longitude")= num [1:1280] 0 0.281 0.563 0.844 1.125 ...
#  ..- attr(*, "latitude")= num [1:640] 89.8 89.5 89.2 88.9 88.7 ...
#  ..- attr(*, "time")= POSIXct[1:4], format: "2000-12-01" "2001-12-01" ...
# $ dims  : Named int [1:9] 1 1 1 1 2 2 640 1280 51
#  ..- attr(*, "names")= chr [1:9] "dat" "var" "sday" "sweek" ...
# $ coords:List of 9
#  ..$ dat      : chr "dat1"
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ var      : chr "tas"
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ sday     : num 1
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ sweek    : num 1
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ syear    : chr [1:2] "2000" "2001"
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ time     : int [1:2] 1 2
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ latitude : num [1:640] 89.8 89.5 89.2 88.9 88.7 ...
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ longitude: num [1:1280] 0 0.281 0.563 0.844 1.125 ...
#  .. ..- attr(*, "indices")= logi FALSE
#  ..$ ensemble : int [1:51] 1 2 3 4 5 6 7 8 9 10 ...
#  .. ..- attr(*, "indices")= logi FALSE
# $ attrs :List of 4
#  ..$ Dates       : POSIXct[1:4], format: "2000-12-01" "2001-12-01" ...
#  ..$ Variable    :List of 1
#  .. ..$ varName: chr "2t"
#  ..$ source_files: chr [1:2] "/esarchive/scratch/aho/tmp/GRIB/GRIB_system5_tas_CORRECTED/tas_20001101.grb" "/esarchive/scratch/aho/tmp/GRIB/GRIB_system5_tas_CORRECTED/tas_20011101.grb"
#  ..$ gribEdition : num 1
# - attr(*, "class")= chr "s2dv_cube"
  .log_memory_usage(recipe$Run$logger, when = "After regridding")
  info(recipe$Run$logger,
       "##### GRIB DATA LOADED SUCCESSFULLY #####")

  return(list(hcst = hcst, fcst = fcst, obs = obs))

}
