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

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

parallel_method = 'T'

result_path = "Times_test_3.1_vertical_interp_{0}_{1:03d}.csv".format(parallel_method, size)
result = pd.DataFrame(index=['read', 'calculate'],
                      columns=['3.1.1.Interp'])

# ======================================================================================================================
# ===============================================  VERTICAL INTERPOLATION  =============================================
# ======================================================================================================================

test_name = '3.1.1.Interp'
if rank == 0:
    print(test_name)

# READ
st_time = timeit.default_timer()

# Original path: /esarchive/exp/monarch/a4dd/original_files/000/2022111512/MONARCH_d01_2022111512.nc
# Rotated grid from MONARCH
source_path = '/gpfs/projects/bsc32/models/NES_tutorial_data/MONARCH_d01_2022111512.nc'

# Read source data
source_data = open_netcdf(path=source_path, info=True)

# Select time and load variables
source_data.sel(time_min=datetime(year=2022, month=11, day=16, hour=0), 
                time_max=datetime(year=2022, month=11, day=16, hour=0))
source_data.keep_vars(['O3', 'mid_layer_height_agl'])
source_data.load()

comm.Barrier()
result.loc['read', test_name] = timeit.default_timer() - st_time

# INTERPOLATE
st_time = timeit.default_timer()
level_list = [0.,50.,100.,250.,500.,750.,1000.,2000.,3000.,5000.]
interp_nes = source_data.interpolate_vertical(level_list, info=True, kind='linear', extrapolate=None)
comm.Barrier()
result.loc['calculate', test_name] = timeit.default_timer() - st_time

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