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

# # Test `velocity_section_and_surface` Module
# 
# Render figure object produced by the `nowcast.figures.research.velocity_section_and_surface` module.
# 
# Set-up and function call replicates as nearly as possible what is done in the `nowcast.workers.make_plots` worker
# to help ensure that the module will work in the nowcast production context.

# This notebook must be run 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.
# Also, it must be run on a workstation that has the Salish Sea Nowcast system `/results/` parition mounted.

# In[1]:


# reload lets us pull in changes from modules that we edit
from importlib import reload


# In[2]:


import datetime
from glob import glob
import logging
import os
from pathlib import Path

import matplotlib

import arrow
import cmocean
from nemo_nowcast import NowcastWorker
import netCDF4 as nc
import scipy.io as sio

from nowcast import lib
from nowcast.figures.research import velocity_section_and_surface


# In[3]:


import io

import yaml


# In[4]:


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


# In[5]:


config = '''
    run_types:
        nowcast-green:
            bathymetry: /results/nowcast-sys/grid/bathymetry_201702.nc
            mesh_mask: /results/nowcast-sys/grid/mesh_mask201702.nc
        nowcast:
            bathymetry: /results/nowcast-sys/grid/bathymetry_201702.nc
            mesh_mask: /results/nowcast-sys/grid/mesh_mask201702.nc

    run:
        results_archive:
            nowcast-green: /results/SalishSea/nowcast-green/
            nowcast: /results/SalishSea/nowcast-blue/
'''

config = yaml.load(io.StringIO(config))


# In[6]:


run_type = 'nowcast'
plot_type = 'research'
run_date = arrow.get('2017-10-23')


# In[7]:


dmy = run_date.format('DDMMMYY').lower()
results_dir = Path(config['run']['results_archive'][run_type], dmy)
bathy = nc.Dataset(config['run_types'][run_type]['bathymetry'])
mesh_mask = nc.Dataset(config['run_types'][run_type]['mesh_mask'])


# In[8]:


yyyymmdd = run_date.format('YYYYMMDD')


# In[9]:


U = nc.Dataset(os.fspath(results_dir/f'SalishSea_1h_{yyyymmdd}_{yyyymmdd}_grid_U.nc'))
V = nc.Dataset(os.fspath(results_dir/f'SalishSea_1h_{yyyymmdd}_{yyyymmdd}_grid_V.nc'))
U_var = U.variables['vozocrtx']
V_var = V.variables['vomecrty']


# In[10]:


get_ipython().run_cell_magic('timeit', '-n1 -r1', '\n# Layout parameters\nsections = (450, 520, 680)\npos = (\n    (0.1, 0.35),\n    (0.4, 0.65),\n    (0.7, 0.95)\n)\nsection_lims = (\n    (235, 318, 0, 445),\n    (192, 277, 0, 445),\n    (127, 197, 0, 445),\n)\n\n# Make figure\nfig = velocity_section_and_surface.make_figure(\n    U_var, V_var, bathy, mesh_mask,\n    sections=sections, pos=pos, section_lims=section_lims\n)\n')


# In[ ]: