#!/usr/bin/env python

import sys
from mpi4py import MPI
import pandas as pd
import timeit
import numpy as np
from nes import *

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

parallel_method = 'Y'

result_path = "Times_test_2.1_spatial_join_{0}_{1:03d}.csv".format(parallel_method, size)
result = pd.DataFrame(index=['read', 'calculate', 'write'],
                      columns=['2.1.1.Existing_file_centroid', '2.1.2.New_file_centroid',
                               '2.1.3.Existing_file_nearest', '2.1.4.New_file_nearest',
                               '2.1.5.Existing_file_intersection', '2.1.6.New_file_intersection'])

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

# Original path: /gpfs/scratch/bsc32/bsc32538/original_files/franco_interp.nc
# Regular lat-lon grid from MONARCH
original_path = '/gpfs/projects/bsc32/models/NES_tutorial_data/franco_interp.nc'

# ======================================================================================================================
# ===================================     CENTROID EXISTING FILE     ===================================================
# ======================================================================================================================

test_name = '2.1.1.Existing_file_centroid'
if rank == 0:
    print(test_name)

# READ
st_time = timeit.default_timer()
nessy = open_netcdf(original_path, parallel_method=parallel_method)
comm.Barrier()
result.loc['read', test_name] = timeit.default_timer() - st_time

# SPATIAL JOIN
# Method can be centroid, nearest and intersection
st_time = timeit.default_timer()
nessy.spatial_join(shapefile_path, method='centroid')
comm.Barrier()
result.loc['calculate', test_name] = timeit.default_timer() - st_time
print('SPATIAL JOIN - Rank {0:03d} - Shapefile: \n{1}'.format(rank, nessy.shapefile))

# ADD TIMEZONES
timezones = nessy.shapefile['tzid'].values.reshape([nessy.lat['data'].shape[0], 
                                                    nessy.lon['data'].shape[-1]])
print('Shape 0', timezones.shape)
timezones = np.repeat(timezones[np.newaxis, :, :], [len(nessy.lev['data'])], axis=0)
print('Shape 1', timezones.shape)
timezones = np.repeat(timezones[np.newaxis, :, :, :], [len(nessy.time)], axis=0)
print('Shape 2', timezones.shape)
nessy.variables['tz'] = {'data': timezones,}

# WRITE
st_time = timeit.default_timer()
nessy.to_netcdf(test_name.replace(' ', '_') + "_{0:03d}.nc".format(size), info=True)
comm.Barrier()
result.loc['write', test_name] = timeit.default_timer() - st_time

comm.Barrier()
if rank == 0:
    print(result.loc[:, test_name])
sys.stdout.flush()

# # ======================================================================================================================
# # ===================================     CENTROID FROM NEW FILE     ===================================================
# # ======================================================================================================================

# test_name = '2.1.2.New_file_centroid'
# if rank == 0:
#     print(test_name)

# # DEFINE PROJECTION
# st_time = timeit.default_timer()
# 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)
# comm.Barrier()
# result.loc['calculate', test_name] = timeit.default_timer() - st_time
# print('FROM SHAPEFILE - Rank {0:03d} - Shapefile: \n{1}'.format(rank, nessy.shapefile))

# # ADD TIMEZONES
# timezones = nessy.shapefile['tzid'].values.reshape([nessy.lat['data'].shape[0], 
#                                                     nessy.lon['data'].shape[-1]])
# timezones = np.repeat(timezones[np.newaxis, :, :], [len(nessy.lev['data'])], axis=0)
# timezones = np.repeat(timezones[np.newaxis, :, :, :], [len(nessy.time)], axis=0)
# nessy.variables['tz'] = {'data': timezones,}

# # WRITE
# st_time = timeit.default_timer()
# nessy.to_netcdf(test_name.replace(' ', '_') + "_{0:03d}.nc".format(size), info=True)
# comm.Barrier()
# result.loc['write', test_name] = timeit.default_timer() - st_time

# comm.Barrier()
# if rank == 0:
#     print(result.loc[:, test_name])
# sys.stdout.flush()

# # ======================================================================================================================
# # ===================================     NEAREST EXISTING FILE      ===================================================
# # ======================================================================================================================

# test_name = '2.1.3.Existing_file_nearest'
# if rank == 0:
#     print(test_name)

# # READ
# st_time = timeit.default_timer()
# nessy = open_netcdf(original_path, parallel_method=parallel_method)
# comm.Barrier()
# result.loc['read', test_name] = timeit.default_timer() - st_time

# # SPATIAL JOIN
# # Method can be centroid, nearest and intersection
# st_time = timeit.default_timer()
# nessy.spatial_join(shapefile_path, method='nearest')
# comm.Barrier()
# result.loc['calculate', test_name] = timeit.default_timer() - st_time
# print('SPATIAL JOIN - Rank {0:03d} - Shapefile: \n{1}'.format(rank, nessy.shapefile))

# # ADD TIMEZONES
# timezones = nessy.shapefile['tzid'].values.reshape([nessy.lat['data'].shape[0], 
#                                                     nessy.lon['data'].shape[-1]])
# timezones = np.repeat(timezones[np.newaxis, :, :], [len(nessy.lev['data'])], axis=0)
# timezones = np.repeat(timezones[np.newaxis, :, :, :], [len(nessy.time)], axis=0)
# nessy.variables['tz'] = {'data': timezones,}

# # WRITE
# st_time = timeit.default_timer()
# nessy.to_netcdf(test_name.replace(' ', '_') + "_{0:03d}.nc".format(size), info=True)
# comm.Barrier()
# result.loc['write', test_name] = timeit.default_timer() - st_time

# comm.Barrier()
# if rank == 0:
#     print(result.loc[:, test_name])
# sys.stdout.flush()

# # ======================================================================================================================
# # ===================================     NEAREST FROM NEW FILE      ===================================================
# # ======================================================================================================================

# test_name = '2.1.4.New_file_nearest'
# if rank == 0:
#     print(test_name)

# # DEFINE PROJECTION
# st_time = timeit.default_timer()
# 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='nearest', 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)
# comm.Barrier()
# result.loc['calculate', test_name] = timeit.default_timer() - st_time
# print('FROM SHAPEFILE - Rank {0:03d} - Shapefile: \n{1}'.format(rank, nessy.shapefile))

# # ADD TIMEZONES
# timezones = nessy.shapefile['tzid'].values.reshape([nessy.lat['data'].shape[0], 
#                                                     nessy.lon['data'].shape[-1]])
# timezones = np.repeat(timezones[np.newaxis, :, :], [len(nessy.lev['data'])], axis=0)
# timezones = np.repeat(timezones[np.newaxis, :, :, :], [len(nessy.time)], axis=0)
# nessy.variables['tz'] = {'data': timezones,}

# # WRITE
# st_time = timeit.default_timer()
# nessy.to_netcdf(test_name.replace(' ', '_') + "_{0:03d}.nc".format(size), info=True)
# comm.Barrier()
# result.loc['write', test_name] = timeit.default_timer() - st_time

# comm.Barrier()
# if rank == 0:
#     print(result.loc[:, test_name])
# sys.stdout.flush()


# # ======================================================================================================================
# # ===================================   INTERSECTION EXISTING FILE   ===================================================
# # ======================================================================================================================

# test_name = '2.1.5.Existing_file_intersection'
# if rank == 0:
#     print(test_name)
    
# # READ
# st_time = timeit.default_timer()
# nessy = open_netcdf(original_path, parallel_method=parallel_method)
# comm.Barrier()
# result.loc['read', test_name] = timeit.default_timer() - st_time

# # SPATIAL JOIN
# # Method can be centroid, nearest and intersection
# st_time = timeit.default_timer()
# nessy.spatial_join(shapefile_path, method='intersection')
# comm.Barrier()
# result.loc['calculate', test_name] = timeit.default_timer() - st_time
# print('SPATIAL JOIN - Rank {0:03d} - Shapefile: \n{1}'.format(rank, nessy.shapefile))

# # ADD TIMEZONES
# timezones = nessy.shapefile['tzid'].values.reshape([nessy.lat['data'].shape[0], 
#                                                     nessy.lon['data'].shape[-1]])
# timezones = np.repeat(timezones[np.newaxis, :, :], [len(nessy.lev['data'])], axis=0)
# timezones = np.repeat(timezones[np.newaxis, :, :, :], [len(nessy.time)], axis=0)
# nessy.variables['tz'] = {'data': timezones,}

# # WRITE
# st_time = timeit.default_timer()
# nessy.to_netcdf(test_name.replace(' ', '_') + "_{0:03d}.nc".format(size), info=True)
# comm.Barrier()
# result.loc['write', test_name] = timeit.default_timer() - st_time

# comm.Barrier()
# if rank == 0:
#     print(result.loc[:, test_name])
# sys.stdout.flush()


# # ======================================================================================================================
# # ===================================   INTERSECTION FROM NEW FILE   ===================================================
# # ======================================================================================================================

# test_name = '2.1.6.New_file_intersection'
# if rank == 0:
#     print(test_name)

# # DEFINE PROJECTION
# st_time = timeit.default_timer()
# 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='intersection', 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)
# comm.Barrier()
# result.loc['calculate', test_name] = timeit.default_timer() - st_time
# print('FROM SHAPEFILE - Rank {0:03d} - Shapefile: \n{1}'.format(rank, nessy.shapefile))

# # ADD TIMEZONES
# timezones = nessy.shapefile['tzid'].values.reshape([nessy.lat['data'].shape[0], 
#                                                     nessy.lon['data'].shape[-1]])
# timezones = np.repeat(timezones[np.newaxis, :, :], [len(nessy.lev['data'])], axis=0)
# timezones = np.repeat(timezones[np.newaxis, :, :, :], [len(nessy.time)], axis=0)
# nessy.variables['tz'] = {'data': timezones,}

# # WRITE
# st_time = timeit.default_timer()
# nessy.to_netcdf(test_name.replace(' ', '_') + "_{0:03d}.nc".format(size), info=True)
# comm.Barrier()
# result.loc['write', test_name] = timeit.default_timer() - st_time

# comm.Barrier()
# if rank == 0:
#     print(result.loc[:, test_name])
# sys.stdout.flush()

# if rank == 0:
#     result.to_csv(result_path)
