# Use a mask for applying different methods in different gridpoints
# Author: Núria Pérez-Zanón
# ------------------------------------------------------------------
# This use case has two parts:
# A) create a mask 
# B) the startR workflow using the mask
# Notice that you can skip step A if you already have a mask, for instance, the Sea-Land Mask of a model, and you want to apply different analysis in the land than in the ocean.
# Notice too that this strategy allows you to chunk in the spatial dimensions.

library(startR)

#####################################################################
# Step A)

# This is an example using toy data. 
#I read one file of system5_m1 to get the latitudes and longitudes in a region, to create a consistent mask:
repos <- '/esarchive/exp/ecmwf/system5_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc'
data <- Start(dat = repos,
              var = 'tas',
              sdate = '20170101',
              ensemble = indices(1),
              time = 'first',
              latitude = indices(300:341),
              longitude = indices(1:40),
              return_vars = list(latitude = 'dat',
                                 longitude = 'dat'),
              retrieve = FALSE)

lats <- attributes(data)$Variables$dat1$longitude 
lons <- attributes(data)$Variables$dat1$latitude

# I create a fake mask with one's and zero's:
mask <- rep(c(rep(1, 21), rep(0, 21)), length(lats))
dim(mask) <- c(latitude = length(lats), longitude = length(lons))

attr(mask, 'longitude') <- lons
attr(mask, 'latitude') <- lats

# Save the mask in a ncdf file:
library(ncdf4)
dims_var <- NULL
dim_lon <- ncdim_def(name = 'longitude', units = 'degrees',
                     vals = as.vector(attributes(data)$Variables$dat1$longitude),
                     longname = 'longitude')
dims_var[[1]] <- dim_lon
dim_lat <- ncdim_def(name = 'latitude', units = 'degrees_north',
                     vals = as.vector(attributes(data)$Variables$dat1$latitude),
                     longname = 'latitude')
dims_var[[2]] <- dim_lat
datanc  <- ncvar_def(name = 'mask',
                     units = 'adimensional',
                     dim = dims_var, missval = -99999)
file_nc <- nc_create("/esarchive/scratch/nperez/mask_system5_m1_harvestmonth.nc",
                     datanc)
ncvar_put(file_nc, datanc, mask)
nc_close(file_nc)

#####################################################################
# --------------------------------------------------------------------------------
#####################################################################
# Step B)

# Read data:
repos <- '/esarchive/exp/ecmwf/system5_m1/monthly_mean/$var$_f6h/$var$_$sdate$.nc'
data <- Start(dat = repos,
              var = 'tas',
              sdate = c('20170101', '20180101'),
              ensemble = indices(1:20),
              time = 'all',
              latitude = indices(300:341),
              longitude = indices(1:40),
              return_vars = list(latitude = 'dat',
                                 longitude = 'dat',
                                 time = 'sdate'),
              retrieve = FALSE)

# Read mask:
path <- '/esarchive/scratch/nperez/$var$_system5_m1_harvestmonth.nc'
mask <-  Start(dat = path,
              var = 'mask',
              latitude = 'all',
              longitude = 'all',
              return_vars = list(latitude = 'dat',
                                 longitude = 'dat'),
              retrieve = FALSE)

# The function does the mean if the mask is 1 or write an NA
MeanMask <- function(x, mask) {
  if (mask == 1) {
    ind <- mean(x)
  } else {
    ind <-  NA
  }
  return(ind)
}

# It is necessary to specify one dimension for mask,
# that's why I have added 'dat' dimension:
stepMask <- Step(fun = MeanMask,
                target_dims = list(x = c('dat', 'ensemble', 'sdate', 'time'),
                                   mask = c('dat')),
                output_dim = NULL)

# Now, there are two data inputs:
wf_mask <- AddStep(list(data, mask), stepMask)

res <- Compute(workflow = wf_mask,
                chunks = list(latitude = 2,
                              longitude = 2))

#####################################################################
# --------------------------------------------------------------------------------
#####################################################################
# Extra lines for output verification:
# Output verification:
dim(res$output1)
#      var  latitude longitude 
#        1        42        40 
summary(res$output1)
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max.    NA's 
#  295.4   300.0   300.2   300.0   300.5   301.0     840 
head(res$output1)
#[1] 300.0644       NA       NA       NA       NA       NA

mask_loaded <-  Start(dat = path,
                      var = 'mask',
                      latitude = 'all',
                      longitude = 'all',
                      return_vars = list(latitude = 'dat',
                                         longitude = 'dat'),
                      retrieve = TRUE)
summary(res$output1[mask_loaded == 0]) # All are NA's: correct
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max.    NA's 
#     NA      NA      NA     NaN      NA      NA     840 
summary(res$output1[mask_loaded == 1]) # There is no NA's: correct
#   Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
#  295.4   300.0   300.2   300.0   300.5   301.0 
sum(mask_loaded == 0) # The number of NA's are 840: correct
#[1] 840

# compare values:
data_loaded <- Start(dat = repos,
                     var = 'tas',
                     sdate = c('20170101', '20180101'),
                     ensemble = indices(1:20),
                     time = 'all',
                     latitude = indices(300:341),
                     longitude = indices(1:40),
                     return_vars = list(latitude = 'dat',
                                        longitude = 'dat',
                                        time = 'sdate'),
                     retrieve = TRUE)
mean(data_loaded[1,1, , , ,1,1])
#[1] 300.0644
res$output1[1,1,1]
#[1] 300.0644
mean(data_loaded[1,1, , , ,1,2])
#[1] 300.3169
res$output1[1,1,2]
#[1] 300.3169

out <- mask_loaded
for (i in 1:(dim(data_loaded)['latitude'])) {
    for (j in 1:(dim(data_loaded)['longitude'])) {
  if (mask_loaded[1,1, i, j] == 1) {
     out[1,1,i, j] <- mean(data_loaded[,,,,,i,j])
  } else {
     out[1,1,i,j] <- NA
  }
    }
}


