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

# In[1]:


import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import datetime as dt
from salishsea_tools import evaltools as et, viz_tools
import os
import datetime as dt
import gsw
import matplotlib.gridspec as gridspec
import matplotlib as mpl
import matplotlib.dates as mdates
import cmocean as cmo
import scipy.interpolate as sinterp
import cmocean
import json
import f90nml
from collections import OrderedDict

fs=16
mpl.rc('xtick', labelsize=fs)
mpl.rc('ytick', labelsize=fs)
mpl.rc('legend', fontsize=fs)
mpl.rc('axes', titlesize=fs)
mpl.rc('axes', labelsize=fs)
mpl.rc('figure', titlesize=fs)
mpl.rc('font', size=fs)
mpl.rc('font', family='sans-serif', weight='normal', style='normal')

import warnings
#warnings.filterwarnings('ignore')
from IPython.display import Markdown, display

get_ipython().run_line_magic('matplotlib', 'inline')



# # Import Puget Sound data and create dataframe

# In[2]:


df=pd.read_excel('/ocean/ksuchy/MOAD/observe/SSMSP_TransBoundary_Extract_2014-2019_Verticals_Keister_Lab.xlsx',engine='openpyxl',sheet_name='SSMSP 2014-2019 Density&Biomass')


# In[3]:


df


# In[4]:


df.keys()


# In[6]:


df.drop(columns=['BugSampleID','Project','Sampling Group','Date Category','Tow Type','Species Size Category Aggregation','Collection Confidence','Carbon Regression EQ','C to DW multipier','DW Regression Equation','C/DW Source'
                       ],inplace=True)


# In[7]:


df


# In[8]:


df.keys()


# In[9]:


df['Broad Group'].unique()


# In[10]:


colListBroad=('Amphipoda', 'Crabs','Krill-Adult Juvenile','Krill-Calyptopis','Krill-Furcilia','Copepod', 'Larvacea',)


# In[39]:


colListGenus=('METRIDIA PACIFICA','NEOCALANUS PLUMCHRUS','CALANOIDA','EUCALANUS BUNGII','CALANUS PACIFICUS','CALANUS MARSHALLAE','NEOCALANUS CRISTATUS','EUCALANUS CALIFORNICUS','EUCALANUS','METRIDIIDAE','NEOCALANUS')


# In[11]:


#l1=set(df['order'].values)
#l2=set(df['genus species'].values)
#l3=l1|l2


# In[12]:


df['Genus species'].unique()


# In[ ]:





# ## Convert date to proper format

# In[11]:


df['Sample Date'][0],df['Sample Time'][0]


# In[12]:


df['Sample Date'][1000:1020]


# In[13]:


df['dtPac']=[dt.datetime.combine(idate, itime) for idate, itime \
                 in zip(df['Sample Date'],df['Sample Time'])]


# In[14]:


df['dtUTC']=[et.pac_to_utc(ii) for ii in df['dtPac']]


# In[17]:


df['dtUTC']


# In[15]:


df.rename(columns={'Sample Code':'Key','Latitude':'Lat','Longitude':'Lon'},inplace=True)


# In[16]:


df.loc[0]


# In[17]:


towIDlist=['Key','Station', 'Site Name', 'Basin', 'Sub Basin', 'Lat',
       'Lon', 'Sample Date', 'Sample Year', 'Sample Month',
       'Sample Time', 'Mesh Size', 'Diameter (cm)', 'Station Depth (m)',
       'Max Tow Depth (m)', 'Min Tow Depth (m)','dtUTC']


# In[18]:


towIDlist2=['Key', 'Station', 'Sample Date','Sample Time','dtUTC']


# In[22]:


len(np.unique(df['Key']))


# In[19]:


len(df.groupby(towIDlist)),len(df.groupby(towIDlist2)),len(df.groupby(['Key']))


# ### Create a biomass dataframe

# In[23]:


biomassDF=df.groupby(towIDlist2,as_index=False).first()\
     .loc[:,towIDlist].copy(deep=True)


# In[26]:


biomassDF.keys()


# In[27]:


biomassDF


# In[30]:


biomassDF2=biomassDF.loc[biomassDF['Sample Year']==2015].copy(deep=True)


# In[31]:


df2=df.loc[df['Sample Year']==2015].copy(deep=True)


# In[35]:


def getbiomass(colname,key,origdf,matchcol):        
    biomassArray=origdf.loc[(origdf.Key==key)&(origdf[matchcol]==colname),
                    ['Final Carbon (mg/m3)']]
    biomass=np.nansum(biomassArray)
    
    return biomass


# In[37]:


for icol in colListBroad:
    biomassDF[icol]=[getbiomass(icol,ikey,df,'Broad Group') for ikey in biomassDF['Key']]
 


# In[41]:


colListGenus


# In[42]:


for icol in colListGenus:
    biomassDF[icol]=[getbiomass(icol,ikey,df,'Genus species') for ikey in biomassDF['Key']]


# In[43]:


biomassDF


# In[ ]:


biomassDF.keys()


# In[ ]:


import netCDF4 as nc


# In[ ]:


#ftemp=nc.Dataset('/data/eolson/results/MEOPAR/NEMO-forcing-new/grid/mesh_mask201702.nc')


# In[ ]:


#ftemp.variables.keys()


# In[ ]:


#ftemp.variables['e3t_0']


# In[ ]:


#ftemp.variables['e3t_1d'][:]


# In[ ]:


fdict={'ptrc_T':1,'grid_T':1}
start_date = dt.datetime(2015,1,1)
end_date = dt.datetime(2016,12,31)
flen=1 # number of days per model output file. always 1 for 201905 and 201812 model runs
namfmt='nowcast' # for 201905 and 201812 model runs, this should always be 'nowcast'
# filemap is dictionary of the form variableName: fileType, where variableName is the name
# of the variable you want to extract and fileType designates the type of 
# model output file it can be found in (usually ptrc_T for biology, grid_T for temperature and 
# salinity)
filemap={'microzooplankton':'ptrc_T','mesozooplankton':'ptrc_T'}
# fdict is a dictionary mappy file type to its time resolution. Here, 1 means hourly output
# (1h file) and 24 means daily output (1d file). In certain runs, multiple time resolutions 
# are available
fdict={'ptrc_T':1,'grid_T':1}


# In[ ]:


PATH= '/results2/SalishSea/nowcast-green.201905/'


# In[ ]:


biomassDF.keys()


# In[ ]:


biomassDF.rename(columns={'Max Tow Depth (m)':'Z_lower','Min Tow Depth (m)':'Z_upper','Station Depth (m)':'Z'},inplace=True)


# In[ ]:


# define log transform function with slight shift to accommodate zero values
def logt(x):
  return np.log10(x+.001)


# In[ ]:


biomassDF['Year']=[ii.year for ii in biomassDF['dtUTC']]
biomassDF['YD']=et.datetimeToYD(biomassDF['dtUTC'])

biomassDF['Amphipoda']=biomassDF['Amphipoda']/8.5
biomassDF['Crabs']=biomassDF['Crabs']/8.5
biomassDF['Krill-Adult Juvenile']=biomassDF['Krill-Adult Juvenile']/8.5
biomassDF['Krill-Calyptopis']=biomassDF['Krill-Calyptopis']/8.5
biomassDF['Krill-Furcilia']=biomassDF['Krill-Furcilia']/8.5
biomassDF['Copepod']=biomassDF['Copepod']/8.5
biomassDF['Larvacea']=biomassDF['Larvacea']/8.5


#biomassDF['METRIDIA PACIFICA']=biomassDF['METRIDIA PACIFICA']/8.5
#biomassDF['NEOCALANUS PLUMCHRUS']=biomassDF['NEOCALANUS PLUMCHRUS']/8.5
#biomassDF['CALANOIDA']=biomassDF['CALANOIDA']/8.5
#biomassDF['EUCALANUS BUNGII']=biomassDF['EUCALANUS BUNGII']/8.5
#biomassDF['CALANUS PACIFICUS']=biomassDF['CALANUS PACIFICUS']/8.5
#biomassDF['CALANUS MARSHALLAE']=biomassDF['CALANUS MARSHALLAE']/8.5
#biomassDF['NEOCALANUS CRISTATUS']=biomassDF['NEOCALANUS CRISTATUS']/8.5
#biomassDF['EUCALANUS CALIFORNICUS']=biomassDF['EUCALANUS CALIFORNICUS']/8.5
#biomassDF['EUCALANUS']=biomassDF['EUCALANUS']/8.5
#biomassDF['METRIDIIDAE']=biomassDF['METRIDIIDAE']/8.5
#biomassDF['NEOCALANUS']=biomassDF['NEOCALANUS']/8.5

# In[ ]:





# In[ ]:


biomassDF


# In[ ]:


biomassDF.keys()


# In[ ]:


data=et.matchData(biomassDF,filemap,fdict,start_date,end_date,'nowcast',PATH,1,quiet=False,method='vertNet');


# In[ ]:


data


# In[ ]:


data.keys()


# In[ ]:





# In[ ]:





# In[ ]:


data['L10Amphipoda']=logt(data['Amphipoda'])
data['L10Crabs']=logt(data['Crabs'])
data['L10Krill-Adult Juvenile']=logt(data['Krill-Adult Juvenile'])
data['L10Krill-Calyptopis']=logt(data['Krill-Calyptopis'])
data['L10Krill-Furcilia']=logt(data['Krill-Furcilia'])
data['L10Copepod']=logt(data['Copepod'])
data['L10Larvacea']=logt(data['Larvacea'])
data['L10mod_microzooplankton']=logt(data['mod_microzooplankton'])
data['L10mod_mesozooplankton']=logt(data['mod_mesozooplankton'])


# In[ ]:


cm1=cmocean.cm.thermal
with nc.Dataset('/data/eolson/results/MEOPAR/NEMO-forcing-new/grid/bathymetry_201702.nc') as bathy:
    bathylon=np.copy(bathy.variables['nav_lon'][:,:])
    bathylat=np.copy(bathy.variables['nav_lat'][:,:])
    bathyZ=np.copy(bathy.variables['Bathymetry'][:,:])


# In[ ]:


fig, ax = plt.subplots(1,1,figsize = (6,6))
with nc.Dataset('/ocean/ksuchy/MOAD/NEMO-forcing/grid/bathymetry_201702.nc') as grid:
    viz_tools.plot_coastline(ax, grid, coords = 'map',isobath=.1)
colors=('teal','green','firebrick','darkorange','darkviolet','fuchsia',
        'royalblue','darkgoldenrod','mediumspringgreen','deepskyblue')
datreg=dict()
for ind, iregion in enumerate(data.Basin.unique()):
    datreg[iregion] = data.loc[data.Basin==iregion]
    ax.plot(datreg[iregion]['Lon'], datreg[iregion]['Lat'],'.',
            color = colors[ind], label=iregion)
ax.set_ylim(47,51)
ax.legend(bbox_to_anchor=[1,.6,0,0])
ax.set_xlim(-126, -120);
ax.set_title('Observation Locations');
ax.legend(bbox_to_anchor=(1.1, 1.05))


# In[ ]:


data.keys()


# In[ ]:


def byRegion(ax,obsvar,modvar,lims):
    PS=[]
    for ind, iregion in enumerate(data.Basin.unique()):
        #ax.plot(datreg[iregion]['Lon'], datreg[iregion]['Lat'],'.',
                #color = colors[ind], label=iregion)
        PS0=et.varvarPlot(ax,datreg[iregion],obsvar,modvar,
                          cols=(colors[ind],),lname=iregion)
        PS.append(PS0)
    l=ax.legend(handles=[ip[0][0] for ip in PS])
    ax.set_xlabel('Obs')
    ax.set_ylabel('Model')
    ax.plot(lims,lims,'k-',alpha=.5)
    ax.set_xlim(lims)
    ax.set_ylim(lims)
    ax.set_aspect(1)
    return PS,l


# In[ ]:


data['Month']=[ii.month for ii in data['dtUTC']]
JF=data.loc[(data.Month==1)|(data.Month==2)]
MAM=data.loc[(data.Month==3)|(data.Month==4)|(data.Month==5)]
JJA=data.loc[(data.Month==6)|(data.Month==7)|(data.Month==8)]
SOND=data.loc[(data.Month==9)|(data.Month==10)|(data.Month==11)|(data.Month==12)]


# In[ ]:


def bySeason(ax,obsvar,modvar,lims):
    for axi in ax:
        axi.plot(lims,lims,'k-')
        axi.set_xlim(lims)
        axi.set_ylim(lims)
        axi.set_aspect(1)
        axi.set_xlabel('Obs')
        axi.set_ylabel('Model')
    PS=et.varvarPlot(ax[0],JF,obsvar,modvar,cols=('crimson','darkturquoise','navy'))
    ax[0].set_title('Winter')
    PS=et.varvarPlot(ax[1],MAM,obsvar,modvar,cols=('crimson','darkturquoise','navy'))
    ax[1].set_title('Spring')
    PS=et.varvarPlot(ax[2],JJA,obsvar,modvar,cols=('crimson','darkturquoise','navy'))
    ax[2].set_title('Summer')
    PS=et.varvarPlot(ax[3],SOND,obsvar,modvar,cols=('crimson','darkturquoise','navy'))
    ax[3].set_title('Autumn')
    return 


# In[ ]:


fig, ax = plt.subplots(1,1,figsize = (16,7))     
PS,l=byRegion(ax,'L10Crabs','L10mod_mesozooplankton',(-2,2))
ax.set_title('Crabs ($\mu$M) By Region')
ax.legend(bbox_to_anchor=(1.1, 1.05))

fig, ax = plt.subplots(1,4,figsize = (16,3.3))
bySeason(ax,'L10Crabs','L10mod_mesozooplankton',(-2,2))



# In[ ]:





# In[ ]:


fig, ax = plt.subplots(1,1,figsize = (16,7))     
PS,l=byRegion(ax,'L10Krill-Adult Juvenile','L10mod_mesozooplankton',(-3,3))
ax.set_title('L10Krill-Adult Juveniles ($\mu$M) By Region')
ax.legend(bbox_to_anchor=(1.1, 1.05))

fig, ax = plt.subplots(1,4,figsize = (16,3.3))
bySeason(ax,'L10Krill-Adult Juvenile','L10mod_mesozooplankton',(-3,3))


# In[ ]:


fig, ax = plt.subplots(1,1,figsize = (16,7))     
PS,l=byRegion(ax,'L10Copepod','L10mod_mesozooplankton',(-1.5,1.5))
ax.set_title('L10Copepod($\mu$M) By Region')
ax.legend(bbox_to_anchor=(1.1, 1.05))

fig, ax = plt.subplots(1,4,figsize = (16,3.3))
bySeason(ax,'L10Copepod','L10mod_mesozooplankton',(-1.5,1.5))



# In[ ]:


fig, ax = plt.subplots(1,1,figsize = (16,7))     
PS,l=byRegion(ax,'L10Larvacea','L10mod_mesozooplankton',(-1.5,1.5))
ax.set_title('L10Larvaceans ($\mu$M) By Region')
ax.legend(bbox_to_anchor=(1.1, 1.05))

fig, ax = plt.subplots(1,4,figsize = (16,3.3))
bySeason(ax,'L10Larvacea','L10mod_mesozooplankton',(-1.5,1.5))


# In[ ]:


fig,ax=plt.subplots(2,4,figsize=(15,7))
fig.subplots_adjust(wspace=.4)
fig.subplots_adjust(hspace=.6)
ax=ax.flatten()

chplc=('Amphipoda', 'Crabs','Krill-Adult Juvenile', 'Krill-Calyptopis', 'Krill-Furcilia','Copepod','Larvacea')

mvar1=data['mod_microzooplankton']
mvar2=data['mod_mesozooplankton']

for ii in range(0,len(chplc)):
    ax[ii].plot(logt(data.loc[:,[chplc[ii]]].values),logt(mvar1),'.',ms=1,color='blue',label='Microzoop')
    ax[ii].plot(logt(data.loc[:,[chplc[ii]]].values),logt(mvar2),'.',ms=1,color='red',label='Mesozoop')
    ax[ii].set_ylabel('Model Class')
    ax[ii].set_xlabel(chplc[ii])
    ax[ii].set_title('log10[$\mu$M +.001]')
    ax[ii].plot((-3,1.5),(-3,1.5),'k-',alpha=.2)
    ax[ii].set_xlim((-3,1.5))
    ax[ii].set_ylim((-3,1.5))
    ax[ii].set_aspect(1)
ax[0].legend(loc=2, fontsize = 'small')


# In[ ]: