# -*- coding: utf-8 -*-
"""Tools for visualizing ADCP data that is read and processed by the adcpy module
This module is imported under the main adcpy, and should be available as 
adcpy.plot. Some methods can be used to visualize flat arrays, independent of
adcpy, and the plots may be created quickly using the IPanel and QPanel 
classes.

This code is open source, and defined by the included MIT Copyright License 

Designed for Python 2.7; NumPy 1.7; SciPy 0.11.0; Matplotlib 1.2.0
2014-09 - First Release; blsaenz, esatel
"""
import numpy as np
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
import scipy.stats.stats as sp
from matplotlib.dates import num2date#,date2num,

import adcpy
from adcpy_recipes import calc_transect_flows_from_uniform_velocity_grid

U_str = 'u'
V_str = 'v'
W_str = 'w'
vel_strs = (U_str,V_str,W_str)


# Common formatting for datenums:
def fmt_dnum(dn):
    return num2date(dn).strftime('%c')


class IPanel(object):
    """
    This object stores and plots a 2D velocity map as an image.  Any of the data
    fields  (kwarg_options) may be specificed as kwargs during initialization. 
    At minimum IPanel requires 'velocity' to be set.
    """
    kwarg_options = ['use_pcolormesh',
                     'minv',      
                     'maxv',       
                     'velocity',          
                     'title',
                     'units',
                     'xlabel',
                     'ylabel',
                     'x',
                     'y',
                     'chop_off_nans',
                     'x_is_mtime',
                     'arrow_color',
                     'xy_is_lonlat',
                     'interpolation',
                     'shading',
                     'my_axes']

    def __init__(self,**kwargs):
        
        # init everything to None
        for kwarg in self.kwarg_options:
                exec("self.%s = None"%kwarg)        
        # set defaults
        self.minv = -0.25
        self.maxv = 0.25
        self.x_is_mtime = False
        self.interpolation = 'nearest'
        self.use_pcolormesh = False
        self.shading = 'flat'
        self.xy_is_lonlat = False
        self.chop_off_nans = False
        
        # read/save arguments
        for kwarg in self.kwarg
_options:
            if kwarg in kwargs:
                exec("self
.%s = kwargs[kwarg]"%kwarg)        
        
    def plot(self,ax=None):
        """
        Plots the data in IPanel onto the axis ax, or if ax is None,
        onto self.my_axes.
        Inputs:
            ax = matplotlib axes object, or None
        Returns:
            Nothing
        """      
        # set desired axes
        if ax is not None:
            plt.sca(ax)
        elif self.my_axes is not None:
            ax = plt.sca(self.my_axes)
        else:
            ax = plt.gca()                
        if self.minv is not None:
            mnv = ",vmin=self.minv"
        else:
            mnv = ""
        if self.minv is not None:
            mxv = ",vmax=self.maxv"
        else:
            mxv = ""
        if self.use_pcolormesh:
            vel_masked = np.ma.array(self.velocity,mask=np.isnan(self.velocity))
            if self.x is not None and self.y is not None:
                xy = "self.x,self.y,"
            else:            
                xy = ""
            plot_cmd = "pc=plt.pcolormesh(%svel_masked.T,shading=self.shading%s%s)"%(xy,mnv,mxv)
            exec(plot_cmd)
        else:
            if self.x is not None and self.y is not None:
                xy = ",extent=[self.x[0],self.x[-1],self.y[-1],self.y[0]]"
            else:            
                xy = ""
            plot_cmd = "pc=plt.imshow(self.velocity.T%s,interpolation=self.interpolation%s%s)"%(xy,mnv,mxv)

        exec(plot_cmd)
        if self.title is not None:
            plt.title(self.title)
        plt.axis('tight')
        if self.chop_off_nans:
            x_test = np.nansum(self.velocity,axis=1)
            x_test = ~np.isnan(x_test)*np.arange(np.size(x_test))
            if self.x is None:
                plt.xlim([np.nanmin(x_test),np.nanmax(x_test)])
            else:
                plt.xlim([self.x[np.nanmin(x_test)],self.x[np.nanmax(x_test)]])
                if self.x[-1] < self.x[0]:
                    plt.xlim(plt.xlim()[::-1])
            y_test = np.nansum(self.velocity,axis=0)
            y_test = ~np.isnan(y_test)*np.arange(np.size(y_test))
            plt.ylim([np.nanmin(y_test),np.nanmax(y_test)])
            if self.y is None:
                plt.ylim([np.nanmin(y_test),np.nanmax(y_test)])
            else:
                plt.ylim([self.y[np.nanmin(y_test)],self.y[np.nanmax(y_test)]])
                if self.y[-1] < self.y[0]:
                    plt.ylim(plt.ylim()[::-1])
        if self.x_is_mtime:
            ax.xaxis_date()
            plt.gcf().autofmt_xdate()
        elif self.xy_is_lonlat:
            ax.yaxis.set_major_formatter(plt.FormatStrFormatter('%7.4f'))       
            ax.xaxis.set_major_formatter(plt.FormatStrFormatter('%7.4f'))
            plt.ylabel('Latitude [degrees N]')
            plt.xlabel('Longitude [degrees E]')            
        if self.xlabel is not None:
            plt.xlabel(self.xlabel)
        if self.ylabel is not None:
            plt.ylabel(self.ylabel)
        plt.colorbar(pc, use_gridspec=True)


class QPanel(object):
    """
    This object stores and plots a 1D or 2D velocity map as a quiver plot.  Any 
    of the data fields (kwarg_options) may be specificed as kwargs during 
    initialization.  At minimum QPanel requires 'velocity' to be set.
    """
    kwarg_options = ['u_vecs',
                     'v_vecs',
                     'velocity',          
                     'title',
                     'units',
                     'xlabel',
                     'ylabel',
                     'x',
                     'y',
                     'v_scale',  # make arrow bigger or smaller, relatively speaking, daults to 1
                     'xpand',    # fractional buffer around xy extent, to capture arrow ends
                     'x_is_mtime',
                     'arrow_color',
                     'xy_is_lonlat',
                     'equal_axes',
                     'my_axes']

    def __init__(self,**kwargs):
        
        # init everything to None
        for kwarg in self.kwarg_options:
                exec("self.%s = None"%kwarg)        
        # set defaults
        self.u_vecs = 50
        self.v_vecs = 50
        self.x_is_mtime = False
        self.xy_is_lonlat = False
        self.arrow_color = 'k'
        self.v_scale = 1.0
        self.xpand = 0.33
        self.equal_axes = False
         
        # read/save arguments
        for kwarg in self.kwarg_options:
            if kwarg in kwargs:
                exec("self.%s = kwargs[kwarg]"%kwarg)

    def plot(self,ax=None):
        """
        Plots the data in QPanel onto the axis ax, or if ax is None,
        onto self.my_axes.
        Inputs:
            ax = matplotlib axes object, or None
        Returns:
            Nothing
        """       
        # set desired axes
        if ax is not None:
            plt.sca(ax)
        elif self.my_axes is not None:
            ax = plt.sca(self.my_axes)
        else:
            ax = plt.gca()
        dims = np.shape(self.velocity)
        u_reduction = max(1,int(dims[0]/self.u_vecs))
        u_indices = np.arange(0,dims[0],u_reduction)
        v_mag = np.sqrt(self.velocity[...,0]**2 + self.velocity[...,1]**2)
        if len(dims) == 2:
            vScale = np.nanmax(v_mag[u_indices])
            local_vel = self.velocity[u_indices,...]
            local_u = local_vel[:,0]
            local_v = local_vel[:,1]
            local_x = self.x[u_indices]
            local_y = self.y[u_indices]
        elif len(dims) == 3:
            v_reduction = max(1,int(dims[1]/self.v_vecs))
            v_indices = np.arange(0,dims[1],v_reduction)
            v_mag = v_mag[u_indices,:]
            v_mag = v_mag[:,v_indices]
            vScale = np.nanmax(np.nanmax(v_mag))
            local_vel = self.velocity[u_indices,:,:]
            l