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

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

parallel_method = 'Y'

result_path = "Times_test_5.3_horiz_interp_conservative.py_{0}_{1:03d}.csv".format(parallel_method, size)
result = pd.DataFrame(index=['read', 'calcul', 'write'],
                      columns=['5.3.1.Only interp', '5.3.2.Create_WM', "5.3.3.Read_WM"])

# NAMEE
src_path = "/gpfs/projects/bsc32/models/NES_tutorial_data/MONARCH_d01_2022111512.nc"
var_list = ['O3']

# ======================================================================================================================
# ======================================        Only interp        =====================================================
# ======================================================================================================================
test_name = '5.3.1.Only interp'
print(test_name)

# READING
st_time = timeit.default_timer()

# Source data
src_nes = open_netcdf(src_path, parallel_method=parallel_method)
src_nes.keep_vars(var_list)
src_nes.load()

# Destination Grid
lat_1 = 37
lat_2 = 43
lon_0 = -3
lat_0 = 40
nx = 397
ny = 397
inc_x = 4000
inc_y = 4000
x_0 = -807847.688
y_0 = -797137.125
dst_nes = create_nes(comm=None, info=False, projection='lcc', lat_1=lat_1, lat_2=lat_2, lon_0=lon_0, lat_0=lat_0,
                     nx=nx, ny=ny, inc_x=inc_x, inc_y=inc_y, x_0=x_0, y_0=y_0, parallel_method=parallel_method)
comm.Barrier()
result.loc['read', test_name] = timeit.default_timer() - st_time

st_time = timeit.default_timer()

interp_nes = src_nes.interpolate_horizontal(dst_grid=dst_nes, kind='Conservative')
comm.Barrier()
result.loc['calcul', test_name] = timeit.default_timer() - st_time

st_time = timeit.default_timer()
interp_nes.to_netcdf(test_name.replace(' ', '_') + ".nc")
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()

# ======================================================================================================================
# ======================================         Create_WM        =====================================================
# ======================================================================================================================
test_name = '5.3.2.Create_WM'
print(test_name)

# READING
st_time = timeit.default_timer()

# Source data
src_nes = open_netcdf(src_path, parallel_method=parallel_method)
src_nes.keep_vars(var_list)
src_nes.load()

# Destination Grid
lat_1 = 37
lat_2 = 43
lon_0 = -3
lat_0 = 40
nx = 397
ny = 397
inc_x = 4000
inc_y = 4000
x_0 = -807847.688
y_0 = -797137.125
dst_nes = create_nes(comm=None, info=False, projection='lcc', lat_1=lat_1, lat_2=lat_2, lon_0=lon_0, lat_0=lat_0,
                     nx=nx, ny=ny, inc_x=inc_x, inc_y=inc_y, x_0=x_0, y_0=y_0, parallel_method=parallel_method)
comm.Barrier()
result.loc['read', test_name] = timeit.default_timer() - st_time

# Cleaning WM
if os.path.exists("NAMEE_to_IP.nc"):
    os.remove("NAMEE_to_IP.nc")

st_time = timeit.default_timer()

interp_nes = src_nes.interpolate_horizontal(dst_grid=dst_nes, kind='Conservative', weight_matrix_path="NAMEE_to_IP.nc")
comm.Barrier()
result.loc['calcul', test_name] = timeit.default_timer() - st_time

st_time = timeit.default_timer()
interp_nes.to_netcdf(test_name.replace(' ', '_') + ".nc")
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()

# ======================================================================================================================
# ======================================           Read_WM         =====================================================
# ======================================================================================================================
test_name = "5.3.3.Read_WM"
print(test_name)

# READING
st_time = timeit.default_timer()

# Source data
src_nes = open_netcdf(src_path, parallel_method=parallel_method)
src_nes.keep_vars(var_list)
src_nes.load()

# Destination Grid
lat_1 = 37
lat_2 = 43
lon_0 = -3
lat_0 = 40
nx = 397
ny = 397
inc_x = 4000
inc_y = 4000
x_0 = -807847.688
y_0 = -797137.125

dst_nes = create_nes(comm=None, info=False, projection='lcc', lat_1=lat_1, lat_2=lat_2, lon_0=lon_0, lat_0=lat_0,
                     nx=nx, ny=ny, inc_x=inc_x, inc_y=inc_y, x_0=x_0, y_0=y_0, parallel_method=parallel_method)
comm.Barrier()
result.loc['read', test_name] = timeit.default_timer() - st_time

st_time = timeit.default_timer()

interp_nes = src_nes.interpolate_horizontal(dst_grid=dst_nes, kind='Conservative', weight_matrix_path="NAMEE_to_IP.nc")
comm.Barrier()
result.loc['calcul', test_name] = timeit.default_timer() - st_time

st_time = timeit.default_timer()
interp_nes.to_netcdf(test_name.replace(' ', '_') + ".nc")
comm.Barrier()
result.loc['write', test_name] = timeit.default_timer() - st_time

result.to_csv(result_path)