from __future__ import division

import datetime
from glob import glob
import os

from IPython.core.display import HTML
import netCDF4 as nc

from salishsea_tools.nowcast import figures

%matplotlib inline

def results_dataset(period, grid, results_dir):
    """Return the results dataset for period (e.g. 1h or 1d)
    and grid (e.g. grid_T, grid_U) from results_dir.
    """
    filename_pattern = 'SalishSea_{period}_*_{grid}.nc'
    filepaths = glob(os.path.join(results_dir, filename_pattern.format(period=period, grid=grid)))
    return nc.Dataset(filepaths[0])

run_date = datetime.datetime(2014,12,21)


# Results dataset location
results_home = '/data/dlatorne/MEOPAR/SalishSea/nowcast/'
results_dir = os.path.join(results_home, run_date.strftime('%d%b%y').lower())
# model winds
model_path = '/ocean/sallen/allen/research/MEOPAR/Operational/'

grid_T_hr = results_dataset('1h', 'grid_T', results_dir)

bathy = nc.Dataset('/data/nsoontie/MEOPAR/NEMO-forcing/grid/bathy_meter_SalishSea2.nc')

from __future__ import division

from cStringIO import StringIO
import datetime
import glob
import os

import arrow
from dateutil import tz
import matplotlib
from matplotlib.backends import backend_agg as backend
import matplotlib.cm as cm
import matplotlib.dates as mdates
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt
import netCDF4 as nc
import numpy as np
import pandas as pd
import requests
from scipy import interpolate as interp

from salishsea_tools import (
    nc_tools,
    viz_tools,
    stormtools,
    tidetools,
)


# Plotting colors
model_c = 'MediumBlue'
observations_c = 'DarkGreen'
predictions_c = 'MediumVioletRed'
stations_c = cm.rainbow(np.linspace(0, 1, 7))

# Time shift for plotting in PST
time_shift = datetime.timedelta(hours=-8)
hfmt = mdates.DateFormatter('%m/%d %H:%M')

# Font format
title_font = {
    'fontname': 'Bitstream Vera Sans', 'size': '20', 'color': 'black',
    'weight': 'medium'
    }
axis_font = {'fontname': 'Bitstream Vera Sans', 'size': '16'}

# Average mean sea level calculated over 1983-2001
# (To be used to centre model output about mean sea level)
MSL_DATUMS = {
    'Point Atkinson': 3.10, 'Victoria': 1.90,
    'Campbell River': 2.89, 'Patricia Bay': 2.30}

back_c = '#14657E'

def axis_colors(ax, plot):
  """ Formats the background colour of plots and colours
  of labels.

  :arg ax: Axis to be formatted.
  :type ax: axis object

  :arg plot: Keyword for background needed for plot.
  :type plot: string

  :returns: axis format
  """

  labels_c = 'white'
  ticks_c = 'white'
  spines_c = 'white'

  if plot == 'blue':
      ax.set_axis_bgcolor(back_c)
  if plot == 'gray':
      ax.set_axis_bgcolor('#DBDEE1')
  if plot == 'white':
      ax.set_axis_bgcolor('white')

  ax.xaxis.label.set_color(labels_c), ax.yaxis.label.set_color(labels_c)
  ax.tick_params(axis='x', colors=ticks_c), ax.tick_params(axis='y', colors=ticks_c)
  ax.spines['bottom'].set_color(spines_c), ax.spines['top'].set_color(spines_c)
  ax.spines['left'].set_color(spines_c), ax.spines['right'].set_color(spines_c)
  ax.title.set_color('white')

def PA_tidal_predictions(grid_T,  PST=1, MSL=0, figsize=(20, 8)):
    """ Plots the tidal cycle at Point Atkinson during a 4 week period
    centred around the simulation start date.

    This function assumes that a tidal prediction file exists in a
    specific directory.
    Tidal predictions were calculated with ttide based on a time series
    from 2013.
    Plots are of predictions caluclated with all consituents.

    :arg grid_T: Hourly tracer results dataset from NEMO.
    :type grid_T: :class:`netCDF4.Dataset`

    :arg PST: Specifies if plot should be presented in PST.
    1 = plot in PST, 0 = plot in UTC.
    :type PST: 0 or 1

    :arg MSL: Specifies if the plot should be centred about mean sea level.
    1=centre about MSL, 0=centre about 0.
    :type MSL: 0 or 1

    :arg figsize: Figure size (width, height) in inches.
    :type figsize: 2-tuple

    :returns: matplotlib figure object instance (fig).
    """

    # Time range
    t_orig,t_end,t_nemo=figures.get_model_time_variables(grid_T)
    timezone=PST*'[PST]' + abs((PST-1))*'[UTC]'

    # Axis limits are set as 2 weeks before and after start date.
    ax_start = t_orig - datetime.timedelta(weeks=2)
    ax_end = t_orig + datetime.timedelta(weeks=2)
    ylims=[-3,3]

    # Figure
    fig = plt.figure(figsize=figsize)
    ax = fig.add_subplot(1, 1, 1)
    fig.patch.set_facecolor(back_c)
    fig.autofmt_xdate()
    ttide=figures.plot_tides(ax,'Point Atkinson',PST,MSL,'black')

    # Line indicating current date
    ax.plot([t_orig +time_shift*PST,t_orig+time_shift*PST],ylims,'-r',lw=2)
    ax.plot([t_end+time_shift*PST,t_end+time_shift*PST],ylims,'-r',lw=2)

    # Axis
    ax.set_xlim([ax_start+time_shift*PST,ax_end+time_shift*PST])
    ax.set_ylim(ylims)
    ax.set_title('Tidal Predictions at Point Atkinson: ' + t_orig.strftime('%d-%b-%Y'),**title_font)
    ax.set_ylabel('Sea Surface Height [m]',**axis_font)
    ax.set_xlabel('Time {}'.format(timezone),**axis_font)
    ax.grid()
    axis_colors(ax, 'gray')
    ax.text(1., -0.2,
            'Tidal predictions calculated with t_tide: http://www.eos.ubc.ca/~rich/#T_Tide',
        horizontalalignment='right',
        verticalalignment='top',
        transform=ax.transAxes, color = 'white')

def compare_tidalpredictions_maxSSH(
    grid_T, grid_B, model_path, PST=1, MSL=0, name='Point Atkinson',
    figsize=(20, 15),
):
    """Plots a map for sea surface height when it was at its maximum
    at Point Atkinson and compares modelled water levels to tidal
    predications over one day.

    It is assummed that the tidal predictions were calculated ahead of
    time and stored in a very specific location.
    The tidal predictions were calculated with all constituents using
    ttide based on a time series from 2013.
    The corrected model takes into account errors resulting in using
    only 8 constituents.
    The residual is calculated as corrected model - tides
    (with all constituents).

    :arg grid_T: Hourly tracer results dataset from NEMO.
    :type grid_T: :class:`netCDF4.Dataset`

    :arg grid_B: Bathymetry dataset for the Salish Sea NEMO model.
    :type grid_B: :class:`netCDF4.Dataset`

    :arg model_path: The directory where the model wind files are stored.
    :type model_path: string

    :arg PST: Specifies if plot should be presented in PST.
    1 = plot in PST, 0 = plot in UTC.
    :type PST: 0 or 1

    :arg MSL: Specifies if the plot should be centred about mean sea level.
    1=centre about MSL, 0=centre about 0.
    :type MSL: 0 or 1

    :arg name: Name of station.
    :type name: string

    :arg figsize:  Figure size (width, height) in inches.
    :type figsize: 2-tuple

    :returns: matplotlib figure object instance (fig).
    """

    # Stations information
    [lats, lons] = figures.station_coords()

    # Time range
    t_orig,t_final,t=figures.get_model_time_variables(grid_T)
    tzone=PST*'[PST]' + abs((PST-1))*'[UTC]'

    # Bathymetry
    bathy, X, Y = tidetools.get_bathy_data(grid_B)

    # Get sea surface height
    [j,i]=tidetools.find_closest_model_point(lons[name],lats[name],X,Y,bathy,allow_land=False)
    ssh = grid_T.variables['sossheig']
    ssh_loc = ssh[:,j,i]

    # Figure
    fig = plt.figure(figsize=figsize)
    fig.patch.set_facecolor(back_c)
    gs = gridspec.GridSpec(3, 2, width_ratios=[2,1])
    gs.update(wspace=0.13, hspace=0.2)
    ax0 = fig.add_subplot(gs[0, 0])  # information box
    axis_colors(ax0, 'blue')
    plt.setp(ax0.spines.values(), visible=False)  # hide axes for information box
    ax0.xaxis.set_visible(False); ax0.yaxis.set_visible(False)
    ax1 = fig.add_subplot(gs[1, 0])  # sea surface height
    ax2 = fig.add_subplot(gs[:, 1])  # map
    ax3 = fig.add_subplot(gs[2, 0])  # residual

    # Sea surface height plot
    ttide=figures.plot_tides(ax1,name,PST,MSL)
    ssh_corr=figures.plot_corrected_model(ax1,t,ssh_loc,ttide,t_orig,t_final,
			PST,MSL,MSL_DATUMS[name])
    ax1.plot(t+PST*time_shift,ssh_loc,'--',c=model_c,linewidth=1,label='Model')

    # Compute residual
    res = figures.compute_residual(ssh_loc,ttide,t_orig,t_final)

    # Find maximim sea surface height and timing
    max_ssh,index,tmax,max_res,max_wind,ind_w = figures.get_maxes(ssh_corr, t, res,
                            lons[name], lats[name], model_path)
    ax0.text(0.05, 0.9, name, fontsize=20,
             horizontalalignment='left',
             verticalalignment='top', color = 'white')
    ax0.text(0.05, 0.75,
         'Max SSH: {:.2f} metres above mean sea level'.format(max_ssh),
             fontsize=15, horizontalalignment='left',
             verticalalignment='top', color = 'white')
    ax0.text(0.05, 0.6,
      'Time of max: {time} {timezone}'.format(time=tmax +PST*time_shift,
           timezone=PST*'[PST]' + abs((PST-1))*'[UTC]'),
             fontsize=15, horizontalalignment='left',
             verticalalignment='top', color = 'white')
    ax0.text(0.05, 0.45,
         'Residual: {:.2f} metres'.format(max_res),
             fontsize=15, horizontalalignment='left',
             verticalalignment='top', color = 'white')
    ax0.text(0.05, 0.3,
         'Wind speed: {:.1f} m/s'.format(float(max_wind)),
             fontsize=15, horizontalalignment='left',
             verticalalignment='top', color = 'white')

    # Mark point for maximum ssh
    ax1.plot(tmax+PST*time_shift, max_ssh, color='white', marker='o',
             markersize=10, markeredgewidth=3, label='Maximum SSH')

    # Axis for sea surface height plot
    ax1.set_xlim(t_orig+PST*time_shift,t_final+PST*time_shift)
    ax1.set_ylim([-3,3])
    ax1.set_title('Hourly Sea Surface Height at ' + name + ': ' + (t_orig).strftime('%d-%b-%Y'),**title_font)
    ax1.set_xlabel('Time {}'.format(tzone),**axis_font)
    ax1.set_ylabel('Water Levels wrt MSL (m)',**axis_font)
    ax1.grid()
    ax1.legend(loc = 0, numpoints = 1)
    axis_colors(ax1, 'gray')
    ax1.xaxis.set_major_formatter(hfmt)

    # Plot Residual
    ax3.plot(t +PST*time_shift,res,'-k',linewidth=2,label='Residual')

    # Axis for residual plot
    ax3.set_xlim(t_orig+PST*time_shift,t_final+PST*time_shift)
    ax3.set_ylim([-1,1])
    ax3.set_xlabel('Time {}'.format(tzone),**axis_font)
    ax3.set_ylabel('Residual (m)',**axis_font)
    ax3.set_yticks(np.arange(-1.0,1.25,0.25))
    ax3.grid()
    ax3.legend(loc = 0, numpoints = 1)
    axis_colors(ax3, 'gray')
    ax3.xaxis.set_major_formatter(hfmt)
    fig.autofmt_xdate()

    # Map of sea surface height
    cs = [-1,-0.5,0.5,1, 1.5,1.6,1.7,1.8,1.9,2,2.1,2.2,2.4,2.6]
    ssh_max_field = np.ma.masked_values(ssh[index], 0)
    mesh=ax2.contourf(ssh_max_field,cs,cmap='nipy_spectral',extend='both',alpha=0.6)
    ax2.contour(ssh_max_field,cs,colors='k',linestyles='--')

    cbar = fig.colorbar(mesh,ax=ax2)
    cbar.set_ticks(cs)
    plt.setp(plt.getp(cbar.ax.axes, 'yticklabels'), color='w')
    cbar.set_label('[m]', color='white')

    ax2.set_title('Sea Surface Height: ' + (tmax+PST*time_shift).strftime('%d-%b-%Y, %H:%M'),**title_font)
    ax2.set_xlabel('X Index',**axis_font)
    ax2.set_ylabel('Y Index',**axis_font)
    ax2.grid()

    axis_colors(ax2, 'white')
    viz_tools.plot_coastline(ax2,grid_B)
    viz_tools.plot_land_mask(ax2,grid_B,color='burlywood')
    ax2.plot(i, j, marker='o', color='white', markersize=10,
             markeredgewidth=3)

def plot_thresholds_all(grid_T, grid_B, model_path, PST=1, MSL=1, figsize=(20, 10)):
  """Plots sea surface height over one day with respect to warning thresholds.

  This function applies only to Point Atkinson, Campbell River, and Victoria.
  There are three different warning thresholds.
  The locations of stations are colored depending on the threshold in
  which they fall: green, yellow, red.

  :arg grid_T: Hourly tracer results dataset from NEMO.
  :type grid_T: :class:`netCDF4.Dataset`

  :arg grid_B: Bathymetry dataset for the Salish Sea NEMO model.
  :type grid_B: :class:`netCDF4.Dataset`

  :arg model_path: The directory where the model wind files are stored.
  :type model_path: string

  :arg PST: Specifies if plot should be presented in PST.
  1 = plot in PST, 0 = plot in UTC.
  :type PST: 0 or 1

  :arg MSL: Specifies if the plot should be centred about mean sea level.
  1=centre about MSL, 0=centre about 0.
  :type MSL: 0 or 1

  :arg figsize:  Figure size (width, height) in inches.
  :type figsize: 2-tuple

  :returns: matplotlib figure object instance (fig).
  """

  # Figure
  fig = plt.figure(figsize=figsize)
  fig.patch.set_facecolor(back_c)
  gs = gridspec.GridSpec(1, 3, width_ratios=[1,1,1])
  gs.update(wspace=0.15, hspace=0.1)

  # Map of region
  fig2 = plt.figure(figsize=figsize)
  fig2.patch.set_facecolor(back_c)
  ax0=fig2.add_subplot(1,1,1)
  figures.plot_map(ax0, grid_B)
  ax0.set_xlim(-125.4, -122.2)
  ax0.set_ylim(48, 50.3)
  ax0.set_title('Station Locations',**title_font)
  axis_colors(ax0, 'gray')
  fig.suptitle('Hourly Sea Surface Height', fontsize=20, color='white')

 # Stations information
  [lats, lons] = figures.station_coords()

  # Bathymetry
  bathy, X, Y = tidetools.get_bathy_data(grid_B)

  m = np.arange(3)
  names = ['Point Atkinson', 'Campbell River', 'Victoria']
  extreme_sshs = [5.61,5.35,3.76]

  for M, name, extreme_ssh in zip(m, names, extreme_sshs):

     # Get sea surface height
     [j,i]=tidetools.find_closest_model_point(lons[name],lats[name],X,Y,bathy,allow_land=False)
     ssh = grid_T.variables['sossheig']
     ssh_loc = ssh[:,j,i]

     # Time range
     t_orig,t_final,t=figures.get_model_time_variables(grid_T)
     tzone=PST*'[PST]' + abs((PST-1))*'[UTC]'

     # Sea surface height plot
     ax = fig.add_subplot(gs[0,M])

     # Plot tides, corrected model and original model
     if name =='Point Atkinson':
       figures.plot_PA_observations(ax,PST)
     ttide=figures.plot_tides(ax,name,PST,MSL)
     ssh_corr=figures.plot_corrected_model(ax,t,ssh_loc,ttide,t_orig,t_final,
			PST,MSL,MSL_DATUMS[name])
     ax.plot(t+PST*time_shift,ssh_loc+MSL_DATUMS[name]*MSL,'--',
			c=model_c,linewidth=1,label='Model')

     # Axis
     ax.set_xlim(t_orig+PST*time_shift,t_final+PST*time_shift)
     ax.set_ylim([-1,6])
     ax.set_title(name + ': ' + (t_orig).strftime('%d-%b-%Y'), **title_font)
     ax.set_xlabel('Time {}'.format(tzone),**axis_font)
     ax.set_ylabel('Water Level above Chart Datum (m)',**axis_font)
     ax.grid()
     axis_colors(ax, 'gray')
     ax.xaxis.set_major_formatter(hfmt)
     fig.autofmt_xdate()

     # Map
     bbox_args = dict(boxstyle='square',facecolor='white',alpha=0.8)
     ax0.annotate(name,(lons[name]-0.05,lats[name]-0.18),fontsize=15,
				color='black',bbox=bbox_args)

     # Define thresholds in sea surface height plots
     [max_tides, mid_tides, extreme_ssh] = figures.plot_threshold_map(ax0,ttide, ssh_corr, 'D', 25, 1.0, name)

     # Plot thresholds in sea surface height plots
     ax.axhline(y=max_tides,color='Gold',linewidth=2,ls='solid',label='Maximum tides')
     ax.axhline(y=mid_tides,color='Red',linewidth=2,ls='solid',label='Extreme water')
     ax.axhline(y=extreme_ssh,color='DarkRed',linewidth=2,ls='solid',label='Historical maximum')

     # Legend
     if M == 2:
	   legend = ax.legend(bbox_to_anchor=(1.5, 1.1), loc='upper right', borderaxespad=0.,
				prop={'size':15}, title=r'Legend')
	   legend.get_title().set_fontsize('20')

  # Citation
  ax0.text(0.3 , -0.1,
        'Tidal predictions calculated with t_tide: http://www.eos.ubc.ca/~rich/#T_Tide \nObserved water levels from Fisheries and Oceans, Canada \nvia Scott Tinis at stormsurgebc.ca',
        horizontalalignment='left',
        verticalalignment='top',
        transform=ax0.transAxes, color = 'white')

def winds_at_max_ssh(grid_T, grid_B, model_path, station, figsize=(20, 15)):
  """ Plots winds at individual stations 4 hours before the
  maxmimum sea surface height at Point Atkinson.

  If that data is not available then the plot is generated at the start of the simulation.
  This function applies to stations at Campbell River, Point Atkinson, Victoria,
  Cherry Point, Neah Bay, and Friday Harbor.

  :arg grid_T: Hourly tracer results dataset from NEMO.
  :type grid_T: :class:`netCDF4.Dataset`

  :arg grid_B: Bathymetry dataset for the Salish Sea NEMO model.
  :type grid_B: :class:`netCDF4.Dataset`

  :arg model_path: The directory where the model files are stored.
  :type model_path: string

  :arg station: Name of the station.
  :type station: string

  :arg figsize:  Figure size (width, height) in inches.
  :type figsize: 2-tuple

  :returns: matplotlib figure object instance (fig).
  """

  # Stations information
  [lats, lons] = figures.station_coords()

  # Map
  fig = plt.figure(figsize=figsize)
  ax = fig.add_subplot(1, 1, 1)
  fig.patch.set_facecolor(back_c)
  figures.plot_map(ax, grid_B)

  # Arrow scale
  scale = 0.1

  # Time range
  [t_orig,t_final,t] = figures.get_model_time_variables(grid_T)

  # Reference arrow
  ax.arrow(-123, 50., 5.*scale, 0.*scale,
              head_width=0.05, head_length=0.1, width=0.02,
              color='white',fc='DarkMagenta', ec='black')
  ax.text(-123, 50.1, "5 m/s")

  # Condition if plotting all stations or a single station
  if station == 'all':
        names = ['Neah Bay', 'Victoria', 'Friday Harbor', 'Cherry Point', 'Sandheads', 'Point Atkinson', 'Campbell River']
        colors=stations_c
  else:
        names=[station]
        colors=['DarkMagenta']

  # Indices for plotting wind vectors
  inds = figures.isolate_wind_timing('Point Atkinson',grid_T,grid_B,model_path,t,4,average=False)
  
  # Loop through all stations to plot arrows and markers
  for name, station_c in zip (names, colors):
    plot_time=figures.plot_wind_vector(ax, name, t_orig, t_final, model_path, inds, scale)
    ax.plot(lons[name], lats[name], marker='D',
	    color=station_c, markersize=10, markeredgewidth=2,label=name) 

  # Time for title and legend
  plot_time=(plot_time[0]+time_shift).strftime('%d-%b-%Y %H:%M')
  legend = ax.legend(numpoints=1, bbox_to_anchor=(-0.1, 1.05), loc=2, borderaxespad=0.,
						prop={'size':15}, title=r'Stations')
  legend.get_title().set_fontsize('20')
  ax.set_title('Modelled winds at \n {time} [PST]'.format(time=plot_time),**title_font)

  # Citation
  ax.text(0.4,-0.05,
    'Modelled winds are from the High Resolution Deterministic Prediction System of Environment Canada.\nhttps://weather.gc.ca/grib/grib2_HRDPS_HR_e.html',
        horizontalalignment='left',
        verticalalignment='top',
        transform=ax.transAxes, color = 'white')

  axis_colors(ax, 'gray')