#!/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()

for projection_type in ['read', 'created']:

    # Open dataset
    if projection_type == 'read':
        test_path = "/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/OUT/stats_bnds/monarch/a45g/regional/daily_max/O3_all/O3_all-000_2021080300.nc"
        nessy = open_netcdf(path=test_path, info=True)

    # Create dataset
    elif projection_type == 'created':
        lat_orig = 41.1
        lon_orig = 1.8
        inc_lat = 0.2
        inc_lon = 0.2
        n_lat = 100
        n_lon = 100
        nessy = 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)
    
    # Add bounds
    nessy.create_spatial_bounds()

    print('NES', projection_type, '-', 'Rank', rank, '-', nessy)
    print('NES', projection_type, '-', 'Rank', rank, '-', 'Lat bounds', 
          nessy.lat_bnds)
    print('NES', projection_type, '-', 'Rank', rank, '-', 'Lon bounds', 
          nessy.lon_bnds)

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