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

paths = {'regular_file': {'path': '/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/original_file/MONARCH_d01_2008123100.nc',
                          'projection': 'regular',
                          'variables': ['O3'],
                          'parallel_methods': ['X', 'Y', 'T']},
         'rotated_file': {'path': '/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/OUT/stats_bnds/monarch/a45g/regional/daily_max/O3_all/O3_all-000_2021080300.nc',
                          'projection': 'rotated',
                          'variables': ['O3_all'],
                          'parallel_methods': ['X', 'Y', 'T']},
         'points_file': {'path': '/esarchive/obs/eea/eionet/hourly/pm10/pm10_202107.nc',
                         'projection': 'points',
                         'variables': [], # all
                         'parallel_methods': ['X', 'T']},
         'points_ghost_file': {'path': '/gpfs/projects/bsc32/AC_cache/obs/ghost/EANET/1.4/daily/sconcso4/sconcso4_201911.nc',
                               'projection': 'points_ghost',
                               'variables': [], # all
                               'parallel_methods': ['X', 'T']},
         'lcc_file': {'path': '/esarchive/exp/wrf-hermes-cmaq/b075/eu/hourly/pm10/pm10_2022062600.nc',
                      'projection': 'lcc',
                      'variables': [], # all
                      'parallel_methods': ['X', 'Y', 'T']},
         'mercator_file': {'path': '/esarchive/scratch/avilanova/software/NES/Jupyter_notebooks/input/mercator_grid_example.nc',
                           'projection': 'mercator',
                           'variables': [], # all
                           'parallel_methods': ['X', 'Y', 'T']}
         }

results = []

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

for name, dict in paths.items():

    path = dict['path']
    projection = dict['projection']
    variables = dict['variables']
    parallel_methods = dict['parallel_methods']

    for parallel_method in parallel_methods:
        
        if rank == 0:
            print('TEST TO USE {0} GRID IN {1} FOR {2} USING {3} NODES'.format(
                projection.upper(), parallel_method, path, size))
            sys.stdout.flush()

        try:

            # Read
            start_time = timeit.default_timer()
            nessy_1 = open_netcdf(path=path, comm=comm, info=True, parallel_method=parallel_method)
            open_time = timeit.default_timer() - start_time

            # Select variables and load
            start_time = timeit.default_timer()
            if len(variables) > 0:
                nessy_1.keep_vars(variables)
            nessy_1.load()
            load_time = timeit.default_timer() - start_time
            comm.Barrier()

            # Write in serial
            if rank == 0:
                print('WRITE IN SERIAL')
                sys.stdout.flush()
            start_time = timeit.default_timer()
            nessy_1.to_netcdf('{0}_{1}_file_{2}_serial.nc'.format(size, projection, parallel_method), 
                              info=True, serial=True)
            serial_time = timeit.default_timer() - start_time
            comm.Barrier()

            # Write in parallel
            if rank == 0:
                print('WRITE IN PARALLEL')
                sys.stdout.flush()
            start_time = timeit.default_timer()
            nessy_1.to_netcdf('{0}_{1}_file_{2}_parallel.nc'.format(size, projection, parallel_method), 
                              info=True)
            parallel_time = timeit.default_timer() - start_time
            comm.Barrier()

            # Write in chunks
            if rank == 0:
                print('WRITE IN CHUNKS')
                sys.stdout.flush()
            start_time = timeit.default_timer()
            nessy_1.to_netcdf('{0}_{1}_file_{2}_chunking.nc'.format(size, projection, parallel_method), 
                              info=True, chunking=True)
            chunking_time = timeit.default_timer() - start_time
            comm.Barrier()

            # Close everything
            del nessy_1

            if rank == 0:
                print('Test was successful for {0} projection in {1}'.format(projection, parallel_method))
                sys.stdout.flush()

            # End timer and save results
            results.append({'Projection': projection, 
                            'Method': parallel_method, 
                            'Open': '{min:02d}:{sec:02.3f}'.format(
                                min=int(open_time // 60), sec=open_time - (int(open_time // 60) * 60)),
                            'Load': '{min:02d}:{sec:02.3f}'.format(
                                min=int(load_time // 60), sec=load_time - (int(load_time // 60) * 60)),
                            'Serial': '{min:02d}:{sec:02.3f}'.format(
                                min=int(serial_time // 60), sec=serial_time - (int(serial_time // 60) * 60)),
                            'Chunking': '{min:02d}:{sec:02.3f}'.format(
                                min=int(chunking_time // 60), sec=chunking_time - (int(chunking_time // 60) * 60)),
                            'Parallel': '{min:02d}:{sec:02.3f}'.format(
                                min=int(parallel_time // 60), sec=parallel_time - (int(parallel_time // 60) * 60))
                            })
            
            comm.Barrier()

        except Exception as e:
            print(e)

        sys.stdout.flush()

comm.Barrier()

if rank == 0:
    table = pd.DataFrame(results)
    print('RESULTS TABLE')
    print(table)
    table.to_csv('{0}_results.csv'.format(size))
    sys.stdout.flush()
