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

mask_tas_tos <- function(input_data, mask_path, lsm_var_name = lsm_var_name, lon, lat, region = region,
                         lon_dim = 'lon', lat_dim = 'lat', 
                         ncores = NULL){
  
  
  mask <- .load_mask(mask_path = mask_path, lsm_var_name = lsm_var_name, 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)
  
  input_data$data <- 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
  
  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(mask_path = mask_path, lsm_var_name = lsm_var_name, land_value = 0, sea_value = 1,
                       lon_dim = 'lon', lat_dim = 'lat', region = region){
  
  lons.min <- region[1]
  lons.max <- region[2]
  lats.min <- region[3]
  lats.max <- region[4]

  data <- startR::Start(dat = mask_path, 
                        var =  lsm_var_name, 
                        lon = values(list(lons.min, lons.max)),
                        lat = values(list(lats.min, lats.max)),
                        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,360),
                        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 >= 50] <- sea_value
  mask$mask[data < 50] <- land_value
  names(dim(mask$mask)) <- c(lon_dim, lat_dim) 
  
  return(mask)
}
