#*******************************************************************************
# Script to test ShapeToMask
# Eva Rifà Rovira 
#*******************************************************************************
devtools::load_all("/esarchive/scratch/erifarov/git/esviz")
#-------------------------------------------------------------------------------
setwd("/esarchive/scratch/erifarov/rpackages/esviz/shapefile/dev-shapetomask-area/out-testing")

################################################################################
# Shapefiles
colombia_cajamarca <- '/esarchive/scratch/erifarov/rpackages/esviz/shapefile/dev-shapetomask-area/scripts-alba/spatial_aggregation_alba/shp_test/colombia_cajamarca.shp'
colombia_sanvicent <- '/esarchive/scratch/erifarov/rpackages/esviz/shapefile/dev-shapetomask-area/scripts-alba/spatial_aggregation_alba/shp_test/colombia_sanvicentedelcaguan.shp'
NUTS_EU <- paste0('/esarchive/shapefiles/NUTS3/NUTS_RG_60M_2021_4326.shp/', 
                  'NUTS_RG_60M_2021_4326.shp')
# st_read
sf_colombia_cajamarca <- st_read(colombia_cajamarca)
sf_colombia_sanvicent <- st_read(colombia_sanvicent)
sf_NUTS_EU <- st_read(NUTS_EU)

# reference grid: dataset = 'era5land'
lat <- array(seq(6, -1, -0.1), c(lat = 71))
lon <- array(seq(-80, -70, 0.1), c(lon = 101))
ref_grid_era5land_sample <- list(lat = lat, lon = lon)
# Data from era5land
load("/esarchive/scratch/erifarov/rpackages/esviz/shapefile/dev-shapetomask-area/scripts-alba/spatial_aggregation_alba/data/data.era5land.RData")
ref_grid_era5land <- list(latitude = attr(data.era5land, "Variable")$dat$latitude, 
                          longitude = attr(data.era5land, "Variable")$dat$longitude)
################################################################################
# Function to test plots
plot_shp <- function(fileout, x, y = NULL) {
  png(file = fileout)
  plot(x)
  if (!is.null(y)) {
    plot(y, add = TRUE)
  }
  dev.off()
}  
################################################################################
#-----------------------------------------------------
# Tests 1: colombia_cajamarca (only 1 region)
#-----------------------------------------------------
source("/esarchive/scratch/erifarov/git/esviz/R/ShapeToMask.R")
# (1.1) Centroid method

# NOTE: We set shp_system = NULL because we don't want to subset any region
era5land_cajamarca_cntr <- ShapeToMask(shp_file = colombia_cajamarca, 
                                       shp_system = NULL, 
                                       ref_grid = ref_grid_era5land, region = TRUE,
                                       find_min_dist = FALSE)
# (1.2) Area method
# target_crs = NULL
era5land_cajamarca_area <- ShapeToMask(shp_file = colombia_cajamarca, 
                                       ref_grid = ref_grid_era5land, 
                                       shp_system = NULL, 
                                       compute_area_coverage = TRUE, 
                                       region = TRUE)
# Error:
# Error in wk_handle.wk_wkb(wkb, s2_geography_writer(oriented = oriented,  : 
#   Loop 0 is not valid: Edge 1 crosses edge 3

# Solution: Set sf::sf_use_s2(FALSE)
sf::sf_use_s2(FALSE)
era5land_cajamarca_area <- ShapeToMask(shp_file = colombia_cajamarca, 
                                       ref_grid = ref_grid_era5land, 
                                       shp_system = NULL, 
                                       compute_area_coverage = TRUE, 
                                       region = TRUE)
sf::sf_use_s2(TRUE)

# Compare results different solutions
sum(era5land_cajamarca_area)
# [1] 4.140404
sum(era5land_cajamarca_cntr)
[1] 4

# Positions
which(era5land_cajamarca_cntr != 0)
# [1] 1561 1662 1663 1762
which(era5land_cajamarca_area != 0)
# [1] 1460 1560 1561 1562 1661 1662 1663 1762 1763 1863 1864

all(which(era5land_cajamarca_cntr != 0) %in% which(era5land_cajamarca_area != 0))
# [1] TRUE

# Apply the results to the data

mean_extract <- function(data_cube, mask) {
  locations <- which(mask != 0)
  res <- mean(data_cube[locations])
  return(res)
}

aggregated_area <- Apply(data = list(data_cube = data.era5land, mask = era5land_cajamarca_area),
                         target_dims = list(data_cube = c('longitude', 'latitude'),
                         mask = c('longitude', 'latitude')),
                         fun = mean_extract)$output1
aggregated_ctr <- Apply(data = list(data_cube = data.era5land, mask = era5land_cajamarca_cntr),
                         target_dims = list(data_cube = c('longitude', 'latitude'),
                         mask = c('longitude', 'latitude')),
                         fun = mean_extract)$output1
dim(aggregated_area)
# time  syear region 
#   12      1      1 
summary(aggregated_area)
#  Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
# 11.58   11.98   12.24   12.20   12.47   12.63 
summary(aggregated_ctr)
#  Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
# 11.49   11.87   12.12   12.09   12.33   12.50 

# Test (1.3): Raster method with exactextractr
# Source the function
source("/esarchive/scratch/erifarov/rpackages/esviz/shapefile/dev-shapetomask-area/scripts-alba/spatial_aggregation_alba/functions/raster_extract_to_array.R")
raster_era5land_cajamarca <- raster.extract(data = data.era5land, 
                                            shp = sf_colombia_cajamarca)
dim(raster_era5land_cajamarca)
# region    syear     time ensemble 
#     10        1       12        1 

summary(raster_era5land_cajamarca)
#  Min. 1st Qu.  Median    Mean 3rd Qu.    Max. 
# 9.124  11.092  11.788  11.978  13.026  15.151 
################################################################################

#-----------------------------------------------------
# Tests 1: colombia_cajamarca (only 1 region)
#-----------------------------------------------------

# Example ShapeToMask: NUTS
shp_file <- paste0('/esarchive/shapefiles/NUTS3/NUTS_RG_60M_2021_4326.shp/', 
                   'NUTS_RG_60M_2021_4326.shp')
ref_grid <- list(lon = seq(10, 40, 0.5), lat = seq(40, 85, 0.5))
NUTS_name <- list(FI = c('Lappi', 'Kainuu'), SI = c('Pomurska', 'Podravska'))

# Mask computed with area coverage
sf::sf_use_s2(FALSE)
mask_area <- ShapeToMask(shp_file = shp_file, ref_grid = ref_grid, 
                         compute_area_coverage = TRUE, reg_names = NUTS_name, 
                         fileout = "mask_area.nc")
mask_cntr <- ShapeToMask(shp_file = shp_file, ref_grid = ref_grid, 
                         compute_area_coverage = FALSE, reg_names = NUTS_name, 
                         fileout = "mask_cntr.nc")
dim(mask_area)
#  lon    lat region 
#   61     91      4 
summary(mask_area)

# Mask computed with centroid approach
mask_cntr <- ShapeToMask(shp_file = shp_file, ref_grid = ref_grid, region = T, 
                         compute_area_coverage = FALSE, reg_names = NUTS_name)
dim(mask_cntr)
#  lon    lat region 
#   61     91      4 

################################################################################

#*****************************************************************************
# Examples function: ShapeToMask
# Source: https://earth.bsc.es/gitlab/es/esviz/-/blob/main/R/ShapeToMask.R 
# An-Chi & Eva (2023)
#*****************************************************************************
# Exmple (1): NUTS
shp_file <- paste0('/esarchive/shapefiles/NUTS3/NUTS_RG_60M_2021_4326.shp/', 
                   'NUTS_RG_60M_2021_4326.shp')
ref_grid <- paste0('/esarchive/recon/ecmwf/era5land/monthly_mean/',
                   'tas_f1h/tas_201006.nc')
# ref_grid <- list(lon = seq(10, 40, 0.5), lat = seq(40, 85, 0.5))

NUTS.id <- paste0("FI1D", c(1:3, 5, 7:9))  
NUTS.name <- list(FI = c('Lappi', 'Kainuu'), SI = c('Pomurska', 'Podravska'))
mask1 <- ShapeToMask(shp_file, ref_grid, reg_ids = NUTS.id)
mask2 <- ShapeToMask(shp_file = shp_file, ref_grid = ref_grid, 
                     reg_names = NUTS.name)

# Exmple (2): GADM
shp_file <- "/esarchive/shapefiles/gadm_country_mask/gadm_country_ISO3166.shp"
ref_grid <- paste0('/esarchive/exp/ecmwf/s2s-monthly_ensfor/weekly_mean/',
                   'tas_f6h/tas_20191212.nc')
GADM.id <- c("ESP", "ITA")
GADM.name <- c("Spain", "Italy")
mask1 <- ShapeToMask(shp_file = shp_file, ref_grid = ref_grid, 
                     reg_ids = GADM.id, shp_system = "GADM")
mask2 <- ShapeToMask(shp_file = shp_file, ref_grid = ref_grid, 
                     reg_names = GADM.name, shp_system = "GADM")

################################################################################