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,
)

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])

def winds_dataset(run_date, model_path):
    """Return the winds dataset for the current date from model_path."""
    year, month, day = run_date.strftime('%G'), run_date.strftime('%m'), run_date.strftime('%d')
    filename_pattern = 'ops_y{yyyy}m{mm}d{dd}.nc'
    filepaths = glob(os.path.join(model_path, filename_pattern.format(yyyy=year,mm=month,dd=day)))
    return nc.Dataset(filepaths[0])

#run_date = datetime.datetime(2015,1,8)
run_date = datetime.date.today()

# 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')

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

time_shift = datetime.timedelta(hours=-8)
hfmt = mdates.DateFormatter('%m/%d %H:%M')

MSL_DATUMS = {
    'Point Atkinson': 3.10, 'Victoria': 1.90,
    'Campbell River': 2.89, 'Patricia Bay': 2.30}

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

import scipy.io as sio

coastline = sio.loadmat('/ocean/rich/more/mmapbase/bcgeo/PNW.mat')

def draw_coast(ax, PNW_coastline):
    coast={}
    coast['lat'] = PNW_coastline['ncst'][:,1]
    coast['lon'] = PNW_coastline['ncst'][:,0]
    ax.plot(coast['lon'],coast['lat'],'-k',rasterized=True,markersize=1) 

def plot_threshold_website(grid_B, grid_T, model_path, scale=0.1,
                           PST=1, figsize = (18, 20)):
    """ Overview image for Salish Sea website. Plots a map of the Salish Sea with
    markers indicating extreme water at Point Atkinson, Victoria nd Campbell
    River. Also plots wind vectors averaged over 4 ours before the max ssh at
    Point Atkinson. Includes text boxes with max water level and timing.

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

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

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

    :arg scale: scale factor or wind arrows
    :type scale: float

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

    :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()

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

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

    # Wind time
    inds = figures.isolate_wind_timing('Point Atkinson', grid_T, grid_B,
                               model_path, t, 4, average=True)

    # Set up loop
    names = ['Point Atkinson', 'Campbell River', 'Victoria']

    # Set up Information
    max_sshs = {}; max_times = {}; max_winds = {}

    # Figure
    fig = plt.figure(figsize=figsize)
    gs = gridspec.GridSpec(2, 3, width_ratios=[1, 1, 1], height_ratios=[6, 1])
    gs.update(hspace=0.1, wspace=0.05)
    ax = fig.add_subplot(gs[0, :])
    ax1 = fig.add_subplot(gs[1, 0])
    ax2 = fig.add_subplot(gs[1, 1])
    ax3 = fig.add_subplot(gs[1, 2])

    # Map
    figures.plot_map(ax, grid_B)
    draw_coast(ax, coastline)
    viz_tools.plot_land_mask(ax, grid_B,color='#DBDEE1',coords='map')
    ax.set_xlim([-126.4,-121.3])
    ax.set_ylim([46.8,51.1])

    for name in names:
        # 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]

        # Get tides and ssh
        ttide = figures.get_tides(name)
        sdt = t_orig.replace(minute=0)
        edt = t_final + datetime.timedelta(minutes=30)
        ssh_corr = stormtools.correct_model(ssh_loc, ttide, sdt, edt)

        # Plot thresholds
        figures.plot_threshold_map(ax, ttide, ssh_corr, 'o', 55, 0.3, name)
        # Plot winds
        twind = figures.plot_wind_vector(ax, name, t_orig, t_final, model_path, inds, scale)

        # Information
        res = figures.compute_residual(ssh_loc, ttide, t_orig, t_final)
        [max_ssh, index_ssh, tmax, max_res, max_wind, ind_w] = figures.get_maxes(ssh_corr, t, res, lons[name], lats[name], model_path)
        max_sshs[name] = max_ssh
        max_times[name] = tmax
        max_winds[name] = max_wind

    # Add winds for other stations
    for name in ['Neah Bay', 'Cherry Point', 'Sandheads', 'Friday Harbor']:
        figures.plot_wind_vector(ax, name, t_orig, t_final, model_path, inds, scale)

    # Reference arrow
    ax.arrow(-121.5, 51, 0.*scale, -5.*scale,
              head_width=0.05, head_length=0.1, width=0.02,
              color='white',fc='DarkMagenta', ec='black')
    ax.text(-121.6, 50.95, "Reference: 5 m/s", rotation=90, fontsize=14)

    # Location labels
    ax.text(-125.6, 48.1, 'Pacific Ocean', fontsize=13)
    ax.text(-122.8, 50.1, 'British Columbia', fontsize=13)
    ax.text(-123.8, 47.3, 'Washington \n State', fontsize=13)

    ax.text(-122.3, 47.6, 'Puget Sound', fontsize=13)
    ax.text(-124.7, 48.45, 'Strait of Juan de Fuca', fontsize=13, rotation=-18)
    ax.text(-123.95, 49.25, 'Strait of \n Georgia', fontsize=13, rotation=-2)

    ax.text(-123.1, 49.4, 'Point \n Atkinson', fontsize=20)
    ax.text(-125.9, 50.05, 'Campbell \n River', fontsize=20)
    ax.text(-123.5, 48.2, 'Victoria', fontsize=20)

    # Figure format
    t = (twind[0]+PST*time_shift).strftime('%A, %B %d, %Y')
    ax.set_title('Marine and Atmospheric Conditions\n {time}'.format(time=t), **title_font)
    fig.patch.set_facecolor('#2B3E50')
    figures.axis_colors(ax, 'gray')

    # Citation
    t1 = (twind[0]+PST*time_shift).strftime('%Y/%m/%d %H:%M')
    t2 = (twind[-1]+PST*time_shift).strftime('%Y/%m/%d %H:%M')
    timezone = PST*'[PST]' + abs((PST-1))*'[UTC]'
    ax.text(0.4,-0.25,
      'Wind vectors averaged over: {time1} to {time2} {tzone}'.format(time1=t1,time2=t2,tzone=timezone),
        horizontalalignment='left',
        verticalalignment='top',
        transform=ax.transAxes, color = 'white',fontsize=14)
    ax.text(0.4, -0.29,
      'Modelled winds are from the High Resolution Deterministic Prediction System\nof Environment Canada: https://weather.gc.ca/grib/grib2_HRDPS_HR_e.html',
        horizontalalignment = 'left',
        verticalalignment = 'top',
        transform=ax.transAxes, color = 'white',fontsize=14)

    # Information_box
    axs = [ax1, ax2, ax3]
    for ax, name in zip (axs, names):
        plt.setp(ax.spines.values(), visible=False)
        ax.xaxis.set_visible(False); ax.yaxis.set_visible(False)
        figures.axis_colors(ax, 'blue')
        display_time=(max_times[name]+PST*time_shift).strftime('%H:%M')


        ax.text(0.05, 0.9, name, fontsize=20,
                horizontalalignment='left', verticalalignment='top', color = 'w')
        ax.text(0.05, 0.7, 'Maximum Water Level: {:.2f} m'.format(max_sshs[name]+MSL_DATUMS[name]),fontsize=15,
                horizontalalignment='left',verticalalignment='top', color = 'w')
        ax.text(0.05, 0.3, 'Time: {time} {tzone}'.format(time=display_time,tzone=timezone),
	        fontsize=15,horizontalalignment='left', verticalalignment='top', color = 'w')
        ax.text(0.05, 0.5,'Wind speed: {:.1f} m/s'.format(float(max_winds[name])),fontsize=15,
                 horizontalalignment='left',verticalalignment='top', color = 'w')

    return fig

from mpl_toolkits.basemap import Basemap
from matplotlib.patches import Polygon

def basemap(grid_B, grid_T, model_path, scale=0.1, PST=1):
    fig=plt.figure()
    fig, ax = plt.subplots(1,1,figsize=(20,15))

    #Map of whole region
    lcc_values=dict(resolution='l',projection='lcc', lat_1=40,lat_2=50,lat_0=45,lon_0=-135,
                    llcrnrlon=-130,llcrnrlat=47, urcrnrlon=-120,urcrnrlat=51,ax=ax)
    base_SoG=Basemap(**lcc_values)

    ## define parallels and meridians to draw.
    parallels=np.arange(-90, 90, 5)
    meridians=np.arange(0, 360, 10)
    base_SoG.drawparallels(parallels, labels=[1, 0, 0, 0],fontsize=10, latmax=90)
    base_SoG.drawmeridians(meridians, labels=[0, 0, 0, 1],fontsize=10, latmax=90)

    #coastlines
    out=base_SoG.drawcoastlines(linewidth=1.5, linestyle='solid', color='k')
    
    #fill
    base_SoG.drawmapboundary(fill_color='LightBlue')
    base_SoG.fillcontinents(color='Tan')
    
    
    #Plotting
    # Stations information
    [lats, lons] = figures.station_coords()
    # Bathymetry
    bathy, X, Y = tidetools.get_bathy_data(grid_B)
    # Time range
    t_orig, t_final, t = figures.get_model_time_variables(grid_T)
    # Wind time
    inds = figures.isolate_wind_timing('Point Atkinson', grid_T, grid_B,
                               model_path, t, 4, average=True)
    # Set up loop
    names = ['Point Atkinson', 'Campbell River', 'Victoria']
    # Set up Information
    max_sshs = {}; max_times = {}; max_winds = {}
    for name in names:
        # 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]
        # Get tides and ssh
        ttide = figures.get_tides(name)
        sdt = t_orig.replace(minute=0)
        edt = t_final + datetime.timedelta(minutes=30)
        ssh_corr = stormtools.correct_model(ssh_loc, ttide, sdt, edt)
        # Plot thresholds
        figures.plot_threshold_map(ax, ttide, ssh_corr, 'o', 55, 0.3, name)
        # Plot winds
        twind = figures.plot_wind_vector(ax, name, t_orig, t_final, model_path, inds, scale)
        # Information
        res = figures.compute_residual(ssh_loc, ttide, t_orig, t_final)
        [max_ssh, index_ssh, tmax, max_res, max_wind, ind_w] = figures.get_maxes(ssh_corr, t, res, lons[name], lats[name], model_path)
        max_sshs[name] = max_ssh
        max_times[name] = tmax
        max_winds[name] = max_wind
    # Add winds for other stations
    for name in ['Neah Bay', 'Cherry Point', 'Sandheads', 'Friday Harbor']:
        figures.plot_wind_vector(ax, name, t_orig, t_final, model_path, inds, scale)
    

    return fig