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

# Testing Dask Distributed With Tides
# ====
# 

# In[1]:


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

import xarray as xr
import matplotlib.pyplot as plt

# from dask.dot import dot_graph
from dask.diagnostics import Profiler, ResourceProfiler, CacheProfiler
from dask.diagnostics import visualize

from bokeh.io import output_notebook
output_notebook()


# In[2]:


from dask.distributed import Client
client = Client(scheduler_file='scheduler.json')
client


# ### Open an Xarray Dataset

# In[3]:


ds = xr.open_dataset('/cxfs/projects/usgs/hazards/cmgp/woodshole/aaretxabaleta/projects/GSB_tides_55nb/ocean_his_gsb_tides_55nb.nc',
                       engine='netcdf4',  chunks={'eta_u':40, 'xi_rho':40}, decode_times=False)


# In[4]:


import netCDF4
nc = netCDF4.Dataset('/cxfs/projects/usgs/hazards/cmgp/woodshole/aaretxabaleta/projects/GSB_tides_55nb/ocean_his_gsb_tides_55nb.nc')


# In[5]:


nc.variables.keys()


# In[6]:


dt, n, m = nc['zeta'].shape
print(dt,n,m)


# In[7]:


t = nc['ocean_time'][:]/(3600*24)


# In[8]:


t[:10]


# In[9]:


from utide import solve
from matplotlib.dates import date2num
import numpy as np

lat = 35.0147

#time = date2num(ds['ocean_time'].to_pydatetime())


# In[10]:


get_ipython().run_cell_magic('time', '', "acoef = solve(t, nc['zeta'][:,10,10], 0*t, lat, \n               trend=False, nodal=False, Rayleigh_min=0.95, method='ols',\n               conf_int='linear')\nval = acoef['Lsmaj']\n")


# In[11]:


acoef.keys()


# In[12]:


acoef['name']


# In[13]:


import dask.array as da
from dask import delayed


# In[14]:


usolve = delayed(solve)


# In[15]:


z = nc['zeta']


# In[16]:


get_ipython().run_cell_magic('time', '', "nsub = 2\ncoefs = [usolve(t, z[:,j*nsub,i*nsub], t*0.0, lat, \n               trend=False, nodal=False, Rayleigh_min=0.95, method='ols',\n               conf_int='linear') for j in range(int(n/nsub)) for i in range(int(m/nsub))]\n")


# In[17]:


arrays = [da.from_delayed(coef['Lsmaj'],           # Construct a small Dask array
                          dtype=val.dtype,   # for every lazy value
                          shape=val.shape)
          for coef in coefs]


# In[18]:


stack = da.stack(arrays, axis=0)                # Stack all small Dask arrays into one


# In[19]:


stack


# In[20]:


get_ipython().run_cell_magic('time', '', 'amps = stack.compute()\n')


# In[21]:


m2amp = amps[:,0].reshape((int(n/nsub),int(m/nsub)))


# In[22]:


get_ipython().run_line_magic('matplotlib', 'inline')
import matplotlib.pyplot as plt


# In[28]:


plt.figure(figsize=(12,8))
plt.pcolormesh(m2amp)
plt.colorbar()
plt.title('M2 Elevation Amplitude');


# In[ ]: