# phaseplot.py - generate 2D phase portraits
#
# Author: Richard M. Murray
# Date: 24 July 2011, converted from MATLAB version (2002); based on
# a version by Kristi Morgansen
#
# Copyright (c) 2011 by California Institute of Technology
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#   1. Redistributions of source code must retain the above copyright
#      notice, this list of conditions and the following disclaimer.
#
#   2. Redistributions in binary form must reproduce the above copyright 
#      notice, this list of conditions and the following disclaimer in the 
#      documentation and/or other materials provided with the distribution.
#
#   3. The name of the author may not be used to endorse or promote products 
#      derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT,
# INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
# HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
# STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING
# IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
# POSSIBILITY OF SUCH DAMAGE.

import numpy as np
import matplotlib.pyplot as mpl
from matplotlib.mlab import frange, find
from exception import ControlNotImplemented
from scipy.integrate import odeint

def phase_plot(odefun, X=None, Y=None, scale=1, X0=None, T=None,
              lingrid=None, lintime=None, logtime=None, timepts=None,
              parms=(), verbose=True):
    """
    Phase plot for 2D dynamical systems

    Produces a vector field or stream line plot for a planar system.

    Call signatures:
      phase_plot(func, X, Y, ...) - display vector field on meshgrid
      phase_plot(func, X, Y, scale, ...) - scale arrows
      phase_plot(func. X0=(...), T=Tmax, ...) - display stream lines
      phase_plot(func, X, Y, X0=[...], T=Tmax, ...) - plot both
      phase_plot(func, X0=[...], T=Tmax, lingrid=N, ...) - plot both
      phase_plot(func, X0=[...], lintime=N, ...) - stream lines with arrows

    Parameters
    ----------
    func : callable(x, t, ...)
        Computes the time derivative of y (compatible with odeint).
        The function should be the same for as used for
        scipy.integrate.  Namely, it should be a function of the form
        dxdt = F(x, t) that accepts a state x of dimension 2 and
        returns a derivative dx/dt of dimension 2.

    X, Y: ndarray, optional
        Two 1-D arrays representing x and y coordinates of a grid.
        These arguments are passed to meshgrid and generate the lists
        of points at which the vector field is plotted.  If absent (or
        None), the vector field is not plotted.

    scale: float, optional
        Scale size of arrows; default = 1

    X0: ndarray of initial conditions, optional
        List of initial conditions from which streamlines are plotted.
        Each initial condition should be a pair of numbers.

    T: array-like or number, optional
        Length of time to run simulations that generate streamlines.
        If a single number, the same simulation time is used for all
        initial conditions.  Otherwise, should be a list of length
        len(X0) that gives the simulation time for each initial
        condition.  Default value = 50.

    lingrid = N or (N, M): integer or 2-tuple of integers, optional
        If X0 is given and X, Y are missing, a grid of arrows is
        produced using the limits of the initial conditions, with N
        grid points in each dimension or N grid points in x and M grid
        points in y.

    lintime = N: integer, optional
        Draw N arrows using equally space time points

    logtime = (N, lambda): (integer, float), optional
        Draw N arrows using exponential time constant lambda

    timepts = [t1, t2, ...]: array-like, optional
        Draw arrows at the given list times

    parms: tuple, optional
        List of parameters to pass to vector field: func(x, t, *parms)

    See also
    --------
    box_grid(X, Y): construct box-shaped grid of initial conditions

    Examples
    --------
    """

    #
    # Figure out ranges for phase plot (argument processing)
    #
    #! TODO: need to add error checking to arguments
    autoFlag = False; logtimeFlag = False; timeptsFlag = False; Narrows = 0;
    if (lingrid != None):
        autoFlag = True;
        Narrows = lingrid;
        if (verbose):
            print 'Using auto arrows\n';

    elif (logtime != None):
        logtimeFlag = True;
        Narrows = logtime[0];
        timefactor = logtime[1];
        if (verbose):
            print 'Using logtime arrows\n';

    elif (timepts != None):
        timeptsFlag = True;
        Narrows = len(timepts);

    else:
        # Figure out the set of points for the quiver plot
        #! TODO: Add sanity checks
        (x1, x2) = np.meshgrid(
            frange(X[0], X[1], float(X[1]-X[0])/X[2]),
            frange(Y[0], Y[1], float(Y[1]-Y[0])/Y[2]));

    if ((not autoFlag) and (not logtimeFlag) and (not timeptsFlag)):
        # Now calculate the vector field at those points
        (nr,nc) = x1.shape;
        dx = np.empty((nr, nc, 2))
        for i in range(nr):
            for j in range(nc):
                dx[i, j, :] = np.squeeze(odefun((x1[i,j], x2[i,j]), 0, *parms))

        # Plot the quiver plot
        #! TODO: figure out arguments to make arrows show up correctly
        if (scale == None):
            mpl.quiver(x1, x2, dx[:,:,1], dx[:,:,2], angles='xy')
        elif (scale != 0):
            #! TODO: optimize parameters for arrows
            #! TODO: figure out arguments to make arrows show up correctly
            xy = mpl.quiver(x1, x2, dx[:,:,0]*np.abs(scale), 
                            dx[:,:,1]*np.abs(scale), angles='xy')
            # set(xy, 'LineWidth', PP_arrow_linewidth, 'Color', 'b');

        #! TODO: Tweak the shape of the plot 
        # a=gca; set(a,'DataAspectRatio',[1,1,1]);
        # set(a,'XLim',X(1:2)); set(a,'YLim',Y(1:2));
        mpl.xlabel('x1'); mpl.ylabel('x2');

    # See if we should also generate the streamlines
    if (X0 == None or len(X0) == 0):
        return

    # Convert initial conditions to a numpy array
    X0 = np.array(X0);
    (nr, nc) = np.shape(X0);

    # Generate some empty matrices to keep arrow information
    x1 = np.empty((nr, Narrows)); x2 = np.empty((nr, Narrows));
    dx = np.empty((nr, Narrows, 2))
  
    # See if we were passed a simulation time
    if (T == None):
        T = 50

    # Parse the time we were passed
    TSPAN = T;
    if (isinstance(T, (int, float))):
        TSPAN = np.linspace(0, T, 100);

    # Figure out the limits for the plot
    if (scale == None):
        # Assume that the current axis are set as we want them
        alim = mpl.axis();
        xmin = alim[0]; xmax = alim[1]; 
        ymin = alim[2]; ymax = alim[3];
    else:
        # Use the maximum extent of all trajectories
        xmin = np.min(X0[:,0]); xmax = np.max(X0[:,0]);
        ymin = np.min(X0[:,1]); ymax = np.max(X0[:,1]);

    # Generate the streamlines for each initial condition
    for i in range(nr):
        state = odeint(odefun, X0[i], TSPAN, args=parms);
        time = TSPAN
        mpl.hold(True);
        mpl.plot(state[:,0], state[:,1])
        #! TODO: add back in colors for stream lines
        # PP_stream_color(np.mod(i-1, len(PP_stream_color))+1));
        # set(h[i], 'LineWidth', PP_stream_linewidth);

        # Plot arrows if quiver parameters were 'auto'
        if (autoFlag or logtimeFlag or timeptsFlag):
            # Compute the locations of the arrows
            #! TODO: check this logic to make sure it works in python
            for j in range(Narrows):
      
                # Figure out starting index; headless arrows start at 0
                k = -1 if scale == None else 0;
      
                # Figure out what time index to use for the next point
                if (autoFlag):
                    # Use a linear scaling based on ODE time vector
                    tind = np.floor((len(time)/Narrows) * (j-k)) + k;
                elif (logtimeFlag):
                    # Use an exponential time vector
                    # MATLAB: tind = find(time < (j-k) / lambda, 1, 'last');
                    tarr = find(time < (j-k) / timefactor);
                    tind = tarr[-1] if len(tarr) else 0;
                elif (timeptsFlag):
                    # Use specified time points
                    # MATLAB: tind = find(time < Y[j], 1, 'last');
                    tarr = find(time < timepts[j]);
                    tind = tarr[-1] if len(tarr) else 0;

                # For tailless arrows, skip the first point
                if (tind == 0 and scale == None):
                    continue;
      
                # Figure out the arrow at this point on the curve
                x1[i,j] = state[tind, 0];
                x2[i,j] = state[tind, 1];

                # Skip arrows outside of initial condition box
                if (scale != None or 
                     (x1[i,j] <= xmax and x1[i,j] >= xmin and
                      x2[i,j] <= ymax and x2[i,j] >= ymin)):
                    v = odefun((x1[i,j], x2[i,j]), 0, *parms)
                    dx[i, j, 0] = v[0]; dx[i, j, 1] = v[1];
                else:
                    dx[i, j, 0] = 0; dx[i, j, 1] = 0;

    # Set the plot shape before plotting arrows to avoid warping
    # a=gca;  
    # if (scale != None):
    #     set(a,'DataAspectRatio', [1,1,1]);
    # if (xmin != xmax and ymin != ymax):
    #     mpl.axis([xmin, xmax, ymin, ymax]);
    # set(a, 'Box', 'on');

    # Plot arrows on the streamlines
    if (scale == None and Narrows > 0):
        # Use a tailless arrow
        #! TODO: figure out arguments to make arrows show up correctly
        mpl.quiver(x1, x2, dx[:,:,0], dx[:,:,1], angles='xy')
    elif (scale != 0 and Narrows > 0):
        #! TODO: figure out arguments to make arrows show up correctly
        xy = mpl.quiver(x1, x2, dx[:,:,0]*abs(scale), dx[:,:,1]*abs(scale),
                        angles='xy')
        # set(xy, 'LineWidth', PP_arrow_linewidth);
        # set(xy, 'AutoScale', 'off');
        # set(xy, 'AutoScaleFactor', 0);

    if (scale < 0):
        bp = mpl.plot(x1, x2, 'b.');		# add dots at base
        # set(bp, 'MarkerSize', PP_arrow_markersize);

    return;

# Utility function for generating initial conditions around a box
def box_grid(xlimp, ylimp):
    """box_grid   generate list of points on edge of box
    
    list = box_grid([xmin xmax xnum], [ymin ymax ynum]) generates a
    list of points that correspond to a uniform grid at the end of the
    box defined by the corners [xmin ymin] and [xmax ymax].
    """

    sx10 = frange(xlimp[0], xlimp[1], float(xlimp[1]-xlimp[0])/xlimp[2])
    sy10 = frange(ylimp[0], ylimp[1], float(ylimp[1]-ylimp[0])/ylimp[2])

    sx1 = np.hstack((0, sx10, 0*sy10+sx10[0], sx10, 0*sy10+sx10[-1]))
    sx2 = np.hstack((0, 0*sx10+sy10[0], sy10, 0*sx10+sy10[-1], sy10))

    return np.transpose( np.vstack((sx1, sx2)) )
