## TODO: remove paths to personal scratchs 
source("/esarchive/scratch/vagudets/repos/csoperational/R/get_regrid_params.R")


adiabatic_correction <- function(recipe) {

  recipe$Analysis$Regrid$type = 'none'
  ## Load exp and obs data
  source("modules/Loading/R/load_seasonal.R")
  data <- load_seasonal(recipe)
  
  ref.name <- recipe$Analysis$Datasets$Reference$name
  exp.name <- recipe$Analysis$Datasets$System$name
  
  archive <- read_yaml("conf/archive.yml")$esarchive
  
  # Define regrid parameters:
  regrid_params <- get_regrid_params(recipe, archive)
  
  ## load orographies in requested grid
  
  ## Exp orography
  orography_exp <- Start(dat = archive$System[[exp.name]]$orography, 
                         var = 'orography',  
                         lon = 'all',
                         lat = 'all',
                         # transform = regrid_params$fcst.transform,
                         # transform_params = list(grid = regrid_params$fcst.gridtype,
                         #                         method = regrid_params$fcst.gridmethod), 
                         # transform_vars = c('lat','lon'),
                         return_vars = list(lat = NULL, lon = NULL), 
                         synonims = list(lon = c('lon','longitude'),
                                         lat = c('lat','latitude')),
                         lat_reorder = Sort(decreasing = TRUE),
                         lon_reorder = CircularSort(0,359.9),
                         num_procs = 1, retrieve = TRUE)
  
  ## Obs orography
  orography_obs <- Start(dat = archive$Reference[[ref.name]]$orography, 
                         var = 'orography',  
                         lon = 'all',
                         lat = 'all',
                         # transform = regrid_params$obs.transform,
                         # transform_params = list(grid = regrid_params$obs.gridtype,
                         #                         method = regrid_params$obs.gridmethod), 
                         # transform_vars = c('lat','lon'),
                         return_vars = list(lat = NULL, lon = NULL), 
                         synonims = list(lon = c('lon','longitude'),
                                         lat = c('lat','latitude')),
                         lat_reorder = Sort(decreasing = TRUE),
                         lon_reorder = CircularSort(0,359.9),
                         num_procs = 1, retrieve = TRUE)
  
  ## Apply lapse rate factor to correct temperature = -0.0065 K/m (-6.5 K/km)
  
  orography_exp_corr <- drop(orography_exp) * 0.0065 
  orography_exp_corr <- Reorder(data = orography_exp_corr, order = c('lat', 'lon'))
  
  orography_obs_corr <- drop(orography_obs) * 0.0065 
  orography_obs_corr <- Reorder(data = orography_obs_corr, order = c('lat', 'lon'))
  
  ## Apply correction to exp temperature data
  for(ens in 1:dim(data$hcst$data)['ensemble']){
    for(time in 1:dim(data$hcst$data)['time']){
      for(syear in 1:dim(data$hcst$data)['syear']){
        data$hcst$data[,,,,syear,time,,,ens] <- data$hcst$data[,,,,syear,time,,,ens] + orography_exp_corr
        
      }
    }  
  }

  ## Apply correction to obs temperature data
  for(ens in 1:dim(data$obs$data)['ensemble']){
    for(time in 1:dim(data$obs$data)['time']){
      for(syear in 1:dim(data$obs$data)['syear']){
        data$obs$data[,,,,syear,time,,,ens] <- data$obs$data[,,,,syear,time,,,ens] + orography_obs_corr
        
      }
    }  
  }
  
  # Regrid data
  #------------------------------------------------------------------- 

  # Regrid reference to system grid:
  if(recipe$Analysis$Regrid$type == 'to_system'){
    
    aux <- CDORemap(data_array = data$obs$data,
                    lons = data$obs$coords$longitude,
                    lats = data$obs$coords$latitude, 
                    grid = "/esarchive/exp/ecmwf/system5c3s/monthly_mean/tas_f6h/tas_20180501.nc", #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)
    
    data$obs$data <- aux$data_array   
    data$obs$coords$longitude <- aux$lons
    data$obs$coords$latitude <- aux$lats
    data$obs$dims['longitude'] <- dim(aux$data_array)['longitude']
    data$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')) {
    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.")
      
    }
  }
  
  info(recipe$Run$logger,
       "##### DATA LOADING COMPLETED SUCCESSFULLY #####")
  
  return(list(hcst = hcst, fcst = fcst, obs = obs))
}

############## METHOD 2 - REGRID OROGRAPHY #############


adiabatic_correction <- function(recipe) {
  
  ## Load exp and obs data
  source("modules/Loading/R/load_seasonal.R")
  data <- load_seasonal(recipe)
  
  ref.name <- recipe$Analysis$Datasets$Reference$name
  exp.name <- recipe$Analysis$Datasets$System$name
  
  archive <- read_yaml("conf/archive.yml")$esarchive
  
  # Define regrid parameters:
  regrid_params <- get_regrid_params(recipe, archive)
  
  ## load orographies in requested grid
  
  ## Exp orography
  orography_exp <- Start(dat = archive$System[[exp.name]]$orography, 
                         var = 'orography',  
                         lon = 'all',
                         lat = 'all',
                         transform = regrid_params$fcst.transform,
                         transform_params = list(grid = regrid_params$fcst.gridtype,
                                                 method = regrid_params$fcst.gridmethod),
                         transform_vars = c('lat','lon'),
                         return_vars = list(lat = NULL, lon = NULL), 
                         synonims = list(lon = c('lon','longitude'),
                                         lat = c('lat','latitude')),
                         lat_reorder = Sort(decreasing = TRUE),
                         lon_reorder = CircularSort(0,359.9),
                         num_procs = 1, retrieve = TRUE)
  
  ## Obs orography
  orography_obs <- Start(dat = archive$Reference[[ref.name]]$orography, 
                         var = 'orography',  
                         lon = 'all',
                         lat = 'all',
                         transform = regrid_params$obs.transform,
                         transform_params = list(grid = regrid_params$obs.gridtype,
                                                 method = regrid_params$obs.gridmethod),
                         transform_vars = c('lat','lon'),
                         return_vars = list(lat = NULL, lon = NULL), 
                         synonims = list(lon = c('lon','longitude'),
                                         lat = c('lat','latitude')),
                         lat_reorder = Sort(decreasing = TRUE),
                         lon_reorder = CircularSort(0,359.9),
                         num_procs = 1, retrieve = TRUE)
  
  
  ## Calculate difference between orographies
   oro_diff <- orography_exp - orography_obs
  
  ## Apply lapse rate factor to correct temperature = -0.0065 K/m (-6.5 K/km)
  
  oro_obs_corr <- oro_diff * -0.0065
  oro_obs_corr <- Reorder(data = drop(oro_obs_corr), order = c('lat', 'lon'))
  
  ## Apply correction to obs temperature data
  for(ens in 1:dim(data$obs$data)['ensemble']){
    for(time in 1:dim(data$obs$data)['time']){
      for(syear in 1:dim(data$obs$data)['syear']){
        data$obs$data[,,,,syear,time,,,ens] <- data$obs$data[,,,,syear,time,,,ens] + oro_obs_corr
        
      }
    }  
  }
  

}

