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

# test_list = [{'in_file': "/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/original_file/MONARCH_d01_2008123100.nc",
#               'prefix': 'small'},
#              {'in_file': "/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/original_file/MONARCH_d01_2021080212.nc",
#               'prefix': 'medium'},
#              {'in_file': "/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/original_file/MONARCH_d01_2017123000.nc",
#               'prefix': 'large'},
#              ]
test_list = [{'in_file': "/gpfs/scratch/bsc32/bsc32538/mr_multiplyby/original_file/MONARCH_d01_2017123000.nc",
              'prefix': 'large'},
             ]
# parallel_method = 'X'
# parallel_method = 'Y'
parallel_method = 'T'


times = pd.DataFrame(index=[0])
# size = MPI.COMM_WORLD.Get_size()

for test in test_list:
    # ======================================================================================================================

    st_time = timeit.default_timer()
    # ===== Reading
    nessy = open_netcdf(path=test['in_file'], parallel_method=parallel_method)
    size = nessy.size
    if nessy.master:
        print("\n=====", test['prefix'], "=====")

    nessy.keep_vars(['O3', 'T', 'U', 'V', 'layer_thickness'])
    # nessy.keep_vars(['O3'])
    # ===== END Reading
    spent_time = timeit.default_timer() - st_time
    if nessy.master:
        print("Init time {0:02d}:{1:02.3f} (tot: {2:.3f} s)".format(
            int(spent_time // 60), spent_time - (int(spent_time // 60) * 60), spent_time))
        sys.stdout.flush()
        times["{0}_Init".format(test['prefix'])] = spent_time

    # ======================================================================================================================

    st_time = timeit.default_timer()
    # ===== Loading
    nessy.load()
    # ===== END Loading
    spent_time = timeit.default_timer() - st_time
    if nessy.master:
        print("Loading time {min:02d}:{sec:02.3f} (tot: {tot:.3f} s ;  var: {var:.3f}s/var)".format(
            min=int(spent_time // 60), sec=spent_time - (int(spent_time // 60) * 60),
            var=spent_time / len(nessy.variables.keys()), tot=spent_time))
        sys.stdout.flush()
        times["{0}_Load".format(test['prefix'])] = spent_time

    # ======================================================================================================================

    if test['prefix'] not in ['large']:
        st_time = timeit.default_timer()
        # ===== Serial Write
        nessy.to_netcdf('nc_test_{0}_{1}_serial.nc'.format(test['prefix'], nessy.size), serial=True)
        # ===== END Serial Write
        spent_time = timeit.default_timer() - st_time
        if nessy.master:
            print("Writing in serial time {min:02d}:{sec:02.3f} (tot: {tot:.3f} s ;  var: {var:.3f}s/var)".format(
                min=int(spent_time // 60), sec=spent_time - (int(spent_time // 60) * 60),
                var=spent_time / len(nessy.variables.keys()), tot=spent_time))
            sys.stdout.flush()
            times["{0}_Serial".format(test['prefix'])] = spent_time

    # ======================================================================================================================

    st_time = timeit.default_timer()
    # ===== Parallel Chunk Write
    nessy.to_netcdf('nc_test_{0}_{1}_parallel_chunked.nc'.format(test['prefix'], nessy.size), chunking=True)
    # ===== END Parallel Chunk Write
    spent_time = timeit.default_timer() - st_time
    if nessy.master:
        print("Writing in chunked parallel time {min:02d}:{sec:02.3f} (tot: {tot:.3f} s ;  var: {var:.3f}s/var)".format(
            min=int(spent_time // 60), sec=spent_time - (int(spent_time // 60) * 60),
            var=spent_time / len(nessy.variables.keys()), tot=spent_time))
        sys.stdout.flush()
        times["{0}_Chunk".format(test['prefix'])] = spent_time

    # ======================================================================================================================

    st_time = timeit.default_timer()
    # ===== Parallel Write
    nessy.to_netcdf('nc_test_{0}_{1}_parallel.nc'.format(test['prefix'], nessy.size))
    # ===== END Parallel Write
    spent_time = timeit.default_timer() - st_time
    if nessy.master:
        print("Writing in parallel time {min:02d}:{sec:02.3f} (tot: {tot:.3f} s ;  var: {var:.3f}s/var)".format(
            min=int(spent_time // 60), sec=spent_time - (int(spent_time // 60) * 60),
            var=spent_time / len(nessy.variables.keys()), tot=spent_time))
        sys.stdout.flush()
        times["{0}_Parallel".format(test['prefix'])] = spent_time

    # ======================================================================================================================


times.to_csv("times_{0}.csv".format(str(size).zfill(3)))
