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

# In[1]:


import xarray as xr
import panel as pn
import hvplot.xarray
from geoviews import tile_sources as gvts
import cf_xarray 
import intake
import pandas as pd


# In[2]:


update_button = pn.widgets.Button(name='Load latest data', button_type='success')


# In[3]:


def plot(ds, var, base_map):
    extra_dims = list(ds[var].dims[:-2])
    da = ds[var].cf.isel(T=slice(-time_vals,-1)).unify_chunks()
    if len(da.shape) == 4:
        mesh = da.hvplot.quadmesh(x=da.cf['longitude'].name, y=da.cf['latitude'].name, 
                              rasterize=True, geo=True, title=var, attr_labels=False, 
                              fields={da.cf['Z'].name: {'default': float(da.cf['Z'].values[-1])}},
                              groupby=extra_dims, cmap='turbo', width=600, grid=True,
                              height=600).opts(alpha=0.7, data_aspect=None, 
                              hooks=[set_tools],
                              active_tools=['pan', 'box_zoom'])
    else:
        mesh = da.hvplot.quadmesh(x=da.cf['longitude'].name, y=da.cf['latitude'].name, 
                              rasterize=True, geo=True, title=var, attr_labels=False, 
                              groupby=extra_dims, cmap='turbo', width=600, grid=True,
                              height=600).opts(alpha=0.7, data_aspect=None, 
                              hooks=[set_tools], active_tools=['pan', 'box_zoom'])
    return pn.panel(mesh * base_map, widgets={'ocean_time': pn.widgets.DatetimePicker})


# In[4]:


url = 'https://mghp.osn.xsede.org/rsignellbucket1/catalogs/ioos_intake_catalog.yml'

show_vars = ['Hwave', 'temp', 'salt', 'zeta', 
             'Uwind', 'Vwind', 'u', 'v', 'ubar', 'vbar', 'Dwave', 'sand_06', 'bstrcwmax']

init_var = 'temp'

time_vals = 48


# In[5]:


def set_tools(plot, element):
    plot.state.toolbar.active_inspect = None


# In[6]:


def map_select(event=None):
    cat = intake.open_catalog(url)
    ds = cat['COAWST-Archive'].to_dask()

    var_select = pn.widgets.Select(name='COAWST Variables:', options=sorted(show_vars), value=init_var)
    base_map_select = pn.widgets.Select(name='Basemap:', options=gvts.tile_sources, value=gvts.OSM)
    
    plot_app = pn.bind(plot, ds=ds, var=var_select, base_map=base_map_select)
    return pn.Column(
        var_select,
        base_map_select,
        pn.panel(plot_app, loading_indicator=True),
    )


# In[7]:


app = pn.Column(
    update_button,
    pn.panel(pn.bind(map_select, update_button), loading_indicator=True),
)


# In[8]:


app.servable('COAWST Explorer')