#!/usr/bin/env python
import sys
from mpi4py import MPI
from nes import *

comm = MPI.COMM_WORLD
rank = comm.Get_rank()
size = comm.Get_size()

# ===== PATH TO MASK ===== #
shapefile_path = '/gpfs/projects/bsc32/models/NES_tutorial_data/timezones_2021c/timezones_2021c.shp'

# ===== FROM EXISTING GRID ===== #
# READ
# Original path: /esarchive/exp/monarch/a4dd/original_files/000/2022111512/MONARCH_d01_2022111512.nc
# Rotated grid from MONARCH
original_path = '/gpfs/projects/bsc32/models/NES_tutorial_data/MONARCH_d01_2022111512.nc'
nessy = open_netcdf(original_path)

# LOAD DATA
nessy.keep_vars('O3')
nessy.load()

# SPATIAL JOIN
# Method can be centroid, nearest and intersection
nessy.spatial_join(shapefile_path, method='centroid')
print('SPATIAL JOIN - Rank: {0}. Shapefile: {1}'.format(rank, nessy.shapefile))

comm.Barrier()
sys.stdout.flush()

# ===== WITH NEW GRID ===== #
# CREATE
# Define projection
projection = 'regular'
lat_orig = 41.1
lon_orig = 1.8
inc_lat = 0.2
inc_lon = 0.2
n_lat = 100
n_lon = 100

# SPATIAL JOIN
# Method can be centroid, nearest and intersection
nessy = from_shapefile(shapefile_path, method='centroid', projection=projection, 
                       lat_orig=lat_orig, lon_orig=lon_orig, 
                       inc_lat=inc_lat, inc_lon=inc_lon, 
                       n_lat=n_lat, n_lon=n_lon)
print('FROM SHAPEFILE - Rank: {0}. Shapefile: {1}'.format(rank, nessy.shapefile))

comm.Barrier()
sys.stdout.flush()
