library(startR)

# 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))







# Output verification:

summary(res$output1)
dim(res$output1)
head(res$output1)

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
summary(res$output1[mask_loaded == 1]) # There is no NA's: correct
sum(mask_loaded == 0) # The number of NA's are 840: correct

# 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])
res$output1[1,1,1]
mean(data_loaded[1,1, , , ,1,2])
res$output1[1,1,2]


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
  }
    }    
}

