source("modules/Loading/R/get_regrid_params.R")
source("modules/Loading/R/check_latlon.R")

# This correction can only be applied to temperature.
# This correction is relevant for the mean_bias metric.
# correction_constant is negative 
# -6.5ºK/km(ºC/km for the environmental lapse rate correction
# -9.8ºK/km (ºC/km) for the adiabatic correction
orography_correction <- function(recipe, data, correction_constant = -6.5) {

  ref.name <- recipe$Analysis$Datasets$Reference$name
  exp.name <- recipe$Analysis$Datasets$System$name
  lats.min <- recipe$Analysis$Region$latmin
  lats.max <- recipe$Analysis$Region$latmax
  lons.min <- recipe$Analysis$Region$lonmin
  lons.max <- recipe$Analysis$Region$lonmax

  archive <- get_archive(recipe)

  # Define regrid parameters:
  regrid_params <- get_regrid_params(recipe, archive)

  circularsort <- check_latlon(lats.min, lats.max, lons.min, lons.max)

  ## Load exp orography
  orography_exp <- Start(dat = archive$System[[exp.name]]$orography,
                         var = 'orography',
                         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'),
                         return_vars = list(lat = NULL, lon = NULL),
                         synonims = list(longitude = c('lon','longitude'),
                                         latitude = c('lat','latitude')),
                         num_procs = 1, retrieve = TRUE)
 if (attributes(orography_exp)$Variables$common$oro$unit != 'm') {
    if (attributes(orography_exp)$Variables$common$oro$unit == 'km') {
      orography_exp <- orography_exp * 1000
      attributes(orography_exp)$Variables$common$oro$unit <- 'm'
    } else {
    stop("Exp orography is not m or km, add option to orography_correction.R or fix the files")
    }
 }

  ## Set negative values to zero
  orography_exp <- pmax(orography_exp, 0)
  ## load obs orography
  orography_obs <- Start(dat = archive$Reference[[ref.name]]$orography,
                         var = 'orography',
                         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'),
                         return_vars = list(latitude = NULL, longitude = NULL),
                         synonims = list(longitude = c('lon','longitude'),
                                         latitude = c('lat','latitude')),
                         num_procs = 1, retrieve = TRUE)
  if (attributes(orography_obs)$Variables$common$oro$unit != 'm') {
    if (attributes(orography_obs)$Variables$common$oro$unit == 'km') {
      orography_obs <- orography_obs * 1000
      attributes(orography_obs)$Variables$common$oro$unit <- 'm'
    } else {
    stop("Obs orography is not m or km, add option to orography_correction.R or fix the files")
    }
  }
  # Convert to zero negative heights:
  orography_obs[orography_obs < 0] <- 0
  orography_exp[orography_exp < 0] <- 0 

  ## Apply lapse rate factor to correct temperature = -0.0065 K/m (-6.5 K/km)
  data$hcst$data <- Apply(list(data$hcst$data),
                          target_dims = c('latitude', 'longitude'),
                          fun = function(x, corr_const,
                                         oro_obs, oro_exp) {
                            x + corr_const/1000 *
                                    (oro_obs - oro_exp)},
                          corr_const = correction_constant,
                          oro_obs = drop(orography_obs),
                          oro_exp = drop(orography_exp))$output1
  data$hcst$data <- Reorder(data$hcst$data, names(dim(data$obs$data)))
  if (!is.null(data$fcst$data)) {
     data$fcst$data <- Apply(list(data$fcst$data),
                          target_dims = c('latitude', 'longitude'),
                          fun = function(x, corr_const,
                                         oro_obs, oro_exp) {
                            x + corr_const/1000 *
                                    (oro_obs - oro_exp)},
                          corr_const = correction_constant,
                          oro_obs = drop(orography_obs),
                          oro_exp = drop(orography_exp))$output1
    data$fcst$data <- Reorder(data$fcst$data, names(dim(data$obs$data)))
  }
  return(data)
}

