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

# In[1]:


import netCDF4 as nc
import datetime as dt
import subprocess
import requests
import matplotlib.pyplot as plt
import cmocean
import numpy as np
import os
import glob
import dateutil as dutil
from salishsea_tools import viz_tools
import pickle
get_ipython().run_line_magic('matplotlib', 'inline')


# In[2]:


recalc=True #False


# In[3]:


t0=dt.datetime(2014,11,15) # 1st start date of run
if recalc:
    with nc.Dataset('/ocean/eolson/MEOPAR/NEMO-forcing/grid/mesh_mask201702_noLPE.nc') as fm:
        tmask=np.copy(fm.variables['tmask'])
        umask=np.copy(fm.variables['umask'])
        vmask=np.copy(fm.variables['vmask'])
        navlon=np.copy(fm.variables['nav_lon'])
        navlat=np.copy(fm.variables['nav_lat'])
        e3t_0=np.copy(fm.variables['e3t_0'])
        e3u_0=np.copy(fm.variables['e3u_0'])
        e3v_0=np.copy(fm.variables['e3v_0'])
        e1t=np.copy(fm.variables['e1t'])
        e2t=np.copy(fm.variables['e2t'])
        e1v=np.copy(fm.variables['e1v'])
        e2u=np.copy(fm.variables['e2u'])
        A=fm.variables['e1t'][0,:,:]*fm.variables['e2t'][0,:,:]*tmask[0,0,:,:]


    #te=dt.datetime(2016,12,1)# last start date of runfnum=18
    stm=np.shape(tmask)
    SiN=2.0
    #nlen=36*2
    nlen=36#100
    dlist=[t0+dt.timedelta(days=ii*10) for ii in range(0,nlen)]

    #sdir0='/results/SalishSea/nowcast-green/'
    sdir1='/results/SalishSea/spinup_v201812/'
    #sdir3='/data/eolson/MEOPAR/SS36runs/CedarRuns/spring2015_HCMZ/'

    tmaskC=np.copy(tmask)
    tmaskC[:,:,370:490,:12]=0
    tmaskC[:,:,887:,30:70]=0


# In[4]:


if recalc:
    tlist=dlist
    SiGlobalTot=dict()
    SiTot=dict()
    BSiTot=dict()
    DiatTot=dict()
    changeSiGlobalTot=dict()
    for idir in (sdir1,):
        fformat1='%d%b%y/'
        if idir.startswith('/data/eolson/MEOPAR/SS36runs/CedarRuns/'):
            fformatT='SalishSea_1d_*_ptrc_T_%Y%m%d-*.nc'
            fformatP='SalishSea_1d_*_ptrc_T_%Y%m%d-*.nc'
        #elif idir==sdir0:
        #    fformatT='SalishSea_1h_%Y%m%d_%Y%m%d_ptrc_T.nc'
        #    fformatP='SalishSea_1h_%Y%m%d_%Y%m%d_grid_T.nc'
        elif idir==sdir1:
            fformatT='SalishSea_1d_%Y%m%d_%Y%m%d_ptrc_T.nc'
            fformatP='SalishSea_1d_%Y%m%d_%Y%m%d_carp_T.nc'
        sumSi=np.zeros((len(tlist),stm[2],stm[3]))
        sumBSi=np.zeros((len(tlist),stm[2],stm[3]))
        sumDiat=np.zeros((len(tlist),stm[2],stm[3]))
        ind=-1
        for idt0 in tlist:
            ind=ind+1
            cdir=idt0.strftime(fformat1).lower()
            iffT=idt0.strftime(fformatT)
            iffP=idt0.strftime(fformatP)
            if idir.startswith('/data/eolson/MEOPAR/SS36runs/CedarRuns/'):
                sffT=idir+iffT
                sffP=idir+iffP
            elif idir.startswith('/results/') or idir.startswith('/results2/'):
                sffT=idir+cdir+iffT
                sffP=idir+cdir+iffP
            print(sffT)
            f=nc.Dataset(glob.glob(sffT)[0])
            fP=nc.Dataset(glob.glob(sffP)[0])
            #if idir==sdir0:
            #    e3t=np.expand_dims((1+fP.variables['sossheig'][0,:,:]/np.sum(e3t_0*tmask,1)),0)*e3t_0
            if idir==sdir1:
                e3t=fP.variables['e3t'][:2,:,:,:]
            Vol=A*e3t
            sumSi[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['silicon'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumBSi[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['biogenic_silicon'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumDiat[ind,:,:]=SiN*1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['diatoms'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            f.close()
            fP.close()
        SiGlobalTot[idir]=np.sum(np.sum(sumSi+sumBSi+sumDiat,2),1)
        SiTot[idir]=np.sum(np.sum(sumSi,2),1)
        BSiTot[idir]=np.sum(np.sum(sumBSi,2),1)
        DiatTot[idir]=np.sum(np.sum(sumDiat,2),1)
        changeSiGlobalTot[idir]=[SiGlobalTot[idir][ii+1]-SiGlobalTot[idir][ii] for ii in range(0,len(tlist)-1)]


# In[5]:


#test=np.sum(np.sum((sumSi[-1,:,:]+np.sum(1e-3*7*tmaskC[0,:,:,:]*Vol[0,:,:,:],0))+sumBSi[-1,:,:]+sumDiat[-1,:,:],1),0)
#test2=np.sum(np.sum(sumSi[-1,:,:]+sumBSi[-1,:,:]+sumDiat[-1,:,:],1),0)


# In[6]:


if recalc:
    fig,ax=plt.subplots(1,1,figsize=(6,5))
    ax.plot(SiGlobalTot[sdir1]-SiGlobalTot[sdir1][0],'b-')
    ax.set_xlabel('10-day intervals since '+t0.strftime('%b $d $Y'))
    ax.set_ylabel('Difference in Total Si')


# In[7]:


# copy restart and add 7 to Si old and new


# ## repeat for N

# In[8]:


if recalc:
    tlist=dlist
    NGlobalTot=dict()
    VolTot=dict()
    NO3Tot=dict()
    NH4Tot=dict()
    PONTot=dict()
    DONTot=dict()
    DiatTot=dict()
    MyriTot=dict()
    NanoTot=dict()
    MiZoTot=dict()
    changeNGlobalTot=dict()
    for idir in (sdir1,):
        fformat1='%d%b%y/'
        if idir.startswith('/data/eolson/MEOPAR/SS36runs/CedarRuns/'):
            fformatT='SalishSea_1h_*_ptrc_T_%Y%m%d-*.nc'
            fformatP='SalishSea_1h_*_ptrc_T_%Y%m%d-*.nc'
        #elif idir==sdir0:
        #    fformatT='SalishSea_1h_%Y%m%d_%Y%m%d_ptrc_T.nc'
        #    fformatP='SalishSea_1h_%Y%m%d_%Y%m%d_grid_T.nc'
        elif idir==sdir1:
            fformatT='SalishSea_1d_%Y%m%d_%Y%m%d_ptrc_T.nc'
            fformatP='SalishSea_1d_%Y%m%d_%Y%m%d_carp_T.nc'
        sumNO3=np.zeros((len(tlist),stm[2],stm[3]))
        sumVol=np.zeros((len(tlist),stm[2],stm[3]))
        sumNH4=np.zeros((len(tlist),stm[2],stm[3]))
        sumPON=np.zeros((len(tlist),stm[2],stm[3]))
        sumDON=np.zeros((len(tlist),stm[2],stm[3]))
        sumDiat=np.zeros((len(tlist),stm[2],stm[3]))
        sumMyri=np.zeros((len(tlist),stm[2],stm[3]))
        sumNano=np.zeros((len(tlist),stm[2],stm[3]))
        sumMiZo=np.zeros((len(tlist),stm[2],stm[3]))
        ind=-1
        for idt0 in tlist:
            ind=ind+1
            cdir=idt0.strftime(fformat1).lower()
            iffT=idt0.strftime(fformatT)
            iffP=idt0.strftime(fformatP)
            if idir.startswith('/data/eolson/MEOPAR/SS36runs/CedarRuns/'):
                sffT=idir+iffT
                sffP=idir+iffP
            elif idir.startswith('/results/') or idir.startswith('/results2/'):
                sffT=idir+cdir+iffT
                sffP=idir+cdir+iffP
            print(sffT)
            f=nc.Dataset(glob.glob(sffT)[0])
            fP=nc.Dataset(glob.glob(sffP)[0])
            #if idir==sdir0:
            #    e3t=np.expand_dims((1+fP.variables['sossheig'][0,:,:]/np.sum(e3t_0*tmask,1)),0)*e3t_0
            if idir==sdir1:
                e3t=fP.variables['e3t'][:2,:,:,:]
            Vol=A*e3t
            sumVol[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumNO3[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['nitrate'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumNH4[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['ammonium'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumPON[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['particulate_organic_nitrogen'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumDON[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['dissolved_organic_nitrogen'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumDiat[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['diatoms'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumMyri[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['ciliates'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            sumMiZo[ind,:,:]=1e-3*np.sum(tmaskC[0,:,:,:]*Vol[0,:,:,:]*f.variables['microzooplankton'][0,:,:,:],0) #mmol/m3*m3*10^-3=mol
            f.close()
            fP.close()
        NGlobalTot[idir]=np.sum(np.sum(sumNO3+sumNH4+sumPON+sumDON+sumDiat+sumMyri+sumNano+sumMiZo,2),1)
        VolTot[idir]=np.sum(np.sum(sumVol,2),1)
        NO3Tot[idir]=np.sum(np.sum(sumNO3,2),1)
        NH4Tot[idir]=np.sum(np.sum(sumNH4,2),1)
        PONTot[idir]=np.sum(np.sum(sumPON,2),1)
        DONTot[idir]=np.sum(np.sum(sumDON,2),1)
        DiatTot[idir]=np.sum(np.sum(sumDiat,2),1)
        MyriTot[idir]=np.sum(np.sum(sumMyri,2),1)
        NanoTot[idir]=np.sum(np.sum(sumNano,2),1)
        MiZoTot[idir]=np.sum(np.sum(sumMiZo,2),1)
        changeNGlobalTot[idir]=[NGlobalTot[idir][ii+1]-NGlobalTot[idir][ii] for ii in range(0,len(tlist)-1)]


# In[9]:


if recalc:
    #plt.plot(SiGlobalTot[sdir0]-SiGlobalTot[sdir0][0],'r-')
    #plt.plot(SiGlobalTot[sdir3]-SiGlobalTot[sdir3][0],'g-')
    #plt.plot(40,test-SiGlobalTot[sdir1][0],'r*')
    #plt.plot(40,test2-SiGlobalTot[sdir1][0],'k*')
    fig,ax=plt.subplots(1,1,figsize=(6,5))
    ax.plot(NGlobalTot[sdir1]-NGlobalTot[sdir1][0],'r-')
    ax.set_xlabel('10-day intervals since '+t0.strftime('%b $d $Y'))
    ax.set_ylabel('Difference in Total N')


# In[10]:


if recalc:
    plt.plot(SiGlobalTot[sdir1]-SiGlobalTot[sdir1][0],'b-')
    plt.plot(NGlobalTot[sdir1]-NGlobalTot[sdir1][0],'r-')


# In[11]:


if recalc:
    pickle.dump(SiGlobalTot[sdir1],open('SiGlobalTot_HCspinup.pkl','wb'))
    pickle.dump(NGlobalTot[sdir1],open('NGlobalTot_HCspinup.pkl','wb'))


# In[12]:


SiGlobalTotHC=pickle.load(open('SiGlobalTot_HCspinup.pkl','rb'))
NGlobalTotHC=pickle.load(open('NGlobalTot_HCspinup.pkl','rb'))


# In[13]:


SiGlobalTotHCold=pickle.load(open('SiGlobalTot_HC.pkl','rb'))
NGlobalTotHCold=pickle.load(open('NGlobalTot_HC.pkl','rb'))


# In[14]:


#SiGlobalTotT3=pickle.load(open('SiGlobalTotT3.pkl','rb'))
#NGlobalTotT3=pickle.load(open('NGlobalTotT3.pkl','rb'))


# In[15]:


#SiGlobalTotZ1=pickle.load(open('SiGlobalTotZ1.pkl','rb'))
#NGlobalTotZ1=pickle.load(open('NGlobalTotZ1.pkl','rb'))


# In[16]:


#tit='spring2015_Z3'
#SiGlobalTotZ3=pickle.load(open('SiGlobalTot_'+tit+'.pkl','rb'))
#NGlobalTotZ3=pickle.load(open('NGlobalTot_'+tit+'.pkl','rb'))


# In[17]:


bdir='/data/eolson/MEOPAR/SS36runs/CedarRuns/'


# In[18]:


tit='spring15spun_Ztest'
SiGlobalTotZtest=pickle.load(open(bdir+tit+'/SiGlobalTot_'+tit+'.pkl','rb'))
NGlobalTotZtest=pickle.load(open(bdir+tit+'/NGlobalTot_'+tit+'.pkl','rb'))


# In[19]:


tit='spring15spun_Z4'
SiGlobalTotZ4=pickle.load(open(bdir+tit+'/SiGlobalTot_'+tit+'.pkl','rb'))
NGlobalTotZ4=pickle.load(open(bdir+tit+'/NGlobalTot_'+tit+'.pkl','rb'))


# In[20]:


tit='spring15spun_Z5'
SiGlobalTotZ5=pickle.load(open(bdir+tit+'/SiGlobalTot_'+tit+'.pkl','rb'))
NGlobalTotZ5=pickle.load(open(bdir+tit+'/NGlobalTot_'+tit+'.pkl','rb'))


# In[21]:


tit='spring16spun_Z6'
SiGlobalTotZ6=pickle.load(open(bdir+tit+'/SiGlobalTot_'+tit+'.pkl','rb'))
NGlobalTotZ6=pickle.load(open(bdir+tit+'/NGlobalTot_'+tit+'.pkl','rb'))


# In[22]:


tit='spring16spun_Z7'
SiGlobalTotZ7=pickle.load(open(bdir+tit+'/SiGlobalTot_'+tit+'.pkl','rb'))
NGlobalTotZ7=pickle.load(open(bdir+tit+'/NGlobalTot_'+tit+'.pkl','rb'))


# In[23]:


fig,ax=plt.subplots(1,1,figsize=(8,16))
plt.plot([t0+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotHC))],SiGlobalTotHC,'y-')
plt.plot([t0+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotHC))],NGlobalTotHC,'y-')
plt.plot([t0+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotHCold))],SiGlobalTotHCold,'r--')
plt.plot([t0+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotHCold))],NGlobalTotHCold,'r--')
plt.plot([dt.datetime(2015,1,11)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ4))],SiGlobalTotZ4,'c--')
plt.plot([dt.datetime(2015,1,11)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ4))],NGlobalTotZ4,'c--')

ax.plot([dt.datetime(2015,1,11)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ5))],SiGlobalTotZ5,'m:')
ax.plot([dt.datetime(2015,1,11)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ5))],NGlobalTotZ5,'m:')
ax.plot([dt.datetime(2015,1,11)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZtest))],SiGlobalTotZtest,'b:')
ax.plot([dt.datetime(2015,1,11)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZtest))],NGlobalTotZtest,'b:')

ax.plot([dt.datetime(2016,1,6)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ6))],SiGlobalTotZ6,'r:')
ax.plot([dt.datetime(2016,1,6)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ6))],NGlobalTotZ6,'r:')
ax.plot([dt.datetime(2016,1,6)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ7))],SiGlobalTotZ7,'g:')
ax.plot([dt.datetime(2016,1,6)+dt.timedelta(10*ii) for ii in range(0,len(SiGlobalTotZ7))],NGlobalTotZ7,'g:')


# In[ ]: