import numpy as np
import xarray as xr
import math
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import os
import f90nml
from salishsea_tools import metric_tools_5x5 as met

%matplotlib inline
plt.rcParams['image.cmap'] = 'jet'
plt.rc('xtick', labelsize=20) 
plt.rc('ytick', labelsize=20) 

reference_namelist_file = '/data/jpetrie/MEOPAR/SS-run-sets/SS-SMELT/namelists/namelist_pisces_cfg_5x5_NewIC'
reference_bio_params = f90nml.read(reference_namelist_file)

tracer_file = 'SS5x5_1h_20150201_20150501_ptrc_T.nc'

param_metrics = pd.DataFrame()

batch_directories = [
                     #'/data/jpetrie/MEOPAR/SalishSea/results/nampiszoo_june_14/',
                     '/data/jpetrie/MEOPAR/SalishSea/results/nampisopt_june_14/',
                     #'/data/jpetrie/MEOPAR/SalishSea/results/nampismes_june_14/',
                     #'/data/jpetrie/MEOPAR/SalishSea/results/nampissink_june_17/',
                     #'/data/jpetrie/MEOPAR/SalishSea/results/nampisprod_june_16/',
                     #'/data/jpetrie/MEOPAR/SalishSea/results/nampismort_june_17/',
                     #'/data/jpetrie/MEOPAR/SalishSea/results/nampisrem_june_17/',
                     #'/data/jpetrie/MEOPAR/SalishSea/results/nampismezo_june_20/',
]

metric_func_list = [
                    met.mean_NH4_at_depth,
                    met.mean_NO3_at_depth,
                    met.mean_DON_at_depth,
                    met.mean_PON_at_depth,
                    met.time_of_peak_PHY2,
                    met.time_surface_NO3_drops_below_4,
                    met.peak_3_day_biomass,
]

for batch_dir in batch_directories:
    for file in os.listdir(batch_dir):
        if os.path.isfile(batch_dir + '/' + file + '/' + tracer_file) and 'zz_frac_waste' not in file:
            last_underscore = file.rfind('_')
            first_underscore = file.find('_')
            param_section = file[:first_underscore]
            param_name = file[(first_underscore+1):last_underscore]
            param_val = float(file[(last_underscore+1):])
            param_desc = param_section + ' ' + param_name
            
            original_val = reference_bio_params[param_section][param_name]
            
            if type(original_val) is list:
                param_scale = round(param_val/original_val[0], 3)
            else:
                param_scale = round(param_val/original_val, 3)
            
            grid_t = xr.open_dataset(batch_dir + '/' + file +'/' + tracer_file)
            
            for metric_func in metric_func_list:
                metric_val = metric_func(grid_t)
                metric_name = metric_func.__name__
                # inefficient to keep appending, but much less expensive than other parts of the loop so it doesn't matter
                param_metrics = param_metrics.append(pd.DataFrame({"PARAM_SECTION":[param_section],
                                                                   "PARAM_NAME":[param_name],
                                                                   "PARAM_DESC":[param_desc],
                                                                   "PARAM_VAL":[param_val],
                                                                   "PARAM_SCALE":[param_scale],
                                                                   "METRIC_NAME":[metric_name],
                                                                   "METRIC_VAL": [metric_val]}))

param_metrics["PARAM_SCALE"] = "PARAM_SCALE_" + param_metrics["PARAM_SCALE"].astype(str)
wide_format_metrics = pd.pivot_table(param_metrics, values='METRIC_VAL', index=['PARAM_SECTION','PARAM_NAME', 'PARAM_DESC', 'METRIC_NAME'], columns=['PARAM_SCALE'])
wide_format_metrics.reset_index(inplace=True)
wide_format_metrics["SLOPE"] = (wide_format_metrics["PARAM_SCALE_1.1"] - wide_format_metrics["PARAM_SCALE_0.9"])/0.2