library(multiApply)
library(startR)
library(s2dv)

mask_tas_tos <- function(input_data, grid, lon, lat, region = region ,
                         lon_dim = 'lon', lat_dim = 'lat', ncores = NULL){
  

  mask <- .load_mask(grid = grid, lon_dim = lon_dim, lat_dim = lat_dim, 
                     sea_value = 1, land_value = 0, region = region)

  
  ## TO DO: improve the check and correct lats
  stopifnot(all(lon == mask$lon))
  stopifnot(max(abs(as.numeric(round(lat,2) - round(mask$lat,2)))) < 0.1) # stopifnot(all(lat == mask$lat))

  tas <- Subset(input_data$data, along = 'var', indices = 1)
  tos <- Subset(input_data$data, along = 'var', indices = 2)
  
  tas_tos <- multiApply::Apply(data = list(tas, tos),
                            target_dims = c(lon_dim, lat_dim),
                            fun = .mask_tas_tos,
                            mask = mask$mask,
                            sea_value = 1,
                            ncores = ncores)$output1
  input_data$data <- tas_tos
 
  return(input_data)
}

.mask_tas_tos <- function(data_tas, data_tos, mask, sea_value){
  data_tas[mask == sea_value] <- data_tos[mask == sea_value]
  return(data_tas)
}

.load_mask <- function(grid, mask_path = NULL, land_value = 0, sea_value = 1,
                       lon_dim = 'lon', lat_dim = 'lat', region = region){
  
  if (is.null(mask_path)){
    mask_sea_land_path <- '/esarchive/exp/ecmwf/system5c3s/constant/lsm/lsm.nc' ## /esarchive/recon/ecmwf/era5land/constant/lsm-r3600x1801cds/lsm.nc'
  } else if (is.character(mask_path)){
   mask_sea_land_path <- mask_path
  } else {
    stop("mask_path must be NULL (to use the default mask and interpolate it to 
         the specified grid) or a string with the mask's path you want to load")
  }
  
  lons.min <- region[1]
  lons.max <- region[2]
  lats.min <- region[3]
  lats.max <- region[4]
  
  ## TO DO: 
  ## Fix region filter for lat and lon
  ## Fix 'number' parameter for mask

  
  data <- startR::Start(dat = mask_sea_land_path, 
                        var = 'lsm',  
                        lon = 'all',
                        lat = 'all',
                        number = 1,               ## needed to add for ensemble member dimension of lsm.nc
                        # lon = values(list(lons.min, lons.max)),
                        # lat = values(list(lats.min, lats.max)),
                        transform = CDORemapper, transform_extra_cells = 2, 
                        transform_params = list(grid = grid, method = 'con', crop = region), 
                        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 = FALSE),
                        lon_reorder = CircularSort(0,359.9),
                        num_procs = 1, retrieve = TRUE)

  mask <- list(mask = drop(data),
               lon = as.numeric(attr(data,'Variables')$common$lon),
               lat = as.numeric(attr(data,'Variables')$common$lat))
  
  mask$mask[data <= 0.5] <- sea_value
  mask$mask[data >  0.5] <- land_value
  names(dim(mask$mask)) <- c(lon_dim, lat_dim)
  
  return(mask)
}
