## TODO: remove paths to personal scratchs 
source("/esarchive/scratch/vagudets/repos/csoperational/R/get_regrid_params.R")
# Load required libraries/funs
source("modules/Loading/R/dates2load.R")
source("modules/Loading/R/get_timeidx.R")
source("modules/Loading/R/check_latlon.R")
## TODO: Move to prepare_outputs.R
source("tools/libs.R")
## TODO: remove these two lines when new as.s2dv_cube() is in CSTools
source('https://earth.bsc.es/gitlab/external/cstools/-/raw/develop-new_s2dv_cube/R/as.s2dv_cube.R')
source('https://earth.bsc.es/gitlab/external/cstools/-/raw/develop-new_s2dv_cube/R/zzz.R')

## TODO: Source new s2dv_cube version
## TODO: Eliminate dim_var dimension (merge_across_dims?)

load_datasets <- 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]][1]
  vars <- 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 esarchive datasets dict:
  ## TODO: Adapt to 'filesystem' option in recipe
  archive <- read_yaml("conf/archive.yml")$esarchive
  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
  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
  ##}

  var_dir_obs <- reference_descrip[[store.freq]][vars]
  var_dir_exp <- exp_descrip[[store.freq]][vars]
  
  # -----------
  obs.path <- paste0(archive$src,
                     obs.dir, store.freq, "/$var$", "$var_dir$",
		     "/$var$_$file_date$.nc")
  
  hcst.path <- paste0(archive$src,
                      hcst.dir, store.freq, "/$var$", "$var_dir$",
		      "$var$_$file_date$.nc")
  
  fcst.path <- paste0(archive$src,
                      hcst.dir, store.freq, "/$var$", "$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 = vars,
		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')),
                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 (recipe$Analysis$Variables$freq == "daily_mean") {
    # 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)
  }
  
  ## Combine tas and tos data into one variable: tas-tos
  if(recipe$Analysis$Variables$name == 'tas tos'){
    #if(recipe$Analysis$Datasets$Reference$name == 'HadCRUT5' || recipe$Analysis$Datasets$Reference$name == 'BEST') {
      source('/esarchive/scratch/nmilders/gitlab/git_clones/s2s-suite/modules/Loading/R/mask_tas_tos.R')
      hcst <- mask_tas_tos(input_data = hcst, region = c(lons.min, lons.max,lats.min, lats.max),
                           grid = 'r360x181', 
                           lon = hcst$coords$longitude,
                           lat = hcst$coords$latitude,
                           lon_dim = 'longitude', lat_dim = 'latitude', ncores = NULL)
      
      hcst$dims[['var']] <- dim(hcst$data)[['var']]
    #}
  }

  # 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 = vars,
		  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')),
    	          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 (recipe$Analysis$Variables$freq == "daily_mean") {
      # 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)
    
    ## tas tos mask
    if (recipe$Analysis$Variables$name == 'tas tos'){
      if (recipe$Analysis$Datasets$Reference$name == 'HadCRUT5'){
        vars <- 'tasanomaly' 
        var_dir_obs <- reference_descrip[[store.freq]][vars]
      }
    }
    
    if (recipe$Analysis$Variables$name == 'tas tos'){
      if (recipe$Analysis$Datasets$Reference$name == 'BEST'){
        vars <- 'tas' 
        var_dir_obs <- reference_descrip[[store.freq]][vars]
      }
    }
    
    obs <- Start(dat = obs.path,
                 var = vars,
		 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 == "daily_mean") {

    # 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 = vars,
		 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)
  
  ## Combine tas and tos data into one variable: tas-tos
  if(recipe$Analysis$Variables$name == 'tas tos'){
    if(recipe$Analysis$Datasets$Reference$name != 'HadCRUT5' & recipe$Analysis$Datasets$Reference$name != 'BEST'){
      source('/esarchive/scratch/nmilders/gitlab/git_clones/s2s-suite/modules/Loading/R/mask_tas_tos.R')
      obs <- mask_tas_tos(input_data = obs, region = c(lons.min, lons.max,lats.min, lats.max),
                           grid = 'r360x181', 
                           lon = obs$coords$longitude,
                           lat = obs$coords$latitude,
                           lon_dim = 'longitude', lat_dim = 'latitude', ncores = NULL)
      
      obs$dims[['var']] <- dim(obs$data)[['var']]
    }
  }

  # Check for consistency between hcst and obs grid
  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.")

    }
  }

  # Remove negative values in accumulative variables
  dictionary <- read_yaml("conf/variable-dictionary.yml")
  for (var_idx in 1:length(vars)) {
    var_name <- vars[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
      }
    }

    # Convert prlr from m/s to mm/day
    ## TODO: Make a unit conversion function
    if (vars[[var_idx]] == "prlr") {
      # Verify that the units are m/s and the same in obs and hcst
      if (((obs$attrs$Variable$metadata[[var_name]]$units == "m s-1") ||
	   (obs$attrs$Variable$metadata[[var_name]]$units == "m s**-1")) &&
	  ((hcst$attrs$Variable$metadata[[var_name]]$units == "m s-1") ||
	   (hcst$attrs$Variable$metadata[[var_name]]$units == "m s**-1"))) {
        info(recipe$Run$logger, "Converting precipitation from m/s to mm/day.")
        obs$data[, var_idx, , , , , , , ] <-
	  obs$data[, var_idx, , , , , , , ]*86400*1000
	obs$attrs$Variable$metadata[[var_name]]$units <- "mm/day"
        hcst$data[, var_idx, , , , , , , ] <-
	  hcst$data[, var_idx, , , , , , , ]*86400*1000
	hcst$attrs$Variable$metadata[[var_name]]$units <- "mm/day"
        if (!is.null(fcst)) {
          fcst$data[, var_idx, , , , , , , ] <- 
	    fcst$data[, var_idx, , , , , , , ]*86400*1000
	  fcst$attrs$Variable$metadata[[var_name]]$units <- "mm/day"
        }
      }
    }
  }
  # Compute anomalies if requested
  # 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 #####")
  
  ############################################################################
  #
  #     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 ##############")
  
  ############################################################################
  ############################################################################
 
  return(list(hcst = hcst, fcst = fcst, obs = obs))

}
