#!/usr/bin/python
# -*- coding: utf-8 -*-
"""Functions to load and plot data from multiple simulation and linearization
files at once

.. _Modelica: http://www.modelica.org/
"""
__author__ = "Kevin Davies"
__email__ = "kdavies4@gmail.com"
__credits__ = ["Kevin Bandy"]
__copyright__ = "Copyright 2012-2013, Georgia Tech Research Corporation"
__license__ = "BSD-compatible (see LICENSE.txt)"


import os
import numpy as np

from modelicares.linres import LinRes
from modelicares.simres import SimRes
from modelicares.base import figure, add_hlines, add_vlines
from control.freqplot import bode, nyquist

from glob import glob
from matplotlib.cbook import iterable
from itertools import cycle


def multiload(locations):
    """Load multiple Modelica_ simulation and/or linearization results.

    **Arguments:**

    - *locations*: Input filename, directory, or list of these

         Wildcards ('*') may be used in the path(s).

    **Returns:**

    1. List of simulations (:class:`simres.SimRes` instances)

    2. List of linearizations (:class:`linres.LinRes` instances)

    Either may be an empty list.

    **Example:**

    .. code-block:: python

       >>> from modelicares import *

       # By file:
       >>> multiload(['examples/ChuaCircuit.mat', 'examples/PID/*/*.mat']) # doctest: +ELLIPSIS
       Valid: SimRes('.../examples/ChuaCircuit.mat')
       Valid: LinRes('.../examples/PID/1/dslin.mat')
       Valid: LinRes('.../examples/PID/2/dslin.mat')
       ([SimRes('.../examples/ChuaCircuit.mat')], [LinRes('.../examples/PID/1/dslin.mat'), LinRes('.../examples/PID/2/dslin.mat')])


       # By directory:
       >>> multiload('examples') # doctest: +ELLIPSIS
       Valid: SimRes('...ChuaCircuit.mat')
       Valid: LinRes('...PID.mat')...
       Valid: SimRes('...ThreeTanks.mat')
       ([SimRes('...ChuaCircuit.mat'), SimRes('...ThreeTanks.mat')], [LinRes('...PID.mat')])
    """

    # Make a list of files.
    fnames = []
    if isinstance(locations, basestring):
        locations = [locations]
    for location in locations:
        if os.path.isdir(location):
            fnames += glob(os.path.join(location, '*.mat'))
        else:
            if '*' in location:
                fnames += glob(location)
            else:
                fnames.append(location)

    # Load the files.
    sims = [] # Simulation results
    lins = [] # Linearization results
    for fname in fnames:
        try:
            sims.append(SimRes(fname))
            print("Valid: " + sims[-1].__repr__())
        except:
            try:
                lins.append(LinRes(fname))
                print("Valid: " + lins[-1].__repr__())
            except:
                print("Could not load simulation or linearization data from "
                      "'%s'." % fname)
    return sims, lins


def multiplot(sims, suffixes='',
              dashes=[(None, None), (3, 3), (1, 1), (3, 2, 1, 2)], **kwargs):
    """Plot data from multiple simulations in 2D Cartesian coordinates.

    This method simply calls :meth:`simres.SimRes.plot` from multiple instances
    of :class:`simres.SimRes`.

    A new figure is created if necessary.

    **Arguments:**

    - *sims*: Simulation result or list of results (instances of
      :class:`simres.SimRes`)

    - *suffixes*: Suffix or list of suffixes for the legends (see
      :meth:`simres.SimRes.plot`)

         If *suffixes* is *None*, then no suffix will be used.  If it is an
         empty string (''), then the base filenames will be used.

    - *\*\*kwargs*: Propagated to :meth:`simres.SimRes.plot` (and thus to
      :meth:`base.plot` and finally :meth:`matplotlib.pyplot.plot`)

         The *dashes* sequence is iterated across all plots.

    **Returns:**

    1. *ax1*: Primary y axes

    2. *ax2*: Secondary y axes

    **Example:**

    .. testsetup::
       >>> from modelicares import closeall
       >>> closeall()

    .. code-block:: python

       >>> from glob import glob
       >>> from modelicares import SimRes, multiplot, saveall

       >>> sims = map(SimRes, glob('examples/ChuaCircuit/*/*.mat'))
       >>> multiplot(sims, title="Chua Circuit", label='examples/ChuaCircuits',
       ...           suffixes=['L.L = %.0f H' % sim.get_IV('L.L')
       ...                     for sim in sims], # Read legend parameters.
       ...           ynames1='L.i', ylabel1="Current") # doctest: +ELLIPSIS
       (<matplotlib.axes._subplots.AxesSubplot object at 0x...>, None)

       >>> saveall()
       Saved examples/ChuaCircuits.pdf
       Saved examples/ChuaCircuits.png

    .. only:: html

       .. image:: ../examples/ChuaCircuits.png
          :scale: 70 %
          :alt: plot of Chua circuit with varying inductance

    .. only:: latex

       .. figure:: ../examples/ChuaCircuits.pdf
          :scale: 70 %

          Plot of Chua circuit with varying inductance
    """
    # Cycle the dashes to make the lines unique.
    if not isinstance(dashes, type(cycle([]))):
        if not iterable(dashes[0]):
            dashes = [dashes]
        dashes = cycle(dashes)
    kwargs['dashes'] = dashes

    # Process the suffixes input.
    if suffixes == '':
        suffixes = [sim.fbase for sim in sims]
    elif suffixes == None:
        suffixes = ['']*len(sims)

    # Generate the plots.
    try:
        ax1, ax2 = sims[0].plot(suffix=suffixes[0], **kwargs)
        kwargs.update({'ax1': ax1, 'ax2': ax2})
        for sim, suffix in zip(sims[1:], suffixes[1:]):
            sim.plot(suffix=suffix, **kwargs)
    except TypeError:
        # sims may be a single simulation result.
        ax1, ax2 = sims.plot(suffix=suffixes, **kwargs)
    return ax1, ax2

def multibode(lins, axes=None, pair=(0, 0), label='bode', title="Bode Plot",
              labels='', colors=['b', 'g', 'r', 'c', 'm', 'y', 'k'],
              styles=[(None,None), (3,3), (1,1), (3,2,1,2)], leg_kwargs={},
              **kwargs):
    r"""Plot multiple linearizations onto a single Bode diagram.

    **Arguments:**

    - *lins*: Linearization result or list of results (instances of
      :class:`linres.LinRes`)

    - *axes*: Tuple (pair) of axes for the magnitude and phase plots

         If *axes* is not provided, then axes will be created in a new figure.

    - *pair*: Tuple of (input index, output index) for the transfer function
      to be chosen from each system (applied to all)

         This is ignored if the system is SISO.

    - *label*: Label for the figure (ignored if axes is provided)

         This will be used as the base filename if the figure is saved.

    - *title*: Title for the figure

    - *labels*: Label or list of labels for the legends

         If *labels* is *None*, then no label will be used.  If it is an
         empty string (''), then the base filenames will be used.

    - *colors*: Color or list of colors that will be used sequentially

         Each may be a character, grayscale, or rgb value.

         .. Seealso:: http://matplotlib.sourceforge.net/api/colors_api.html

    - *styles*: Line/dash style or list of line/dash styles that will be
      used sequentially

         Each style is a string representing a linestyle (e.g., "--") or a
         tuple of on/off lengths representing dashes.  Use "" for no line
         and "-" for a solid line.

         .. Seealso::
            http://matplotlib.sourceforge.net/api/collections_api.html

    - *leg_kwargs*: Dictionary of keyword arguments for
      :meth:`matplotlib.pyplot.legend`

         If *leg_kwargs* is *None*, then no legend will be shown.

    - *\*\*kwargs*: Additional arguments for :meth:`control.freqplot.bode`

    **Returns:**

    1. *axes*: Tuple (pair) of axes for the magnitude and phase plots

    **Example:**

    .. testsetup::
       >>> from modelicares import closeall
       >>> closeall()

    .. code-block:: python

       >>> import os

       >>> from glob import glob
       >>> from modelicares import LinRes, multibode, saveall, read_params
       >>> from numpy import pi, logspace

       >>> lins = map(LinRes, glob('examples/PID/*/*.mat'))
       >>> labels = ["Ti=%g" % read_params('Ti', os.path.join(lin.dir, 'dsin.txt'))
       ...           for lin in lins]
       >>> multibode(lins,
       ...           title="Bode Plot of Modelica.Blocks.Continuous.PID",
       ...           label='examples/PIDs-bode', omega=2*pi*logspace(-2, 3),
       ...           labels=labels, leg_kwargs=dict(loc='lower right')) # doctest: +ELLIPSIS
       (<matplotlib.axes._subplots.AxesSubplot object at 0x...>, <matplotlib.axes._subplots.AxesSubplot object at 0x...>)

       >>> saveall()
       Saved examples/PIDs-bode.pdf
       Saved examples/PIDs-bode.png

    .. only:: html

       .. image:: ../examples/PIDs-bode.png
          :scale: 70 %
          :alt: Bode plot of PID with varying parameters

    .. only:: latex

       .. figure:: ../examples/PIDs-bode.pdf
          :scale: 70 %

          Bode plot of PID with varying parameters
    """
    # Create axes if necessary.
    if not axes:
        fig = figure(label)
        axes = (fig.add_subplot(211), fig.add_subplot(212))

    # Process the lins input.
    if not iterable(lins):
        lins = [lins]

    # Process the labels input.
    if labels == '':
        labels = [lin.fbase for lin in lins]
    elif labels == None:
        labels = ['']*len(lins)

    # Set up the color(s) and line style(s).
    if not iterable(colors):
        # Use the single color for all plots.
        colors = (colors,)
    if not iterable(styles):
        # Use the single line style for all plots.
        styles = [styles]
    elif type(styles[0]) is int:
        # One dashes tuple has been provided; use its value for all plots.
        styles = [styles]
    n_colors = len(colors)
    n_styles = len(styles)

    # Create the plots.
    for i, (lin, label) in enumerate(zip(lins, labels)):
        if lin.sys.inputs > 1 or lin.sys.outputs > 1:
            # Extract the SISO TF. TODO: Is there a better way to do this?
            sys = ss(self.sys.A, self.sys.B[:, pair[0]], self.sys.C[pair[1], :],
                     self.sys.D[pair[1], pair[0]])
        else:
            sys = lin.sys
        bode(sys, Hz=True, label=label,
             color=colors[np.mod(i, n_colors)], axes=axes,
             style=styles[np.mod(i, n_styles)], **kwargs)

    # Decorate and finish.
    axes[0].set_title(title)
    if leg_kwargs is not None:
        loc = leg_kwargs.pop('loc', 'best')
        axes[0].legend(loc=loc, **leg_kwargs)
        axes[1].legend(loc=loc, **leg_kwargs)
    return axes

def multinyquist(lins, ax=None, pair=(0, 0), label='nyquist',
                 title="Nyquist Plot",  xlabel="Real Axis",
                 ylabel="Imaginary Axis", labels='',
                 colors=['b', 'g', 'r', 'c', 'm', 'y', 'k'],
                 leg_kwargs={}, **kwargs):
    """Plot multiple linearizations onto a single Nyquist diagram.

    **Arguments:**

    - *lins*: Linearization result or list of results (instances of
      :class:`linres.LinRes`)

    - *ax*: Axes onto which the Nyquist diagrams should be plotted

         If *ax* is not provided, then axes will be created in a new figure.

    - *pair*: Tuple of (input index, output index) for the transfer function
      to be chosen from each system (applied to all)

         This is ignored if the system is SISO.

    - *label*: Label for the figure (ignored if axes is provided)

         This will be used as the base filename if the figure is saved.

    - *title*: Title for the figure

        - *xlabel*: x-axis label

        - *ylabel*: y-axis label

    - *labels*: Label or list of labels for the legends

         If *labels* is *None*, then no label will be used.  If it is an
         empty string (''), then the base filenames will be used.

    - *colors*: Color or list of colors that will be used sequentially

         Each may be a character, grayscale, or rgb value.

         .. Seealso:: http://matplotlib.sourceforge.net/api/colors_api.html

    - *leg_kwargs*: Dictionary of keyword arguments for
      :meth:`matplotlib.pyplot.legend`

         If *leg_kwargs* is *None*, then no legend will be shown.

    - *\*\*kwargs*: Additional arguments for :meth:`control.freqplot.nyquist`

         If *textFreq* is not specified, then only the frequency points of the
         first system will have text labels.

    **Returns:**

        1. *ax*: Axes of the Nyquist plot

    **Example:**

    .. testsetup::
       >>> from modelicares import closeall
       >>> closeall()

    .. code-block:: python

       >>> import os

       >>> from glob import glob
       >>> from modelicares import LinRes, multinyquist, saveall, read_params
       >>> from numpy import pi, logspace

       >>> lins = map(LinRes, glob('examples/PID/*/*.mat'))
       >>> labels = ["Td=%g" % read_params('Td', os.path.join(lin.dir, 'dsin.txt'))
       ...           for lin in lins]
       >>> multinyquist(lins,
       ...              title="Nyquist Plot of Modelica.Blocks.Continuous.PID",
       ...              label='examples/PIDs-nyquist', textFreq=True,
       ...              omega=2*pi*logspace(-1, 3, 81), labelFreq=20,
       ...              labels=labels) # doctest: +ELLIPSIS
       <matplotlib.axes._subplots.AxesSubplot object at 0x...>

       >>> saveall()
       Saved examples/PIDs-nyquist.pdf
       Saved examples/PIDs-nyquist.png

    .. only:: html

       .. image:: ../examples/PIDs-nyquist.png
          :scale: 70 %
          :alt: Nyquist plot of PID with varying parameters

    .. only:: latex

       .. figure:: ../examples/PIDs-nyquist.pdf
          :scale: 70 %

          Nyquist plot of PID with varying parameters
    """
    # Create axes if necessary.
    if not ax:
        fig = figure(label)
        ax = fig.add_subplot(111, aspect='equal')

    # Process the lins input.
    if not iterable(lins):
        lins = [lins]


    # Process the labels input.
    if labels == '':
        labels = [lin.fbase for lin in lins]
    elif labels == None:
        labels = ['']*len(lins)

    # Set up the color(s).
    if not iterable(colors):
        # Use the single color for all plots.
        colors = (colors,)
    n_colors = len(colors)

    # Create the plots.
    textFreq = kwargs.pop('textFreq', None)
    for i, (lin, label) in enumerate(zip(lins, labels)):
        if lin.sys.inputs > 1 or lin.sys.outputs > 1:
            # Extract the SISO TF. TODO: Is there a better way to do this?
            sys = ss(self.sys.A, self.sys.B[:, pair[0]], self.sys.C[pair[1], :],
                     self.sys.D[pair[1], pair[0]])
        else:
            sys = lin.sys
        nyquist(sys, mark=False, label=label, ax=ax,
                textFreq=i==0 if textFreq is None else textFreq,
                color=colors[np.mod(i, n_colors)], **kwargs)

    # Decorate and finish.
    add_hlines(ax, color='k', linestyle='--', linewidth=0.5)
    add_vlines(ax, color='k', linestyle='--', linewidth=0.5)
    ax.set_title(title)
    if xlabel: # Without this check, xlabel=None will give a label of "None".
        ax.set_xlabel(xlabel)
    if ylabel: # Same purpose
        ax.set_ylabel(ylabel)
    if leg_kwargs is not None:
        loc = leg_kwargs.pop('loc', 'best')
        ax.legend(loc=loc, **leg_kwargs)
    return ax

if __name__ == '__main__':
    """Test the contents of this file."""
    import doctest
    doctest.testmod()
    exit()
