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

# # Develop `tracer_thalweg_and_surface` Figure Module
# 
# Development of functions for `nowcast.figures.research.tracer_thalweg_and_surface` web site figure module.
# 
# This is an example of developing the functions for a web site figure module in a notebook.
# It follows the function organization patterns described in
# [Creating a Figure Module](https://salishsea-nowcast.readthedocs.io/en/latest/figures/create_fig_module.html) docs.

# Notebooks like this should be developed in a
# [Nowcast Figures Development Environment](https://salishsea-nowcast.readthedocs.io/en/latest/figures/fig_dev_env.html)
# so that all of the necessary dependency packages are installed.
# The development has to be done on a workstation that has the Salish Sea Nowcast system `/results/` parition mounted.
# 
# If you choose to develop your figure in a notebook,
# the next step is to create a module like
# `nowcast.figures.research.tracer_thalweg_and_surface`
# and a notebook like
# [TestTracerThalwegAndSurfaceModule](https://nbviewer.jupyter.org/urls/bitbucket.org/salishsea/salishseanowcast/raw/tip/notebooks/figures/research/TestTracerThalwegAndSurface.ipynb)
# that tests it in the nowcast system context.

# In[1]:


import datetime
import os
from types import SimpleNamespace

import cmocean
import matplotlib.pyplot as plt
from matplotlib import gridspec
import netCDF4 as nc
import numpy as np
import scipy.io as sio

import salishsea_tools
from salishsea_tools import visualisations as vis
from salishsea_tools import tidetools, viz_tools

import nowcast
import nowcast.figures.website_theme


# In[2]:


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


# We know for sure that we need to develop:
# 
# * a `_prep_plot_data()` function
# * a `_prep_fig_axes()` function
# * a `make_figure()` function that the `make_plots` worker will call
# 
# Since the figure will have 2 axes
# (a vertical thalweg slice, and a horizontal surface slice)
# we also know that we need to develop 2 axes plotting functions:
# 
# * `_plot_tracer_thalweg()`
# * `_plot_tracer_surface()`

# ## `_prep_plot_data()` Function

# In[3]:


def _prep_plot_data(tracer_var, mesh_mask, depth_integrated):
    hr = 19
    sj, ej = 200, 800
    si, ei = 20, 395
    
    tracer_hr = tracer_var[hr]
    masked_tracer_hr = np.ma.masked_where(
        mesh_mask["tmask"][0, ...] == 0, tracer_hr)
    surface_hr = masked_tracer_hr[0, sj:ej, si:ei]
    
    if depth_integrated:
        grid_heights = mesh_mask.variables['e3t_1d'][:][0].reshape(
            tracer_hr.shape[0], 1, 1)
        height_weighted = masked_tracer_hr[:, sj:ej, si:ei] * grid_heights
        surface_hr = height_weighted.sum(axis=0)
        
    return SimpleNamespace(
        tracer_var=tracer_var,
        tracer_hr=tracer_hr,
        surface_hr=surface_hr,
        surface_j_limits=(sj, ej),
        surface_i_limits=(si, ei),
        thalweg_depth_limits=(0, 450),
        thalweg_length_limits=(0, 632),
    )


# In[4]:


bathy = nc.Dataset('/results/nowcast-sys/NEMO-forcing/grid/bathy_downonegrid2.nc')

mesh_mask = nc.Dataset('/results/nowcast-sys/NEMO-forcing/grid/mesh_mask_downbyone2.nc')

ptrc_T = nc.Dataset('/results/SalishSea/nowcast-green/25apr17/SalishSea_1h_20170425_20170425_ptrc_T.nc')

tracer_var = ptrc_T.variables['PHY2']


# Confirm that `_prep_plot_data()` is doing what we want.
# Lots of tests are possible,
# but we'll just confirm that the shapes of the data arrays to be plotted
# are what we expect.

# In[5]:


plot_data = _prep_plot_data(tracer_var, mesh_mask, depth_integrated=True)

print(plot_data.tracer_hr.shape)
print(plot_data.surface_hr.shape)


# ## `_prep_fig_axes() Function`

# In[6]:


def _prep_fig_axes(figsize, theme):
    fig = plt.figure(
        figsize=figsize, facecolor=theme.COLOURS['figure']['facecolor'])

    gs = gridspec.GridSpec(1, 2, width_ratios=[1.618, 1])

    ax_thalweg = fig.add_subplot(gs[0])
    ax_thalweg.set_axis_bgcolor(theme.COLOURS['axes']['background'])

    ax_surface = fig.add_subplot(gs[1])
    ax_surface.set_axis_bgcolor(theme.COLOURS['axes']['background'])

    return fig, (ax_thalweg, ax_surface)


# ## `_plot_tracer_thalweg()` Function
# 
# Because our figure is going to show 2 views of the tracer
# (a vertical slice along the thalweg,
# and a horizontal slice on the ocean surface)
# it would be nice to use the same contour colour levels if appropriate.
# But it won't be appropriate if the values contoured on the surface are
# depth averaged.
# It's also possible that the tracer concentration ranges will be different
# enough in the 2 slices that common contour levels don't make sense.
# 
# So,
# we come up with a function that calculates contour levels for the two axes
# and decides whether whether the levels are similar enough that one colour bar
# is sufficient for the figure, or if each axes requires one.

# In[7]:


def _calc_clevels(plot_data):
    """Calculates contour levels for the two axes and decides whether whether
    the levels are similar enough that one colour bar is sufficient for the 
    figure, or if each axes requires one.
    """
    percent_98_surf = np.percentile(plot_data.surface_hr.compressed(), 98)
    percent_2_surf = np.percentile(plot_data.surface_hr.compressed(), 2)

    percent_98_grid = np.percentile(
        np.ma.masked_values(plot_data.tracer_hr, 0).compressed(), 98)
    percent_2_grid = np.percentile(
        np.ma.masked_values(plot_data.tracer_hr, 0).compressed(), 2)

    overlap = (
        max(0, min(percent_98_surf, percent_98_grid)
            - max(percent_2_surf, percent_2_grid)))
    magnitude = (
        (percent_98_surf - percent_2_surf) + (percent_98_grid - percent_2_grid))
    if 2 * overlap / magnitude > 0.5:
        max_clevel = max(percent_98_surf, percent_98_grid)
        min_clevel = min(percent_2_surf, percent_2_grid)
        clevels_thalweg = np.arange(
            min_clevel, max_clevel, (max_clevel - min_clevel) / 20.0)
        clevels_surface = clevels_thalweg
        show_thalweg_cbar = False
    else:
        clevels_thalweg = np.arange(
            percent_2_grid, percent_98_grid,
            (percent_98_grid - percent_2_grid) / 20.0)
        clevels_surface = np.arange(
            percent_2_surf, percent_98_surf,
            (percent_98_surf - percent_2_surf) / 20.0)
        show_thalweg_cbar = True
    return clevels_thalweg, clevels_surface, show_thalweg_cbar


# In[8]:


def _plot_tracer_thalweg(ax, plot_data, bathy, mesh_mask, cmap, clevels):
    cbar = vis.contour_thalweg(
        ax, plot_data.tracer_hr, bathy, mesh_mask, clevels=clevels, cmap=cmap,
        thalweg_file='/results/nowcast-sys/tools/bathymetry/thalweg_working.txt',
        cbar_args={'fraction': 0.030, 'pad': 0.04, 'aspect': 45}
    )
    return cbar


# ## `_thalweg_axes_labels()` Function
# 
# We also put the code for labeling and prettifying the axes in a separate function.
# Along the way,
# we pull the code that we use to label either colour bar out into a separate function
# that both `_thalweg_axes_labels()` and `_surface_axes_labels()` can make use of.

# In[9]:


def _cbar_labels(cbar, contour_intervals, theme, label):
    cbar.set_ticks(contour_intervals)
    cbar.ax.axes.tick_params(labelcolor=theme.COLOURS['cbar']['tick labels'])
    cbar.set_label(
        label,
        fontproperties=theme.FONTS['axis'],
        color=theme.COLOURS['text']['axis'])


# In[10]:


def _thalweg_axes_labels(
    ax, plot_data, show_thalweg_cbar, clevels, cbar, theme
):
    ax.set_xlim(plot_data.thalweg_length_limits)
    ax.set_ylim(
        plot_data.thalweg_depth_limits[1], plot_data.thalweg_depth_limits[0])
    if show_thalweg_cbar:
        label = (
            f'{plot_data.tracer_var.long_name} [{plot_data.tracer_var.units}]')
        _cbar_labels(cbar, clevels[::2], theme, label)
    else:
        cbar.remove()
    ax.set_xlabel(
        'Distance along thalweg [km]', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    ax.set_ylabel(
        'Depth [m]', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    theme.set_axis_colors(ax)


# ## `_plot_tracer_surface()` Function

# In[11]:


def _plot_tracer_surface(ax, plot_data, cmap, clevels):
    x, y = np.meshgrid(
        np.arange(*plot_data.surface_i_limits, dtype=int),
        np.arange(*plot_data.surface_j_limits, dtype=int))
    mesh = ax.contourf(
        x, y, plot_data.surface_hr, levels=clevels, cmap=cmap, extend='both')
    cbar = plt.colorbar(mesh, ax=ax, fraction=0.034, pad=0.04, aspect=45)
    return cbar


# ## `_surface_axes_labels()` Function

# In[12]:


def _surface_axes_labels(
    ax, tracer_var, depth_integrated, clevels, cbar, theme
):
    cbar_units = (
        f'{tracer_var.units}*m' if depth_integrated
        else f'{tracer_var.units}')
    cbar_label = f'{tracer_var.long_name} [{cbar_units}]'
    _cbar_labels(cbar, clevels[::2], theme, cbar_label)
    ax.set_xlabel(
        'Grid x', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    ax.set_ylabel(
        'Grid y', color=theme.COLOURS['text']['axis'],
        fontproperties=theme.FONTS['axis'])
    ax.set_axis_bgcolor('burlywood')
    viz_tools.set_aspect(ax)
    theme.set_axis_colors(ax)


# ## `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[13]:


def make_figure(
    tracer_var, bathy, mesh_mask, cmap,  depth_integrated,
    figsize=(16, 9), theme=nowcast.figures.website_theme
):
    plot_data = _prep_plot_data(tracer_var, mesh_mask, depth_integrated)
    fig, (ax_thalweg, ax_surface) = _prep_fig_axes(figsize, theme)

    clevels_thalweg, clevels_surface, show_thalweg_cbar = _calc_clevels(
        plot_data)

    cbar_thalweg = _plot_tracer_thalweg(
        ax_thalweg, plot_data, bathy, mesh_mask, cmap, clevels_thalweg)
    _thalweg_axes_labels(
        ax_thalweg, plot_data, show_thalweg_cbar, clevels_thalweg,
        cbar_thalweg, theme)

    cbar_surface = _plot_tracer_surface(
        ax_surface, plot_data, cmap, clevels_surface)
    _surface_axes_labels(
        ax_surface, tracer_var, depth_integrated, clevels_surface, cbar_surface,
        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.

# In[14]:


get_ipython().run_cell_magic('timeit', '-n1 -r1', '\nfig = make_figure(\n    tracer_var, bathy, mesh_mask, cmap=cmocean.cm.algae, depth_integrated=True)\n')


# In[15]:


tracers = {
    'DOC': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'MICZ': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'MYRI': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'NH4': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'NO3': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'O2': {'cmap': cmocean.cm.oxy, 'depth_integrated': False},
    'PHY': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'PHY2': {'cmap': cmocean.cm.algae, 'depth_integrated': True},
    'POC': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'Si': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
    'bSi': {'cmap': cmocean.cm.matter, 'depth_integrated': False},
}
for tracer in tracers:
    tracer_var = ptrc_T.variables[tracer]
    fig = make_figure(
        tracer_var, bathy, mesh_mask, 
        cmap=tracers[tracer]['cmap'], 
        depth_integrated=tracers[tracer]['depth_integrated']
    )


# In[ ]: