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

# In[1]:


import numpy as np
import matplotlib.pyplot as plt
import netCDF4 as nc
import datetime as dt
from salishsea_tools import evaltools as et, places
import xarray as xr

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


# #### Example 1: load a time series at one location by accessing netCDF4 files stored on /results or /results2

# In[2]:


get_ipython().run_cell_magic('time', '', "start= dt.datetime(2017,3,1)\nend=dt.datetime(2017,4,1) # the code called below (evaltools.index_model_files) includes the end date \n                          # in the values returned\nbasedir='/results/SalishSea/nowcast-green.201812/'\nnam_fmt='nowcast'\nflen=1 # files contain 1 day of data each\nftype= 'ptrc_T' # load bio files\ntres=24 # 1: hourly resolution; 24: daily resolution   <- try changing to 1 and loading hourly data\nflist=et.index_model_files(start,end,basedir,nam_fmt,flen,ftype,tres)\n# flist contains paths: file pathes; t_0 timestemp of start of each file; t_n: timestamp of start of next file\nprint(flist)\n")


# In[3]:


# reminder of variable names in ptrc files:
with nc.Dataset(flist.loc[0,['paths']].values[0]) as ff: #  <-when you access elements of a pandas array, sometimes
                                                         #    you get an array output, even if it only contains one
                                                         #    element. To get the element rather than the array 
                                                         #     containing it, use [0]
    print(ff.variables.keys())
    # also grab time reference:
    torig=dt.datetime.strptime(ff.variables['time_centered'].time_origin,'%Y-%m-%d %H:%M:%S')
    print('time origin:',torig)


# In[4]:


# get model i,j of location S3 from places
ij,ii=places.PLACES['S3']['NEMO grid ji']
ik=0 # choose surface level

nlen= int(tres*24/flen) # number of data points per file


# In[5]:


get_ipython().run_cell_magic('time', '', "# create empty numpy arrays of shape of desired timeseries (1 dimension, length of flist)\ntt= np.zeros((flist.shape[0]*nlen,),dtype=object) # array to hold times\nmicZ= np.zeros((flist.shape[0]*nlen,)) # array to hold microzo conc\ndiat= np.zeros((flist.shape[0]*nlen,)) # array to hold diatom conc\nfor ind, row in flist.iterrows():\n    with nc.Dataset(row['paths']) as ff:\n        tt[ind*nlen:(ind+1)*nlen]=[torig+dt.timedelta(seconds=xx) for xx in ff.variables['time_centered'][:]]\n        micZ[ind*nlen:(ind+1)*nlen]=ff.variables['microzooplankton'][:,ik,ij,ii]\n        diat[ind*nlen:(ind+1)*nlen]=ff.variables['diatoms'][:,ik,ij,ii]\n")


# In[6]:


get_ipython().run_cell_magic('time', '', "fig,ax=plt.subplots(1,1,figsize=(12,3))\nax.plot(tt,diat,'c-',label='diatoms')\nax.plot(tt,micZ,'-',color='darkorange',label='microzooplankton')\nax.legend(loc=2);\nax.set_ylabel('Concentration ($\\mu$M N)')\nax.set_xlim(tt[0],tt[-1])\n")


# ### repeat with hourly data over a shorter interval

# In[7]:


start= dt.datetime(2017,3,1)
end=dt.datetime(2017,3,5) # the code called below (evaltools.index_model_files) includes the end date 
                          # in the values returned
tres=1 # 1: hourly resolution; 24: daily resolution   <- try changing to 1 and loading hourly data
flist=et.index_model_files(start,end,basedir,nam_fmt,flen,ftype,tres)
# flist contains paths: file pathes; t_0 timestemp of start of each file; t_n: timestamp of start of next file
flist


# In[8]:


# get model i,j of location S3 from places
ij,ii=places.PLACES['S3']['NEMO grid ji']
ik=0 # choose surface level

nlen= int(tres*24/flen) # number of data points per file


# In[9]:


get_ipython().run_cell_magic('time', '', "# create empty numpy arrays of shape of desired timeseries (1 dimension, length of flist)\ntt= np.zeros((flist.shape[0]*nlen,),dtype=object) # array to hold times\nmicZ= np.zeros((flist.shape[0]*nlen,)) # array to hold microzo conc\ndiat= np.zeros((flist.shape[0]*nlen,)) # array to hold diatom conc\nfor ind, row in flist.iterrows():\n    with nc.Dataset(row['paths']) as ff:\n        tt[ind*nlen:(ind+1)*nlen]=[torig+dt.timedelta(seconds=xx) for xx in ff.variables['time_centered'][:]]\n        micZ[ind*nlen:(ind+1)*nlen]=ff.variables['microzooplankton'][:,ik,ij,ii]\n        diat[ind*nlen:(ind+1)*nlen]=ff.variables['diatoms'][:,ik,ij,ii]\n")


# In[10]:


get_ipython().run_cell_magic('time', '', "fig,ax=plt.subplots(1,1,figsize=(12,3))\nax.plot(tt,diat,'c-',label='diatoms')\nax.plot(tt,micZ,'-',color='darkorange',label='microzooplankton')\nax.legend(loc=2);\nax.set_ylabel('Concentration ($\\mu$M N)')\nax.set_xlim(tt[0],tt[-1])\n")


# In[ ]: