# 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/mask_tas_tos.R')
source("modules/Loading/R/compare_exp_obs_grids.R")

load_tas_tos <- function(recipe, retrieve = TRUE) {
  
  # -------------------------------------------
  # 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
  ncores <- recipe$Analysis$ncores

  variable <- c("tas", "tos", "sic")
  store.freq <- recipe$Analysis$Variables$freq
  
  if(is.null(recipe$Analysis$Variables$sic_threshold)){
    sic.threshold <- 0.15
  } else {
    sic.threshold <- recipe$Analysis$Variables$sic_threshold
  }
  
  data_order <- c('dat', 'var', 'sday', 'sweek', 'syear', 'time', 'latitude', 'longitude', 'ensemble')
  
  
  # get sdates array
  sdates <- dates2load(recipe)
  
  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)
  }
  
  # get esarchive datasets dict:
  archive <- get_archive(recipe)
  exp_descrip <- archive$System[[exp.name]]
  
  reference_descrip <- archive$Reference[[ref.name]]
  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
 
  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 data without regrid
  #-------------------------------------------------------------------
  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(decreasing = TRUE),
                longitude = values(list(lons.min, lons.max)),
                longitude_reorder = circularsort,
                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 (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
  }
  
  # Define sea-ice grid points based of sea-ice concentration threshold
  ice_hcst <- hcst[,3,,,,,,,] >= sic.threshold
  
  # Replace Tos with Tas for data points with sea ice
  hcst[,2,,,,,,,][ice_hcst] <- hcst[,1,,,,,,,][ice_hcst] 
  
  
  # 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 hcst tas and tos data
  #-------------------------------------------------------------------

  hcst <- mask_tas_tos(input_data = hcst, 
                       region = c(lons.min, lons.max,lats.min, lats.max),
                       mask_path = archive$System[[exp.name]]$land_sea_mask,
                       lsm_var_name = 'lsm',
                       lon = as.vector(hcst$coords$longitude), 
                       lat = as.vector(hcst$coords$latitude),
                       lon_dim = 'longitude', 
                       lat_dim = 'latitude', 
                       ncores = ncores)
  
  hcst$dims[['var']] <- dim(hcst$data)[['var']]
  hcst$attrs$Variable$varName <- 'tas-tos'

  hcst$data <- Reorder(hcst$data, data_order)
    
  # Load forecast data without regrid
  #-------------------------------------------------------------------
  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(decreasing = TRUE),
                  longitude = values(list(lons.min, lons.max)),
                  longitude_reorder = circularsort,
                  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 (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
    }
    
    # Define sea-ice grid points based of sea-ice concentration threshold
    ice_fcst <- fcst[,3,,,,,,,] >= sic.threshold 
    
    # Replace Tos with Tas for datapoints with sea ice
    fcst[,2,,,,,,,][ice_fcst] <- fcst[,1,,,,,,,][ice_fcst] 
    
    
    # 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)
    }
    
    # Combine fcst tas and tos data
    #-------------------------------------------------------------------
    
    fcst <- mask_tas_tos(input_data = fcst, region = c(lons.min, lons.max,lats.min, lats.max),
                         mask_path = archive$System[[exp.name]]$land_sea_mask, lsm_var_name = 'lsm',
                         lon = fcst$coords$longitude, lat = fcst$coords$latitude,
                         lon_dim = 'longitude', lat_dim = 'latitude', ncores = NULL)
    
    fcst$dims[['var']] <- dim(fcst$data)[['var']]
    fcst$attrs$Variable$varName <- 'tas-tos'
    
    fcst$data <- Reorder(fcst$data, data_order)
    
  } else {
    fcst <- NULL
  }
  
  # Load obs data without regrid
  #-------------------------------------------------------------------
  
  # 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 (frequency == "monthly_mean") {
    
    dates_file <- format(as.Date(dates, '%Y%m%d'), "%Y%m")
    dim(dates_file) <- dim(dates)
    
    
    # Define variables for blended tas-tos datasets
    if (recipe$Analysis$Datasets$Reference$name == 'BEST'){
      variable <- 'tas' 
      var_dir_obs <- reference_descrip[[frequency]][variable]
    }
    
    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(decreasing = TRUE),
                 longitude = values(list(lons.min, lons.max)),
                 longitude_reorder = circularsort,
                 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 (frequency == "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 = 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(decreasing = TRUE),
                 longitude = values(list(lons.min, lons.max)),
                 longitude_reorder = circularsort,
                 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
  
  if(!recipe$Analysis$Datasets$Reference$name %in% c('HadCRUT4','HadCRUT5','BEST','GISTEMPv4')){
    
    # Define sea-ice grid points based of sea-ice concentration threshold
    ice_obs <- (obs[,3,,,,,,,]) >= sic.threshold
    
    # Replace NA values with False
    ice_obs[is.na(ice_obs)] <- FALSE

    # Replace Tos with Tas for datapoints with sea ice
    obs[,2,,,,,,,][ice_obs] <- obs[,1,,,,,,,][ice_obs] 
  }
  
  # Convert obs to s2dv_cube
  obs <- as.s2dv_cube(obs)
  
  
  # Combine obs tas and tos data
  #-------------------------------------------------------------------
  ## TODO: Ask about this list 
  if(!recipe$Analysis$Datasets$Reference$name %in% c('HadCRUT4','HadCRUT5','BEST','GISTEMPv4')){
      
    obs <- mask_tas_tos(input_data = obs, 
                        region = c(lons.min, lons.max,lats.min, lats.max),
                        mask_path = archive$Reference[[ref.name]]$land_sea_mask, 
                        lsm_var_name = 'sftof',
                        lon = as.vector(obs$coords$longitude), 
                        lat = as.vector(obs$coords$latitude),
                        lon_dim = 'longitude', lat_dim = 'latitude', ncores = NULL)

    obs$dims[['var']] <- dim(obs$data)[['var']]
    obs$attrs$Variable$varName <- 'tas-tos'

  } ## close if on reference name
  
  obs$data <- Reorder(obs$data, data_order)
  
  
  # Regrid data
  #------------------------------------------------------------------- 
  
  # Regrid reference to system grid:
  if(recipe$Analysis$Regrid$type == 'to_system'){
    
    aux <- CDORemap(data_array = obs$data, ## Not regridding to desired grid when latitudes are ordered descending 
                    lons = obs$coords$longitude,
                    lats = obs$coords$latitude, 
                    grid = regrid_params$obs.gridtype,
                    method = recipe$Analysis$Regrid$method, 
                    avoid_writes = TRUE,
                    crop = c(lons.min, lons.max,lats.min, lats.max),
                    force_remap = TRUE)
    
    obs$data <- aux$data_array   
    obs$coords$longitude <- aux$lons
    obs$coords$latitude <- aux$lats
    obs$dims['longitude'] <- dim(aux$data_array)['longitude']
    obs$dims['latitude'] <- dim(aux$data_array)['latitude']
    rm(aux)
  }
  
  # Regrid system to reference grid:
  if(recipe$Analysis$Regrid$type == 'to_reference'){
    
    aux <- CDORemap(data_array = hcst$data, 
                    lons = hcst$coords$longitude, lats = hcst$coords$latitude, 
                    grid = regrid_params$fcst.gridtype, method = recipe$Analysis$Regrid$method, 
                    avoid_writes = TRUE, crop = TRUE,
                    force_remap = TRUE)
    
    hcst$data <- aux$data_array
    hcst$coords$longitude <- aux$lons
    hcst$coords$latitude <- aux$lats
    hcst$dims['longitude'] <- dim(aux$data_array)['longitude']
    hcst$dims['latitude'] <- dim(aux$data_array)['latitude']
    rm(aux)
    
    if (!is.null(recipe$Analysis$Time$fcst_year)) {
      aux <- CDORemap(data_array = fcst$data, 
                      lons = fcst$coords$longitude,
                      lats = fcst$coords$latitude,
                      grid = regrid_params$fcst.gridtype,
                      method = recipe$Analysis$Regrid$method,
                      avoid_writes = TRUE, crop = TRUE,
                      force_remap = TRUE)
      
      fcst$data <- aux$data_array
      fcst$coords$longitude <- aux$lons
      fcst$coords$latitude <- aux$lats
      fcst$dims['longitude'] <- dim(aux$data_array)['longitude']
      fcst$dims['latitude'] <- dim(aux$data_array)['latitude']
      rm(aux)
    }
  }
  
  # Regrid all data to user defined grid:
  if(!recipe$Analysis$Regrid$type %in% c('to_system','to_reference')){
  
    aux <- CDORemap(data_array = hcst$data, 
                    lons = hcst$coords$longitude, lats = hcst$coords$latitude, 
                    grid = regrid_params$fcst.gridtype, method = recipe$Analysis$Regrid$method, 
                    avoid_writes = TRUE, crop = TRUE,
                    force_remap = TRUE)
    
    hcst$data <- aux$data_array
    hcst$coords$longitude <- aux$lons
    hcst$coords$latitude <- aux$lats
    hcst$dims['longitude'] <- dim(aux$data_array)['longitude']
    hcst$dims['latitude'] <- dim(aux$data_array)['latitude']
    rm(aux)
    
    if (!is.null(recipe$Analysis$Time$fcst_year)) {
      aux <- CDORemap(data_array = fcst$data, 
                      lons = fcst$coords$longitude,
                      lats = fcst$coords$latitude, 
                      grid = regrid_params$fcst.gridtype,
                      method = recipe$Analysis$Regrid$method, 
                      avoid_writes = TRUE, crop = TRUE,
                      force_remap = TRUE)
      
      fcst$data <- aux$data_array
      fcst$coords$longitude <- aux$lons
      fcst$coords$latitude <- aux$lats
      fcst$dims['longitude'] <- dim(aux$data_array)['longitude']
      fcst$dims['latitude'] <- dim(aux$data_array)['latitude']
      rm(aux)
    }
    
    aux <- CDORemap(data_array = obs$data, 
                    lons = obs$coords$longitude, 
                    lats = obs$coords$latitude, 
                    grid = regrid_params$obs.gridtype,
                    method = recipe$Analysis$Regrid$method, 
                    avoid_writes = TRUE, crop = TRUE,
                    force_remap = TRUE)
    
    obs$data <- aux$data_array
    obs$coords$longitude <- aux$lons
    obs$coords$latitude <- aux$lats
    obs$dims['longitude'] <- dim(aux$data_array)['longitude']
    obs$dims['latitude'] <- dim(aux$data_array)['latitude']
    rm(aux)
  }
  
  
  # Check for consistency between hcst and obs grid
  if (!(recipe$Analysis$Regrid$type == 'none')) {
    compare_exp_obs_grids(exp = hcst, obs = obs)
  }
  
  info(recipe$Run$logger, retrieve = retrieve,
       "##### DATA LOADING COMPLETED SUCCESSFULLY #####")
  
  return(list(hcst = hcst, fcst = fcst, obs = obs))
}
