#!/usr/bin/env python
# coding: utf-8

# This notebook compares the output of four biological simulations ran on a 5x5 grid.
# 
# The initial nutrient and plankton ditributions are different for each of these simulations but everything else is the same.

# In[1]:


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

get_ipython().run_line_magic('matplotlib', 'inline')
plt.rcParams['image.cmap'] = 'Paired'
#sns.set(font_scale=2)


# In[2]:


# Parent directory for the results of this run
group_results_directory = '/data/jpetrie/MEOPAR/SalishSea/results/IC_comparison/'
tracer_file = 'SS5x5_1h_20150201_20150501_ptrc_T.nc'
individual_directories = []
months = []
for file in os.listdir(group_results_directory):
    if file.endswith("_IC"):
        month = file.split("_IC")[0]
        individual_directories.append(file)
        months.append(month)


tracer_datasets = [xr.open_dataset(group_results_directory + '/' + file +'/' + tracer_file) for file in individual_directories]
dataset_dict = dict(zip(months, tracer_datasets))


# In[11]:


# Plot depth/time/PHY heatmap for each value of alpha
for month in dataset_dict:
    z = dataset_dict[month].deptht.values
    t= np.array([float(x) for x in dataset_dict[month].time_centered.values])
    tz,zt=np.meshgrid((t[:] - t[0])/10**9/3600/24,z[:])
    
    fig, ax = plt.subplots(1, 1,figsize=(10,4))
    mesh=plt.pcolormesh(tz,zt,dataset_dict[month].NO3.values[:,:,2,2].T, vmin = 10, vmax = 40)
    fig.colorbar(mesh)
    plt.gca().invert_yaxis()
    plt.title("NO3 (Initial Month = " + str(month)+")", size = 20)
    plt.xlabel("Days", size =18)
    plt.ylabel("Depth", size = 18)


# In[10]:


# Plot depth/time/PHY heatmap for each value of alpha
for month in dataset_dict:
    z = dataset_dict[month].deptht.values
    t= np.array([float(x) for x in dataset_dict[month].time_centered.values])
    tz,zt=np.meshgrid((t[:] - t[0])/10**9/3600/24,z[:])
    
    fig, ax = plt.subplots(1, 1,figsize=(10,4))
    mesh=plt.pcolormesh(tz,zt,dataset_dict[month].Si.values[:,:,2,2].T, vmin = 30, vmax = 70)
    fig.colorbar(mesh)
    plt.gca().invert_yaxis()
    plt.title("Si (Initial Month = " + str(month)+")", size = 20)
    plt.xlabel("Days", size =18)
    plt.ylabel("Depth", size = 18)


# In[5]:


# Plot depth/time/PHY heatmap for each value of alpha
for month in dataset_dict:
    z = dataset_dict[month].deptht.values
    t= np.array([float(x) for x in dataset_dict[month].time_centered.values])
    tz,zt=np.meshgrid((t[:] - t[0])/10**9/3600/24,z[:])
    
    fig, ax = plt.subplots(1, 1,figsize=(10,4))
    mesh=plt.pcolormesh(tz,zt,dataset_dict[month].PHY.values[:,:,2,2].T, vmin = 0, vmax = 3)
    fig.colorbar(mesh)
    plt.gca().invert_yaxis()
    plt.title("PHY (Initial Month = " + str(month)+")", size = 20)
    plt.xlabel("Days", size =18)
    plt.ylabel("Depth", size = 18)


# In[14]:


# Plot depth/time/PHY heatmap for each value of alpha
for month in dataset_dict:
    z = dataset_dict[month].deptht.values
    t= np.array([float(x) for x in dataset_dict[month].time_centered.values])
    tz,zt=np.meshgrid((t[:] - t[0])/10**9/3600/24,z[:])
    
    fig, ax = plt.subplots(1, 1,figsize=(10,4))
    mesh=plt.pcolormesh(tz,zt,dataset_dict[month].PHY2.values[:,:,2,2].T, vmin = 0, vmax = 5)
    fig.colorbar(mesh)
    plt.gca().invert_yaxis()
    plt.title("PHY2 (Initial Month = " + str(month)+")", size = 20)
    plt.xlabel("Days", size =18)
    plt.ylabel("Depth", size = 18)


# In[13]:


# Convert list of xarrays to a single pandas dataframe. Easier to manipulate and plot
dataframe_list = []
for month in dataset_dict:
    df = dataset_dict[month].to_dataframe()
    df['MONTH_IC'] = month
    dataframe_list.append(df)

df = pd.concat(dataframe_list)

df_subset = df.query('x == 2 and y == 2 and axis_nbounds == 0') # Subset for certain grids of the 5x5
df_subset.reset_index(inplace=True) # Get rid of that multiindex madness


# In[14]:


# Create a dataframe with depth averaged tracer values. If all grid heights are the same this is the mean value

depth_averaged = df_subset.groupby(['time_counter','MONTH_IC','x','y']).mean()
depth_averaged.reset_index(inplace=True)


# In[29]:


plot_vars = [ 'DOC','MICZ', 'MYRI', 'NH4', 'NO3', 'O2', 'PHY', 'PHY2', 'POC', 'Si', 'bSi']

for var in plot_vars:
    fg = sns.FacetGrid(data=depth_averaged, hue='MONTH_IC',size =8)
    fg.map(plt.plot, 'time_counter', var).add_legend()
    fg.set_xticklabels(rotation=-35)
    plt.subplots_adjust(top=0.9)
    fg.fig.suptitle(var)