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

# # Develop `second_narrows_current` Figure Module
# 
# Development of functions for `nowcast.figures.fvcom.second_narrows_current` web site figure module.

# In[1]:


from pathlib import Path
import shlex
import subprocess
from types import SimpleNamespace

import arrow
import matplotlib.dates
import matplotlib.pyplot as plt
import numpy
import requests
import xarray
from salishsea_tools import unit_conversions

from nowcast.figures import shared
import nowcast.figures.website_theme


# In[2]:


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


# ## `_prep_plot_data()` Function

# In[3]:


def _prep_plot_data(place, fvcom_stn_dataset, obs_dataset):
    # FVCOM velocity component datasets
    stations = [
        name.decode().strip().split(maxsplit=1)[1] 
        for name in fvcom_stn_dataset.name_station.values
    ]
    fvcom_u = fvcom_stn_dataset.u.isel(siglay=0, station=stations.index(place))
    fvcom_u.attrs.update({
        'long_name': 'u Velocity',
        'units': 'm/s',
    })
    fvcom_v = fvcom_stn_dataset.v.isel(siglay=0, station=stations.index(place))
    # FVCOM current speed and direction
    fvcom_speed = numpy.sqrt(fvcom_u**2 + fvcom_v**2)
    fvcom_speed.name = 'fvcom_current_speed'
    fvcom_speed.attrs.update({
        'long_name': 'Current Speed',
        'units': 'm/s',
        'label': 'Model',
    })
    direction = numpy.arctan2(fvcom_v, fvcom_u)
    fvcom_dir = numpy.rad2deg(direction + (direction < 0) * 2 * numpy.pi)
    fvcom_dir.name = 'fvcom_current_direction'
    fvcom_dir.attrs.update({
        'long_name': 'Current To Direction',
        'units': '°CCW from East',
        'label': 'Model',
    })
    # HADCP observations dataset
    obs = xarray.Dataset(
        data_vars={
            "speed": (
                {"time": obs_dataset.data_vars["s.time"].size}, 
                obs_dataset.data_vars["s.speed"], 
                obs_dataset.data_vars["s.speed"].attrs,
            ),
            "direction": (
                {"time": obs_dataset.data_vars["s.time"].size}, 
                obs_dataset.data_vars["s.direction"],
                obs_dataset.data_vars["s.direction"].attrs,
            ),
        },
        coords={"time": obs_dataset.data_vars["s.time"].values},
    )
    fvcom_time_range = slice(
        str(fvcom_speed.time.values[0]), str(fvcom_speed.time.values[-1])
    )
    obs_speed = obs.speed.sel(time=fvcom_time_range)
    obs_speed.attrs.update({"label": "HADCP Observed"})
    rad_ccw_from_east = numpy.deg2rad(90 - obs.direction.sel(time=fvcom_time_range))
    obs_dir = numpy.rad2deg(rad_ccw_from_east + (rad_ccw_from_east<0)*2*numpy.pi)
    obs_dir.attrs.update({"label": "HADCP Observed"})
    shared.localize_time(fvcom_u)
    shared.localize_time(fvcom_speed)
    shared.localize_time(fvcom_dir)
    shared.localize_time(obs_speed)
    shared.localize_time(obs_dir)
    return SimpleNamespace(
        fvcom_u=fvcom_u,
        fvcom_speed=fvcom_speed,
        fvcom_dir=fvcom_dir,
        obs_speed=obs_speed,
        obs_dir=obs_dir,
    )


# Test `_prep_plot_data()`:

# In[4]:


run_type = 'nowcast'
run_date = arrow.get('2018-10-25')

fvcom_stn_dataset_path_tmpl = (
    '/opp/fvcom/{run_type}/{ddmmmyy}/vhfr_low_v2_station_timeseries.nc'
)
if run_type == 'nowcast':
    fvcom_stn_dataset_path = Path(
        fvcom_stn_dataset_path_tmpl.format(
            run_type=run_type, ddmmmyy=run_date.format("DDMMMYY").lower())
    )
else:
    nowcast_dataset_path = Path(
        fvcom_stn_dataset_path_tmpl.format(
            run_type='nowcast', ddmmmyy=run_date.format("DDMMMYY").lower())
    )
    forecast_dataset_path = Path(
        fvcom_stn_dataset_path_tmpl.format(
            run_type='forecast', ddmmmyy=run_date.format("DDMMMYY").lower())
    )
    fvcom_stn_dataset_path = Path('/tmp/vhfr_low_v2_station_timeseries_forecast.nc')
    cmd = (
        f'ncrcat -O {nowcast_dataset_path} {forecast_dataset_path} '
        f'-o {fvcom_stn_dataset_path}'
    )
    subprocess.check_output(shlex.split(cmd))
cmd = (
    f'ncrename -O -v siglay,sigma_layer -v siglev,sigma_level '
    f'{fvcom_stn_dataset_path} /tmp/{fvcom_stn_dataset_path.name}')
subprocess.check_output(shlex.split(cmd))
fvcom_stn_dataset = xarray.open_dataset(f'/tmp/{fvcom_stn_dataset_path.name}')

obs_dataset = xarray.open_dataset(
    "https://salishsea.eos.ubc.ca/erddap/tabledap/ubcVFPA2ndNarrowsCurrent2sV1"
)

plot_data = _prep_plot_data('2nd Narrows', fvcom_stn_dataset, obs_dataset)


# In[5]:


print(plot_data.fvcom_u)
plot_data.fvcom_u.plot()


# In[6]:


print(plot_data.fvcom_speed)
plot_data.fvcom_speed.plot()


# In[7]:


print(plot_data.obs_speed)
plot_data.obs_speed.plot()


# In[8]:


print(plot_data.fvcom_dir)
plot_data.fvcom_dir.plot()


# In[9]:


print(plot_data.obs_dir)
plot_data.obs_dir.plot()


# ## `_prep_fig_axes()` Function

# In[10]:


def _prep_fig_axes(figsize, theme):
    fig, (ax_speed, ax_dir, ax_u) = plt.subplots(
        3, 1, figsize=figsize, facecolor=theme.COLOURS['figure']['facecolor']
    )
    ax_speed = {'mps': ax_speed}
    ax_speed['knots'] = ax_speed['mps'].twinx()
    ax_u = {'mps': ax_u}
    ax_u['knots'] = ax_u['mps'].twinx()
    fig.autofmt_xdate()
    return fig, (ax_speed, ax_dir, ax_u)


# ## `_plot_current_speed_time_series()` Function

# In[11]:


def _plot_current_speed_time_series(ax, plot_data, theme):
    plot_data.obs_speed.plot(
        ax=ax["mps"],
        marker=".",
        linestyle="None",
        label=plot_data.obs_speed.attrs['label'],
        markerfacecolor=theme.COLOURS['time series']['2nd Narrows observed current speed'],
        markeredgecolor=theme.COLOURS['time series']['2nd Narrows observed current speed'],
    )
    plot_data.fvcom_speed.plot(
        ax=ax['mps'],
        linewidth=2,
        color=theme.COLOURS['time series']['2nd Narrows model current speed'],
        label=plot_data.fvcom_speed.attrs['label'],
        alpha=0.8,
    )


# ## `_current_speed_axes_labels()` Function

# In[12]:


def _current_speed_axes_labels(ax, plot_data, theme):
    ax['mps'].set_title(
        'Current at 2nd Narrows',
        fontproperties=theme.FONTS['axes title'],
        color=theme.COLOURS['text']['axes title']
    )
    mps_limits = numpy.array((0, 5))
    ax['mps'].set_ylabel(
        f'{plot_data.fvcom_speed.attrs["long_name"]} '
        f'[{plot_data.fvcom_speed.attrs["units"]}]',
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis']
    )
    ax['mps'].set_ylim(mps_limits)
    ax['knots'].set_ylabel(
        f'{plot_data.fvcom_speed.attrs["long_name"]} [knots]',
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis']
    )
    ax['knots'].set_ylim(unit_conversions.mps_knots(mps_limits))
    ax['mps'].legend(loc='best')
    ax['mps'].grid(axis='both')
    for k in ax:
        theme.set_axis_colors(ax[k])


# ## `_plot_current_direction_time_series()` Function

# In[13]:


def _plot_current_direction_time_series(ax, plot_data, theme):
    plot_data.obs_dir.plot(
        ax=ax,
        marker=".",
        linestyle="None",
        label=plot_data.obs_speed.attrs['label'],
        markerfacecolor=theme.COLOURS['time series']['2nd Narrows observed current direction'],
        markeredgecolor=theme.COLOURS['time series']['2nd Narrows observed current direction'],
    )
    plot_data.fvcom_dir.plot(
        ax=ax,
        linewidth=2,
        color=theme.COLOURS['time series']['2nd Narrows model current direction'],
        label=plot_data.fvcom_speed.attrs['label'],
        alpha=0.8,
    )


# ## `_current_direction_axes_labels()` Function

# In[14]:


def _current_direction_axes_labels(ax, plot_data, theme):
    ax.set_ylim(0, 360)
    ax.set_yticks((0, 45, 90, 135, 180, 225, 270, 315, 360))
    ax.set_yticklabels(('E', 'NE', 'N', 'NW', 'W', 'SW', 'S', 'SE', 'E'))
    ax.set_ylabel(
        f'{plot_data.fvcom_dir.attrs["long_name"]} ',
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis']
    )
    ax.legend(loc='best')
    ax.grid(axis='both')
    theme.set_axis_colors(ax)


# ## `_plot_u_velocity_time_series()` Function

# In[15]:


def _plot_u_velocity_time_series(ax, plot_data, theme):
    plot_data.fvcom_u.plot(
        ax=ax['mps'],
        linewidth=2,
        color=theme.COLOURS['time series']['2nd Narrows model current speed'],
    )


# ## `_u_velocity_axes_labels()` Function

# In[16]:


def _u_velocity_axes_labels(ax, plot_data, theme):
    ax['mps'].set_xlabel(
        f'Time [{plot_data.fvcom_u.attrs["tz_name"]}]',
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis']
    )
    ax['mps'].xaxis.set_major_formatter(matplotlib.dates.DateFormatter('%d%b %H:%M'))
    mps_limits = numpy.array((-4, 4))
    ax['mps'].set_ylabel(
        f'{plot_data.fvcom_u.attrs["long_name"]} '
        f'[{plot_data.fvcom_u.attrs["units"]}]',
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis']
    )
    ax['mps'].set_ylim(mps_limits)
    ax['knots'].set_ylabel(
        f'{plot_data.fvcom_u.attrs["long_name"]} [knots]',
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis']
    )
    ax['knots'].set_ylim(unit_conversions.mps_knots(mps_limits))
    ax['mps'].grid(axis='both')
    for k in ax:
        theme.set_axis_colors(ax[k])


# ## `make_figure()` Function
# 
# This is is the function that will be called by the `nowcast.workers.make_plots` worker to return a `matplotlib.figure.Figure` object.

# In[17]:


def make_figure(
    place, fvcom_stn_dataset, obs_dataset,
    figsize=(16, 9),
    theme=nowcast.figures.website_theme,
):
    plot_data = _prep_plot_data(place, fvcom_stn_dataset, obs_dataset)
    fig, (ax_speed, ax_dir, ax_u) = _prep_fig_axes(figsize, theme)
    _plot_current_speed_time_series(ax_speed, plot_data, theme)
    _current_speed_axes_labels(ax_speed, plot_data, theme)
    _plot_current_direction_time_series(ax_dir, plot_data, theme)
    _current_direction_axes_labels(ax_dir, plot_data, theme)
    _plot_u_velocity_time_series(ax_u, plot_data, theme)
    _u_velocity_axes_labels(ax_u, plot_data, theme)
    return fig


# ## Render the Figure
# 
# The `%%timeit` cell magic lets us keep an eye on how log the figure takes to process.
# Setting `-n1 -r1` prevents it from processing the figure more than once
# as it might try to do to generate better statistics.

# ### Nowcast Figure

# In[18]:


get_ipython().run_cell_magic('timeit', '-n1 -r1', '\nfrom importlib import reload\nfrom nowcast.figures import website_theme\nfrom salishsea_tools import places\nreload(website_theme)\nreload(places)\n\nrun_type = \'nowcast\'\nrun_date = arrow.get(\'2018-10-25\')\n\nfvcom_stn_dataset_path_tmpl = (\n    \'/opp/fvcom/{run_type}/{ddmmmyy}/vhfr_low_v2_station_timeseries.nc\'\n)\nif run_type == \'nowcast\':\n    fvcom_stn_dataset_path = Path(\n        fvcom_stn_dataset_path_tmpl.format(\n            run_type=run_type, ddmmmyy=run_date.format("DDMMMYY").lower())\n    )\nelse:\n    nowcast_dataset_path = Path(\n        fvcom_stn_dataset_path_tmpl.format(\n            run_type=\'nowcast\', ddmmmyy=run_date.format("DDMMMYY").lower())\n    )\n    forecast_dataset_path = Path(\n        fvcom_stn_dataset_path_tmpl.format(\n            run_type=\'forecast\', ddmmmyy=run_date.format("DDMMMYY").lower())\n    )\n    fvcom_stn_dataset_path = Path(\'/tmp/vhfr_low_v2_station_timeseries_forecast.nc\')\n    cmd = (\n        f\'ncrcat {nowcast_dataset_path} {forecast_dataset_path} \'\n        f\'-o {fvcom_stn_dataset_path}\'\n    )\n    subprocess.check_output(shlex.split(cmd))\ncmd = (\n    f\'ncrename -O -v siglay,sigma_layer -v siglev,sigma_level \'\n    f\'{fvcom_stn_dataset_path} /tmp/{fvcom_stn_dataset_path.name}\')\nsubprocess.check_output(shlex.split(cmd))\nfvcom_stn_dataset = xarray.open_dataset(f\'/tmp/{fvcom_stn_dataset_path.name}\')\n\nobs_dataset = xarray.open_dataset(\n    "https://salishsea.eos.ubc.ca/erddap/tabledap/ubcVFPA2ndNarrowsCurrent2sV1"\n)\n\nfig = make_figure(\'2nd Narrows\', fvcom_stn_dataset, obs_dataset)\n')


# ### Forecast Figure

# In[19]:


get_ipython().run_cell_magic('timeit', '-n1 -r1', '\nfrom importlib import reload\nfrom nowcast.figures import website_theme\nfrom salishsea_tools import places\nreload(website_theme)\nreload(places)\n\nrun_type = \'forecast\'\nrun_date = arrow.get(\'2018-10-25\')\n\nfvcom_stn_dataset_path_tmpl = (\n    \'/opp/fvcom/{run_type}/{ddmmmyy}/vhfr_low_v2_station_timeseries.nc\'\n)\nif run_type == \'nowcast\':\n    fvcom_stn_dataset_path = Path(\n        fvcom_stn_dataset_path_tmpl.format(\n            run_type=run_type, ddmmmyy=run_date.format("DDMMMYY").lower())\n    )\nelse:\n    nowcast_dataset_path = Path(\n        fvcom_stn_dataset_path_tmpl.format(\n            run_type=\'nowcast\', ddmmmyy=run_date.format("DDMMMYY").lower())\n    )\n    forecast_dataset_path = Path(\n        fvcom_stn_dataset_path_tmpl.format(\n            run_type=\'forecast\', ddmmmyy=run_date.format("DDMMMYY").lower())\n    )\n    fvcom_stn_dataset_path = Path(\'/tmp/vhfr_low_v2_station_timeseries_forecast.nc\')\n    cmd = (\n        f\'ncrcat -O {nowcast_dataset_path} {forecast_dataset_path} -o {fvcom_stn_dataset_path}\')\n    subprocess.check_output(shlex.split(cmd))\ncmd = (\n    f\'ncrename -O -v siglay,sigma_layer -v siglev,sigma_level \'\n    f\'{fvcom_stn_dataset_path} /tmp/{fvcom_stn_dataset_path.name}\')\nsubprocess.check_output(shlex.split(cmd))\nfvcom_stn_dataset = xarray.open_dataset(f\'/tmp/{fvcom_stn_dataset_path.name}\')\n\nobs_dataset = xarray.open_dataset(\n    "https://salishsea.eos.ubc.ca/erddap/tabledap/ubcVFPA2ndNarrowsCurrent2sV1"\n)\n\nfig = make_figure(\'2nd Narrows\', fvcom_stn_dataset, obs_dataset)\n')


# In[ ]: