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

# # Read ADCIRC output from Zarr (cloud-optimized NetCDF)

# In[1]:


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

from dask.distributed import Client, progress, LocalCluster
from dask_kubernetes import KubeCluster
import xarray as xr
import s3fs


# Start a dask cluster to crunch the data

# In[2]:


cluster = KubeCluster.from_yaml('/home/jovyan/custom-worker-template.yaml')
cluster.scale(30);


# In[3]:


client = Client(cluster)


# In[9]:


client


# In[ ]:


# jetstream s3
# url='https://iu.jetstream-cloud.org:8080'
# fs = s3fs.S3FileSystem(client_kwargs=dict(endpoint_url=url), anon=True)
# s3map = s3fs.S3Map('rsignell/adcirc_test01', s3=fs)


# In[10]:


# AWS s3
fs = s3fs.S3FileSystem(anon=True)
s3map = s3fs.S3Map('rsignell/adcirc_test01', s3=fs)


# In[11]:


ds = xr.open_zarr(s3map)


# In[12]:


ds


# In[13]:


ds['zeta']


# In[14]:


ds['zeta'].nbytes/1.e9


# In[ ]:





# In[ ]:





# In[15]:


ds['zeta'].max(dim='time')


# In[16]:


max_var = ds['zeta'].max(dim='time').persist()
progress(max_var)


# # Visualize data on mesh using Datashader

# In[4]:


import numpy as np
import datashader as dshade
import holoviews as hv
import geoviews as gv
import cartopy.crs as ccrs

from holoviews.operation.datashader import datashade, rasterize
from colorcet import cm_n
from matplotlib.cm import jet

datashade.precompute = True

hv.extension('bokeh')
get_ipython().run_line_magic('opts', 'Image RGB VectorField [width=800 height=600]')


# In[5]:


import warnings
warnings.filterwarnings("ignore")


# In[6]:


import pandas as pd


# In[17]:


v = np.vstack((ds['x'], ds['y'], max_var)).T
verts = pd.DataFrame(v, columns=['x','y','z'])


# In[18]:


points = gv.operation.project_points(gv.Points(verts, vdims=['z']))


# In[19]:


tris = pd.DataFrame(ds['element'].values.astype('int')-1, columns=['v0','v1','v2'])


# In[20]:


tiles = gv.WMTS('https://maps.wikimedia.org/osm-intl/{Z}/{X}/{Y}@2x.png')
value = 'max water level'
label = '{} (m)'.format(value)
trimesh = gv.TriMesh((tris, points), label=label)


# In[21]:


get_ipython().run_cell_magic('opts', 'Image [colorbar=True] (cmap=jet)', "\nerror_meshes = rasterize(trimesh,aggregator=dshade.mean('z'))\ntiles * error_meshes\n")


# In[22]:


get_ipython().run_cell_magic('time', '', "ds['zeta'][:,800000].plot()\n")


# In[ ]: