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

# In[1]:


import netCDF4 as nc
from matplotlib import pyplot as plt
import numpy as np
import glob
import pickle
from salishsea_tools import evaltools as et
import datetime as dt
import os
import re
import cmocean
get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


with nc.Dataset('/ocean/eolson/MEOPAR/NEMO-forcing/grid/mesh_mask201702_noLPE.nc') as mesh:
    tmask=mesh.variables['tmask'][0,:,:,:]
    e1t=np.expand_dims(mesh.variables['e1t'][:,:,:],1)
    e2t=np.expand_dims(mesh.variables['e2t'][:,:,:],1)
SOGtmaskPath='/ocean/eolson/MEOPAR/northernNO3PaperCalcs/save/SOGtmask.pkl'
(tmaskSOG,_,_,_,_)=pickle.load(open(SOGtmaskPath,'rb'))


# In[3]:


idir='/data/eolson/results/MEOPAR/SS36runs/CedarRuns/shortTestLeap/'


# In[4]:


fP=nc.Dataset(glob.glob(idir+'/SalishSea_1d_*_dian_T*.nc')[0])


# In[5]:


fP2=nc.Dataset(glob.glob('/data/eolson/results/MEOPAR/SS36runs/CedarRuns/shortTestSNAPe3tn/SalishSea_1d_*_dian_T*.nc')[0])


# In[6]:


fP.variables.keys()


# In[7]:


fP2.variables.keys()


# In[8]:


allSum=np.sum(np.sum(np.sum(tmaskSOG*fP.variables['ALLTRNO3'][:,:,:,:],3),2),1)
afiltSum=np.sum(np.sum(np.sum(tmaskSOG*fP.variables['ATF_NO3'][:,:,:,:],3),2),1)
no3Sum=np.sum(np.sum(np.sum(tmaskSOG*e1t*e2t*fP.variables['NO3_E3TSNAP'][:,:,:,:],3),2),1)
nitrSum=np.sum(np.sum(np.sum(tmaskSOG*fP.variables['NITR'][:,:,:,:],3),2),1)
PPSum=np.sum(np.sum(np.sum(tmaskSOG*(fP.variables['PPDIATNO3V'][:,:,:,:]+\
                                             fP.variables['PPPHYNO3V'][:,:,:,:]+\
                                             fP.variables['PPMRUBNO3V'][:,:,:,:]),3),2),1)
rivSum=np.sum(np.sum(np.sum(tmaskSOG*fP.variables['RIVNO3'][:,:,:,:],3),2),1)-nitrSum+PPSum
physSum=np.sum(np.sum(np.sum(tmaskSOG*fP.variables['PHYSTRNO3'][:,:,:,:],3),2),1)-rivSum-nitrSum+PPSum


# In[9]:


physSum2=np.sum(np.sum(np.sum(tmaskSOG*fP2.variables['PHYSTRNO3'][:,:,:,:],3),2),1)
allSum2=np.sum(np.sum(np.sum(tmaskSOG*fP2.variables['ALLTRNO3'][:,:,:,:],3),2),1)
afiltSum2=np.sum(np.sum(np.sum(tmaskSOG*fP2.variables['AFILTNO3'][:,:,:,:],3),2),1)
no3Sum2=np.sum(np.sum(np.sum(tmaskSOG*e1t*e2t*fP2.variables['NO3_E3TSNAP'][:,:,:,:],3),2),1)
nitrSum2=np.sum(np.sum(np.sum(tmaskSOG*fP2.variables['NITR'][:,:,:,:],3),2),1)
PPSum2=np.sum(np.sum(np.sum(tmaskSOG*(fP2.variables['PPDIATNO3V'][:,:,:,:]+\
                                             fP2.variables['PPPHYNO3V'][:,:,:,:]+\
                                             fP2.variables['PPMRUBNO3V'][:,:,:,:]),3),2),1)
rivSum2=np.sum(np.sum(np.sum(tmaskSOG*fP2.variables['RIVNO3'][:,:,:,:],3),2),1)


# In[10]:


no3diff=(no3Sum[1:]-no3Sum[:-1])/(24*3600)
no3diff2=(no3Sum2[1:]-no3Sum2[:-1])/(24*3600)
#no3diff2=np.concatenate((no3diff,[(np.sum(np.sum(np.sum(tmaskSOG*e1t*e2t*fT.variables['NO3_E3T'][-1,:,:,:])))-no3sum[-1])/(23*3600)]))


# In[11]:


afiltSum


# In[ ]:





# In[12]:


fig,ax=plt.subplots(1,3,figsize=(16,5))
ax[0].plot(range(1,5),no3diff,'k-',label='no3diff')
ax[0].plot(range(0,5),allSum,'g--',label='all')
#ax[0].plot(range(0,5),afiltSum,'b--',label='afilt')
ax[0].legend()
ax[1].plot(range(0,5),allSum,'g-',label='all')
ax[1].plot(range(0,5),physSum,'k--',label='phys+bio+afilt')
ax[1].plot(range(0,5),afiltSum,'y--',label='afilt')
ax[1].plot(range(0,5),nitrSum-PPSum,'m--',label='bio')
ax[1].plot(range(0,5),physSum+nitrSum-PPSum+rivSum,'c--',label='phys+bio+riv')
#ax[1].plot(range(0,5),bioSum,'k-',label='bio')
#ax[1].plot(range(0,5),nitrSum-PPSum,'r--',label='NITR-PP')
ax[1].legend()
ax[2].plot(range(0,5),nitrSum-PPSum,'r--',label='NITR-PP')
ax[2].legend()


# In[31]:


fig,ax=plt.subplots(1,3,figsize=(16,5))
ax[0].plot(range(1,5),no3diff,'k-',label='no3diff')
ax[0].plot(range(1,5),no3diff2,'r--',label='no3diff2')
ax[0].plot(range(0,5),allSum,'g-',label='all')
ax[0].plot(range(0,5),allSum2,'c--',label='all2')
ax[0].plot(0,allSum[0]*1.1,'k*')
ax[0].legend()
ax[0].set_ylim(-1200000,1200000)
ax[1].plot(range(0,5),physSum,'k-',label='phys')
ax[1].plot(range(0,5),physSum2,'r-',label='phys2')
ax[1].plot(range(0,5),nitrSum-PPSum,'m-',label='bio')
ax[1].plot(range(0,5),nitrSum2-PPSum2,'c--',label='bio2')
#ax[1].plot(range(0,5),afiltSum,'y--',label='afilt')
#ax[1].plot(range(0,5),nitrSum-PPSum,'m--',label='bio')
#ax[1].plot(range(0,5),physSum+nitrSum-PPSum+rivSum,'c--',label='phys+bio+riv')
#ax[1].plot(range(0,5),bioSum,'k-',label='bio')
#ax[1].plot(range(0,5),nitrSum-PPSum,'r--',label='NITR-PP')
ax[1].legend()
ax[1].set_ylim(-1200000,1600000)
ax[2].plot(range(0,5),nitrSum-PPSum,'r--',label='NITR-PP')
ax[2].plot(range(0,5),rivSum,'y-',label='riv')
ax[2].plot(range(0,5),rivSum2,'g--',label='riv2')
ax[2].legend()
ax[2].set_ylim(-800000,100000)


# In[29]:


(nitrSum2-PPSum2)/(nitrSum-PPSum)


# In[ ]:





# In[ ]:





# In[14]:


with nc.Dataset(glob.glob(idir+'/SalishSea_1d_*_Malaspina_U*.nc')[0]) as f:
    malUA=np.sum(np.sum(f.variables['NO3TVDX'][:,:,:,0],2),1)
    malUD=np.sum(np.sum(f.variables['ULDFNO3'][:,:,:,0],2),1)


# In[15]:


with nc.Dataset(glob.glob(idir+'/SalishSea_1d_*_Haro_V*.nc')[0]) as f:
    harVA=np.sum(np.sum(f.variables['NO3TVDY'][:,:,0,:],2),1)
    harVD=np.sum(np.sum(f.variables['VLDFNO3'][:,:,0,:],2),1)


# In[16]:


with nc.Dataset(glob.glob(idir+'/SalishSea_1d_*_SJC_V*.nc')[0]) as f:
    sjcVA=np.sum(np.sum(f.variables['NO3TVDY'][:,:,0,:],2),1)
    sjcVD=np.sum(np.sum(f.variables['VLDFNO3'][:,:,0,:],2),1)


# In[17]:


with nc.Dataset(glob.glob(idir+'/SalishSea_1d_*_Rosario_V*.nc')[0]) as f:
    rosVA=np.sum(np.sum(f.variables['NO3TVDY'][:,:,0,:],2),1)
    rosVD=np.sum(np.sum(f.variables['VLDFNO3'][:,:,0,:],2),1)


# In[18]:


with nc.Dataset(glob.glob(idir+'/SalishSea_1d_*_Sutil_V*.nc')[0]) as f:
    sutVA=np.sum(np.sum(f.variables['NO3TVDY'][:,:,0,:],2),1)
    sutVD=np.sum(np.sum(f.variables['VLDFNO3'][:,:,0,:],2),1)


# In[19]:


with nc.Dataset(glob.glob(idir+'/SalishSea_1d_*_Discovery_V*.nc')[0]) as f:
    disVA=np.sum(np.sum(f.variables['NO3TVDY'][:,:,0,:],2),1)
    disVD=np.sum(np.sum(f.variables['VLDFNO3'][:,:,0,:],2),1)


# In[20]:


fig,ax=plt.subplots(1,3,figsize=(16,6))
ax[0].plot(range(0,5),-1*malUA,'r-',label='Malaspina')
ax[0].plot(range(0,5),-1*sutVA,'m-',label='Sutil')
ax[0].plot(range(0,5),-1*disVA,'y-',label='Discovery')
ax[0].plot(range(0,5),harVA,'b-',label='Haro')
ax[0].plot(range(0,5),sjcVA,'c-',label='SJC')
ax[0].plot(range(0,5),rosVA,'g-',label='Rosario')
adv=harVA+sjcVA+rosVA-disVA-sutVA-malUA
ax[0].plot(range(0,5),adv,'k-',label='sum')
ax[0].legend()
ax[0].set_title('Advection')
ax[1].plot(range(0,5),-1*malUD,'r-',label='Malaspina')
ax[1].plot(range(0,5),-1*sutVD,'m-',label='Sutil')
ax[1].plot(range(0,5),-1*disVD,'y-',label='Discovery')
ax[1].plot(range(0,5),harVD,'b-',label='Haro iso')
ax[1].plot(range(0,5),sjcVD,'c-',label='SJC')
ax[1].plot(range(0,5),rosVD,'g-',label='Rosario')
dif=harVD+sjcVD+rosVD-disVD-sutVD-malUD
ax[1].plot(range(0,5),dif,'k-',label='sum')
ax[1].legend()
ax[1].set_title('Diffusion')
ax[2].plot(range(0,5),physSum,'k-',label='phys')
ax[2].plot(range(0,5),adv+dif,'c--',label='lateral')
ax[2].plot(range(0,5),adv+dif+rivSum,'r--',label='rivers+lateral')
ax[2].plot(0,adv[0]+dif[0]+rivSum[0],'k*')
#ax[2].plot(np.arange(0.5,4,1),no3diff,'k--',label='no3diff')
ax[2].legend()


# In[21]:


adv[0],dif[0],rivSum[0],adv[0]+dif[0],adv[0]+dif[0]+rivSum[0],physSum[0]

np.max(np.abs(harVA-harTVDY))/np.max(np.abs(harTVDY))*100np.max(np.abs(harVD-harYM))/np.max(np.abs(harYM))*100
# In[22]:


## phys difference
adv+dif+rivSum-physSum


# In[ ]:





# In[ ]:





# In[ ]: