#!/usr/bin/env python
"""
  map.py

  State Estimation and Analysis for PYthon

  Utilities for dealing with basemap plotting. These routnes are simply
  abstractions over the existing basemap to make it quicker for generating
  basemap plots and figures.

    Examples
    -------

    Assume you have longitude, latitude, and sst values:

    >>> m=seapy.mapping.map(llcrnrlon=lon[0,0],llcrnrlat=lat[0,0],
    >>>     urcrnrlon=lon[-1,-1],urcrnrlat=lat[-1,-1],dlat=2,dlon=2)
    >>> m.pcolormesh(lon,lat,sst,vmin=22,vmax=26,cmap=plt.cm.bwr)
    >>> m.land()
    >>> m.colorbar(label="Sea Surface Temp [$^\circ$C]",cticks=[22,23,24,25,26])
    >>> m.ax.patch.set_facecolor("aqua")
    >>> m.ax.patch.set_alpha(1)
    >>> m.fig.patch.set_alpha(0.0)
    >>> m.fig.savefig("sst.png",dpi=100)


  Written by Brian Powell on 9/4/14
  Copyright (c)2017 University of Hawaii under the BSD-License.
"""
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from seapy.model import asgrid


def gen_coastline(lon, lat, bathy, depth=0):
    """
    Given lon, lat, and bathymetry, generate vectors of line segments
    of the coastline. This can be exported to matlab (via savemat) to be
    used with the 'editmask' routine for creating grid masks.

    Input
    -----
    lon : array,
        longitudes of bathymetry locations
    lat : array,
        latitudes of bathymetry locations
    bathy : array,
        bathymetry (negative for ocean, positive for land) values
    depth : float,
        depth to use as the definition of the coast

    Returns
    -------
    lon : ndarray,
        vector of coastlines, separated by nan (matlab-style)
    lat : ndarray,
        vector of coastlines, separated by nan (matlab-style)
    """
    CS = plt.contour(lon, lat, bathy, [depth - 0.25, depth + 0.25])
    lon = list()
    lat = list()
    for col in CS.collections:
        for path in col.get_paths():
            lon.append(path.vertices[:, 0])
            lon.append(np.nan)
            lat.append(path.vertices[:, 1])
            lat.append(np.nan)
    return (np.hstack(lon), np.hstack(lat))


class map(object):

    def __init__(self, grid=None, llcrnrlon=-180, llcrnrlat=-40, urcrnrlon=180,
                 urcrnrlat=40, proj='lcc', resolution='c', figsize=(8., 6.),
                 dlat=1, dlon=2, fig=None, ax=None, fill_color="aqua"):
        """
        map class for abstracting the basemap methods for quick and easy creation
        of geographically referenced data figures


        Parameters
        ----------
        grid: seapy.model.grid or string, optional:
            grid to use to define boundaries
        llcrnrlon: float, optional
            longitude of lower, left corner
        llcrnrlat: float, optional
            latitude of lower, left corner
        urcrnrlon: float, optional
            longitude of upper, right corner
        urcrnrlat: float, optional
            latitude of upper, right corner
        proj: string, optional
            projection to use for map
        resolution: character
            resolution to use for coastline, etc. From Basemap:
            'c' (crude), 'l' (low), 'i' (intermediate),
            'h' (high), 'f' (full), or None
        figsize: list, optional
            dimensions to use for creation of figure
        dlat: float, optional
            interval to mark latitude lines (e.g., if dlat=0.5 every 0.5deg mark)
        dlon: float, optional
            interval to mark longitude lines (e.g., if dlon=0.5 every 0.5deg mark)
        fig: matplotlib.pyplot.figure object, optional
            If you want to plot on a pre-configured figure, pass the figure object
            along with the axis object.
        ax: matplotlib.pyplot.axis object, optional
            If you want to plot on a pre-configured figure, pass the axis object
            along with the figure object.
        fill_color: string, optional
            The color to use for the axis background

        Returns
        -------
        None

        """
        if grid is not None:
            grid = asgrid(grid)
            llcrnrlat = np.min(grid.lat_rho)
            urcrnrlat = np.max(grid.lat_rho)
            llcrnrlon = np.min(grid.lon_rho)
            urcrnrlon = np.max(grid.lon_rho)

        self.basemap = Basemap(llcrnrlon=llcrnrlon, llcrnrlat=llcrnrlat,
                               urcrnrlon=urcrnrlon, urcrnrlat=urcrnrlat,
                               projection=proj,
                               lat_0=urcrnrlat - (urcrnrlat - llcrnrlat) / 2.,
                               lon_0=urcrnrlon - (urcrnrlon - llcrnrlon) / 2.,
                               resolution=resolution, area_thresh=0.0, ax=ax)

        self.figsize = figsize
        self.dlon = dlon
        self.dlat = dlat
        self.fig = fig
        self.ax = ax
        self.fill_color = fill_color
        reset = True if fig is None else False
        self.new_figure(reset=reset)

    def new_figure(self, fill_color=None, reset=False):
        """
        Create or update a figure for plotting

        Parameters
        ----------
        fill_color: string, optional
           Color to fill the background of the axes with
        reset: bool, optional
           Reset the figure
        """
        if reset:
            if self.ax:
                self.ax.set_axis_off()
            if self.fig:
                plt.close(self.fig)

        if self.fig is None or self.ax is None:
            self.fig = plt.figure(figsize=self.figsize)
            self.ax = self.fig.add_axes([-0.01, 0.25, 1.01, 0.7])

        if fill_color is None:
            fill_color = self.fill_color

        self.basemap.drawmapboundary(fill_color=fill_color)
        # Create the longitude lines
        nticks = int((self.basemap.urcrnrlon - self.basemap.llcrnrlon) /
                     self.dlon)
        md = np.mod(self.basemap.llcrnrlon, self.dlon)
        if md:
            slon = self.basemap.llcrnrlon + self.dlon - md
        else:
            slon = self.basemap.llcrnrlon
            nticks += 1
        lon_lines = np.arange(nticks) * self.dlon + slon
        self.basemap.drawmeridians(lon_lines, color="0.5",
                                   linewidth=0.25, dashes=[1, 1, 0.1, 1],
                                   labels=[0, 0, 0, 1], fontsize=12)

        # Create the latitude lines
        nticks = int((self.basemap.urcrnrlat - self.basemap.llcrnrlat) /
                     self.dlat)
        md = np.mod(self.basemap.llcrnrlat, self.dlat)
        if md:
            slat = self.basemap.llcrnrlat + self.dlat - md
        else:
            slat = self.basemap.llcrnrlat
            nticks += 1
        lat_lines = np.arange(nticks) * self.dlat + slat
        self.basemap.drawparallels(lat_lines, color="0.5",
                                   linewidth=0.25, dashes=[1, 1, 0.1, 1],
                                   labels=[1, 0, 0, 0], fontsize=12)

    def land(self, color="black"):
        """
        Draw the land mask

        Parameters
        ----------
        color: string, optional