#!/usr/bin/env python
import sys
import timeit
import pandas as pd
from mpi4py import MPI
from nes import *

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

# ===== FILE WITH EXISTING BOUNDS ===== #
# READ
# Original path: /gpfs/scratch/bsc32/bsc32538/HERMESv3/OUT_Complete_single/GFAS_p13h/HERMESv3_GR_GFAS_d01_2022050100.nc
# Rotated grid from HERMES
path_1 = '/gpfs/projects/bsc32/models/NES_tutorial_data/HERMESv3_GR_GFAS_d01_2022050100.nc'
nessy_1 = open_netcdf(path=path_1, info=True)

# EXPLORE BOUNDS
print('FILE WITH EXISTING BOUNDS - Rank', rank, '-', 'Lat bounds', nessy_1.lat_bnds)
print('FILE WITH EXISTING BOUNDS - Rank', rank, '-', 'Lon bounds', nessy_1.lon_bnds)

# WRITE
nessy_1.to_netcdf('bounds_file_1.nc', info=True)

# REOPEN
nessy_2 = open_netcdf('bounds_file_1.nc', info=True)

# LOAD DATA AND EXPLORE BOUNDS
print('FILE WITH EXISTING BOUNDS AFTER WRITE - Rank', rank, '-', 'Lat bounds', nessy_2.lat_bnds)
print('FILE WITH EXISTING BOUNDS AFTER WRITE - Rank', rank, '-', 'Lon bounds', nessy_2.lon_bnds)

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

# ===== FILE WITHOUT EXISTING BOUNDS ===== #
# Original path: /gpfs/scratch/bsc32/bsc32538/mr_multiplyby/OUT/stats_bnds/monarch/a45g/regional/daily_max/O3_all/O3_all-000_2021080300.nc
# Rotated grid from MONARCH
path_3 = "/gpfs/projects/bsc32/models/NES_tutorial_data/O3_all-000_2021080300.nc"
nessy_3 = open_netcdf(path=path_3, info=True)

# CREATE BOUNDS
nessy_3.create_spatial_bounds()

# EXPLORE BOUNDS
print('FILE WITHOUT EXISTING BOUNDS - Rank', rank, '-', 'Lat bounds', nessy_3.lat_bnds)
print('FILE WITHOUT EXISTING BOUNDS - Rank', rank, '-', 'Lon bounds', nessy_3.lon_bnds)

# WRITE
nessy_3.to_netcdf('bounds_file_2.nc', info=True)

# REOPEN
nessy_4 = open_netcdf('bounds_file_2.nc', info=True)

# LOAD DATA AND EXPLORE BOUNDS
print('FILE WITH EXISTING BOUNDS AFTER WRITE - Rank', rank, '-', 'Lat bounds', nessy_4.lat_bnds)
print('FILE WITH EXISTING BOUNDS AFTER WRITE - Rank', rank, '-', 'Lon bounds', nessy_4.lon_bnds)

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

# ===== FROM NEW GRID ===== #
# CREATE GRID
lat_orig = 41.1
lon_orig = 1.8
inc_lat = 0.2
inc_lon = 0.2
n_lat = 100
n_lon = 100
nessy_5 = create_nes(comm=None, info=True, projection='regular',
                     lat_orig=lat_orig, lon_orig=lon_orig, 
                     inc_lat=inc_lat, inc_lon=inc_lon, 
                     n_lat=n_lat, n_lon=n_lon)

# CREATE BOUNDS
nessy_5.create_spatial_bounds()

# EXPLORE BOUNDS
print('FROM NEW GRID - Rank', rank, '-', 'Lat bounds', nessy_5.lat_bnds)
print('FROM NEW GRID - Rank', rank, '-', 'Lon bounds', nessy_5.lon_bnds)

# WRITE
nessy_5.to_netcdf('bounds_file_3.nc', info=True)

# REOPEN
nessy_6 = open_netcdf('bounds_file_3.nc', info=True)

# LOAD DATA AND EXPLORE BOUNDS
print('FROM NEW GRID AFTER WRITE - Rank', rank, '-', 'Lat bounds', nessy_6.lat_bnds)
print('FROM NEW GRID AFTER WRITE - Rank', rank, '-', 'Lon bounds', nessy_6.lon_bnds)

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