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

# Hide warning
pd.options.mode.chained_assignment = None

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

results = []
for method in ['spatial_overlay', 'spatial_join']:
    for projection in ['regular', 'rotated']:
        for projection_type in ['created', 'read']:

            # Regular projection
            if projection == 'regular':
                # Create dataset and get shapefile
                if 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
                    coordinates = 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)
                    coordinates.create_shapefile()

                # Open dataset and get shapefile
                elif projection_type == 'read':
                    coordinates_path = '/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/original_file/MONARCH_d01_2008123100.nc'
                    coordinates = open_netcdf(path=coordinates_path, info=True)
                    coordinates.create_shapefile()
                    coordinates.keep_vars(['O3'])
                    coordinates.load()
                    coordinates.shapefile['O3'] = coordinates.variables['O3']['data'][-1, -1, :].ravel()

            # Rotated projection
            elif projection == 'rotated':
                # Create dataset and get shapefile
                if projection_type == 'created':
                    centre_lat = 51
                    centre_lon = 10
                    west_boundary = -35
                    south_boundary = -27
                    inc_rlat = 0.2
                    inc_rlon = 0.2
                    coordinates = create_nes(comm=None, info=True, 
                                             projection='rotated',
                                             centre_lat=centre_lat, 
                                             centre_lon=centre_lon,
                                             west_boundary=west_boundary, 
                                             south_boundary=south_boundary,
                                             inc_rlat=inc_rlat, 
                                             inc_rlon=inc_rlon)
                    coordinates.create_shapefile()

                # Open dataset and get shapefile
                elif projection_type == 'read':
                    coordinates_path = '/gpfs/scratch/bsc32/bsc32538/original_files/CAMS_MONARCH_d01_2022070412.nc'
                    coordinates = open_netcdf(path=coordinates_path, info=True)
                    coordinates.create_shapefile()
                    coordinates.keep_vars(['O3'])
                    coordinates.load()
                    coordinates.shapefile['O3'] = coordinates.variables['O3']['data'][-1, -1, :].ravel()

            coordinates.write_shapefile('coordinates_{0}_{1}'.format(projection, 
                                                                     projection_type))

            mask_path = '/gpfs/projects/bsc32/models/NES_tutorial_data/timezones_2021c/timezones_2021c.shp'
            mask = gpd.read_file(mask_path)

            # Spatial overlay (old method)
            if method == 'spatial_overlay':
                start_time = timeit.default_timer()
                intersection = coordinates.shapefile.copy()
                #intersection['area'] = intersection.geometry.area
                intersection = coordinates.spatial_overlays(intersection, mask)
                #intersection.rename(columns={'idx1': 'FID', 'idx2': 'shp_id'}, inplace=True)
                #intersection['fraction'] = intersection.geometry.area / intersection['area']
                #intersection.sort_values('fraction', ascending=False, inplace=True)
                #intersection = intersection.drop_duplicates(subset='FID', keep="first")
                #intersection.set_index('FID', inplace=True)
                #coordinates.loc[intersection.index, coordinates.shp_colname] = intersection[coordinates.shp_colname]
                time = timeit.default_timer() - start_time

            # Spatial join (new method)
            elif method == 'spatial_join':
                start_time = timeit.default_timer()
                coordinates.spatial_join(mask, method='intersection')
                time = timeit.default_timer() - start_time

            coordinates.write_shapefile('masked_coordinates_{0}_{1}_{2}'.format(projection, 
                                                                                projection_type, 
                                                                                method))

        results = []
        results.append({'Projection': projection,
                        'Projection type': projection_type,
                        'Method': '{min:02d}:{sec:02.3f}'.format(
                            min=int(time // 60), sec=time - (int(time // 60) * 60))
                        })

        comm.Barrier()

comm.Barrier()

if rank == 0:
    table = pd.DataFrame(results)
    print('RESULTS TABLE')
    print(table)
    sys.stdout.flush()
