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

# In[141]:


import numpy as np
import xarray as xr
import matplotlib.pyplot as plt

plt.rcParams['image.cmap'] = 'jet'
get_ipython().run_line_magic('matplotlib', 'inline')


# In[142]:


result_directory = '/data/jpetrie/MEOPAR/SalishSea/results/remin_param_sweep_new_depth/nampisrem_zz_remin_D_PON_1.4512018923044442e-06/'
tracer_file = 'SS5x5_1h_20150201_20150501_ptrc_T.nc'
mesh_mask_file = 'mesh_mask.nc'


# In[143]:


tracer_dataset = xr.open_dataset(result_directory + tracer_file)
mesh_mask = xr.open_dataset(result_directory + mesh_mask_file)


# In[151]:


nitrogen_tracers = ['NO3','NH4','DOC','POC','PHY','PHY2','MICZ','MYRI']

tracer_sum = np.zeros(tracer_dataset[nitrogen_tracers[0]].values.shape)

for tracer in nitrogen_tracers:
    tracer_sum = tracer_sum + tracer_dataset[tracer].values
# tracer_sum now contains total nitrogen concentration by t,z,y,x


# In[152]:


# Sum across x and y coordinates. Gives 2D array with nitrogen concentration by time and depth
x_y_summed = np.sum(np.sum(tracer_sum,axis = 3),axis=2)


# In[153]:


# A vector with the height of each grid element
grid_heights = np.array(mesh_mask['e3t_1d'][0])

# 2D array (time,depth) with the total quantity of nitrogen for that time and depth
grid_scaled = np.multiply(grid_heights,x_y_summed)

# Total nitrogen by time only (1D array)
depth_summed = np.sum(grid_scaled,axis = 1)


# In[154]:


plt.plot(range(len(depth_summed)),depth_summed)
plt.ylabel("Total Nitrogen")
plt.xlabel("Time Index")
plt.title("5x5 Total Nitrogen vs. Time")


# In[155]:


z = tracer_dataset.deptht.values
t= np.array([float(x) for x in tracer_dataset.time_centered.values])
tz,zt=np.meshgrid((t[:] - t[0])/10**9/3600/24,z[:])

fig, ax = plt.subplots(1, 1,figsize=(18,8))
mesh=plt.pcolormesh(tz,zt,x_y_summed.T)
fig.colorbar(mesh)
plt.gca().invert_yaxis()
#plt.title("NH4 (DON remineralization Param = " + str(DON_remin_const)+")", size = 20)
plt.xlabel("Days", size =18)
plt.ylabel("Depth", size = 18)


# In[156]:


plt.plot(tracer_dataset.deptht.values,x_y_summed[2149,:], 'ro')
plt.ylabel("Total Nitrogen")
plt.xlabel("Depth")


# In[157]:


for n in range(0,40):
    #x = np.sum(grid_scaled[:,n],axis = 1)
    fig, ax = plt.subplots(1, 1,figsize=(10,4))
    plt.plot((t[:] - t[0])/10**9/3600/24,x_y_summed[:,n])
    plt.ylabel("Total Nitrogen")
    plt.xlabel("Days")
    plt.title("Depth = " + str(z[n]))
    plt.ylim(0,22)