"""
Classes to plot TimeSeries w/ matplotlib.

:author: Pierre GF Gerard-Marchant & Matt Knox
:contact: pierregm_at_uga_dot_edu - mattknow_ca_at_hotmail_dot_com
:date: $Date: 2007-07-18 20:52:17 -0700 (Wed, 18 Jul 2007) $
:version: $Id: mpl_timeseries.py 3175 2007-07-19 03:52:17Z pierregm $
"""
__author__ = "Pierre GF Gerard-Marchant & Matt Knox ($Author: pierregm $)"
__version__ = '1.0'
__revision__ = "$Revision: 3175 $"
__date__     = '$Date: 2007-07-18 20:52:17 -0700 (Wed, 18 Jul 2007) $'


import matplotlib
from matplotlib import pylab, rcParams
from matplotlib import _pylab_helpers
from matplotlib.artist import setp
from matplotlib.axes import Subplot, PolarSubplot
from matplotlib.cbook import flatten
from matplotlib.collections import LineCollection
from matplotlib.contour import ContourSet
from matplotlib.dates import DayLocator, MonthLocator, YearLocator, \
                             DateFormatter
from matplotlib.figure import Figure
from matplotlib.legend import Legend
from matplotlib.mlab import meshgrid
from matplotlib.ticker import Formatter, ScalarFormatter, FuncFormatter, \
                              Locator, FixedLocator
#from matplotlib.transforms import nonsingular

import numpy
import maskedarray as MA

import timeseries
from timeseries import date_array, Date, DateArray, TimeSeries, get_freq_group
from timeseries import const as _c

import warnings

#####---------------------------------------------------------------------------
#---- --- Matplotlib extensions ---
#####---------------------------------------------------------------------------

def add_generic_subplot(figure_instance, *args, **kwargs):
    """Generalizes the `add_subplot` figure method to generic subplots.
The specific Subplot object class to add is given through the keywords
`SubplotClass` or `class`.

:Parameters:
    `figure_instance` : Figure object
        Figure to which the generic subplot should be attached.
    `args` : Misc
        Miscellaneous arguments to the subplot.
    `kwargs` : Dictionary
        Keywords. Same keywords as `Subplot`, with the addition of
        - `SubplotClass` : Type of subplot
        - `subclass` : Shortcut to `SubplotClass`.
        - any keyword required by the `SubplotClass` subclass.
    """

    key = figure_instance._make_key(*args, **kwargs)
    #TODO: Find why, sometimes, key is not hashable (even if tuple)
    # else, there's a fix below
    try:
        key.__hash__()
    except TypeError:
        key = str(key)
    #
    if figure_instance._seen.has_key(key):
        ax = figure_instance._seen[key]
        figure_instance.sca(ax)
        return ax
    #
    if not len(args):
        return
#    if hasattr(args[0], '__array__'):
#        fixedargs = args[1:]
#    else:
#        fixedargs = args
    #
    SubplotClass = kwargs.pop("SubplotClass", Subplot)
    SubplotClass = kwargs.pop("subclass",SubplotClass)
    if isinstance(args[0], Subplot) or isinstance(args[0], PolarSubplot):
        a = args[0]
        assert(a.get_figure() is figure_instance)
#        a.set_figure(figure_instance)
    else:
        ispolar = kwargs.pop('polar', False)
        if ispolar:
            a = PolarSubplot(figure_instance, *args, **kwargs)
        else:
            a = SubplotClass(figure_instance, *args, **kwargs)

    figure_instance.axes.append(a)
    figure_instance._axstack.push(a)
    figure_instance.sca(a)
    figure_instance._seen[key] = a
    return a


def nonsingular(vmin, vmax, expander=0.001, tiny=1e-15, increasing=True):
    '''
    Ensure the endpoints of a range are not too close together.

    "too close" means the interval is smaller than 'tiny' times
            the maximum absolute value.

    If they are too close, each will be moved by the 'expander'.
    If 'increasing' is True and vmin > vmax, they will be swapped.
    '''
    #TODO: Remove that when matplotlib incorporate it by default
    swapped = False
    if vmax < vmin:
        vmin, vmax = vmax, vmin
        swapped = True
    if vmax - vmin <= max(abs(vmin), abs(vmax)) * tiny:
        if vmin == 0.0:
            vmin = -expander
            vmax = expander
        else:
            vmin -= expander*abs(vmin)
            vmax += expander*abs(vmax)
    if swapped and not increasing:
        vmin, vmax = vmax, vmin
    return vmin, vmax

##### -------------------------------------------------------------------------
#---- --- Locators ---
##### -------------------------------------------------------------------------

def _get_default_annual_spacing(nyears):
    """Returns a default spacing between consecutive ticks for annual data."""
    if nyears < 11:
        (min_spacing, maj_spacing) = (1, 1)
    elif nyears < 20:
        (min_spacing, maj_spacing) = (1, 2)
    elif nyears < 50:
        (min_spacing, maj_spacing) = (1, 5)
    elif nyears < 100:
        (min_spacing, maj_spacing) = (5, 10)
    elif nyears < 200:
        (min_spacing, maj_spacing) = (5, 25)
    elif nyears < 600:
        (min_spacing, maj_spacing) = (10, 50)
    else:
        factor = nyears // 1000 + 1
        (min_spacing, maj_spacing) = (factor*20, factor*100)
    return (min_spacing, maj_spacing)


def period_break(dates, period):
    """Returns the indices where the given period changes.

:Parameters:
    dates : DateArray
        Array of dates to monitor.
    period : string
        Name of the period to monitor.
    """
    current = getattr(dates, period)
    previous = getattr(dates-1, period)
    return (current - previous).nonzero()[0]

def has_level_label(label_flags):
    """returns true if the label_flags indicate there is at least
one label for this level"""
    if label_flags.size == 0 or \
       (label_flags.size == 1 and label_flags[0] == 0):
        return False
    else:
        return True

def _daily_finder(vmin, vmax, freq, aslocator):

    if freq == _c.FR_BUS:
        periodsperyear = 261
        periodspermonth = 19
    elif freq == _c.FR_DAY:
        periodsperyear = 365
        periodspermonth = 28
    elif get_freq_group(freq) == _c.FR_WK:
        periodsperyear = 52
        periodspermonth = 3
    else:
        raise ValueError("unexpected frequency")

    (vmin, vmax) = (int(vmin), int(vmax))
    span = vmax - vmin + 1
    dates = date_array(start_date=Date(freq,vmin),
                       end_date=Date(freq, vmax))
    default = numpy.arange(vmin, vmax+1)
    # Initialize the output
    if not aslocator:
        format = numpy.empty(default.shape, dtype="|S10")
        format.flat = ''

    def first_label(label_flags):
        if label_flags[0] == 0: return label_flags[1]
        else: return label_flags[0]

    # Case 1. Less than a month
    if span <= periodspermonth:
        month_start = period_break(dates,'month')
        if aslocator:
            major = default[month_start]
            minor = default
        else:
            year_start = period_break(dates,'year')
            format[:] = '%d'
            format[month_start] = '%d\n%b'
            format[year_start] = '%d\n%b\n%Y'
            if not has_level_label(year_start):
                if not has_level_label(month_start):
                    if dates.size > 1:
                        idx = 1
                    else:
                        idx = 0
                    format[idx] = '%d\n%b\n%Y'
                else:
                    format[first_label(month_start)] = '%d\n%b\n%Y'
    # Case 2. Less than three months
    elif span <= periodsperyear//4:
        month_start = period_break(dates,'month')
        if aslocator:
            major = default[month_start]
            minor = default
        else:
            week_start = period_break(dates,'week')
            year_start = period_break(dates,'year')

            format[week_start] = '%d'
            format[month_start] = '\n\n%b'
            format[year_start] = '\n\n%b\n%Y'
            if not has_level_label(year_start):
                if not has_level_label(month_start):
                    format[first_label(week_start)] = '\n\n%b\n%Y'
                else:
                    format[first_label(month_start)] = '\n\n%b\n%Y'
    # Case 3. Less than 14 months ...............
    elif span <= 1.15 * periodsperyear:

        if aslocator:
            d_minus_1 = dates-1

            month_diff = numpy.abs(dates.month - d_minus_1.month)
            week_diff = numpy.abs(dates.week - d_minus_1.week)
            minor_idx = (month_diff + week_diff).nonzero()[0]

            major = default[month_diff != 0]
            minor = default[minor_idx]
        else:
            year_start = period_break(dates,'year')
            month_start = period_break(dates,'month')

            format[month_start] = '%b'
            format[year_start] = '%b\n%Y'
            if not has_level_label(year_start):
                format[first_label(month_start)] = '%b\n%Y'
    # Case 4. Less than 2.5 years ...............
    elif span <= 2.5 * periodsperyear:
        year_start = period_break(dates,'year')
        if aslocator:
            month_start = period_break(dates, 'quarter')
            major = default[year_start]
            minor = default[month_start]
        else:
            quarter_start = period_break(dates, 'quarter')
            format[quarter_start] = '%b'
            format[year_start] = '%b\n%Y'
    # Case 4. Less than 4 years .................
    elif span <= 4 * periodsperyear:
        year_start = period_break(dates,'year')
        month_start = period_break(dates, 'month')
        if aslocator:
            major = default[year_start]
            minor = default[month_start]
        else:
            month_break = dates[month_start].month
            jan_or_jul = month_start[(month_break == 1) | (month_break == 7)]
            format[jan_or_jul] = '%b'
            format[year_start] = '%b\n%Y'
    # Case 5. Less than 11 years ................
    elif span <= 11 * periodsperyear:
        year_start = period_break(dates,'year')
        if aslocator:
            quarter_start = period_break(dates, 'quarter')
            major = default[year_start]
            minor = default[quarter_start]
        else:
            format[year_start] = '%Y'
    # Case 6. More than 12 years ................
    else:
        year_start = period_break(dates,'year')
        year_break = dates[year_start].years
        nyears = span/periodsperyear
        (min_anndef, maj_anndef) = _get_default_annual_spacing(nyears)
        major_idx = year_start[(year_break % maj_anndef == 0)]
        if aslocator:
            major = default[major_idx]
            minor_idx = year_start[(year_break % min_anndef == 0)]
            minor = default[minor_idx]
        else:
            format[major_idx] = '%Y'
    #............................................
    if aslocator:
        return minor, major
    else:
        formatted = (format != '')
        return dict([(d,f) for (d,f) in zip(default[formatted],format[formatted])])
#...............................................................................
def _monthly_finder(vmin, vmax, freq, aslocator):
    if freq != _c.FR_MTH:
        raise ValueError("Unexpected frequency")
    periodsperyear = 12

    (vmin, vmax) = (int(vmin), int(vmax))
    span = vmax - vmin + 1
    #............................................
    dates = numpy.arange(vmin, vmax+1)
    format = numpy.empty(span, dtype="|S8")
    format.flat = ''
    year_start = (dates % 12 == 1).nonzero()[0]
    #............................................
    if span <= 1.15 * periodsperyear:
        if aslocator:
            major = dates[year_start]
            minor = dates
        else:

            format[:] = '%b'
            format[year_start] = '%b\n%Y'

            if not has_level_label(year_start):
                if dates.size > 1:
                    idx = 1
                else:
                    idx = 0
                format[idx] = '%b\n%Y'
    #........................
    elif span <= 2.5 * periodsperyear:
        if aslocator:
            major = dates[year_start]
            minor = dates
        else:
            quarter_start = (dates % 3 == 1).nonzero()
            format[quarter_start] = '%b'
            format[year_start] = '%b\n%Y'
    #.......................
    elif span <= 4 * periodsperyear:
        if aslocator:
            major = dates[year_start]
            minor = dates
        else:
            jan_or_jul = (dates % 12 == 1) | (dates % 12 == 7)
            format[jan_or_jul] = '%b'
            format[year_start] = '%b\n%Y'
    #........................
    elif span <= 11 * periodsperyear:
        if aslocator:
            quarter_start = (dates % 3 == 1).nonzero()
            major = dates[year_start]
            minor = dates[quarter_start]
        else:
            format[year_start] = '%Y'
   #.........................
    else:
        nyears = span/periodsperyear
        (min_anndef, maj_anndef) = _get_default_annual_spacing(nyears)
        years = dates[year_start]//12 + 1
        major_idx = year_start[(years % maj_anndef == 0)]
        if aslocator:
            major = dates[major_idx]
            minor = dates[year_start[(years % min_anndef == 0)]]
        else:
            format[major_idx] = '%Y'
    #........................
    if aslocator:
        return minor, major
    else:
        formatted = (format != '')
        return dict([(d,f) for (d,f) in zip(dates[formatted],format[formatted])])
#...............................................................................
def _quarterly_finder(vmin, vmax, freq, aslocator):
    if get_freq_group(freq) != _c.FR_QTR:
        raise ValueError("Unexpected frequency")
    periodsperyear = 4
    (vmin, vmax) = (int(vmin), int(vmax))
    span = vmax - vmin + 1
    #............................................
    dates = numpy.arange(vmin, vmax+1)
    format = numpy.empty(span, dtype="|S8")
    format.flat = ''
    year_start = (dates % 4 == 1).nonzero()[0]
    #............................................
    if span <= 3.5 * periodsperyear:
        if aslocator:
            major = dates[year_start]
            minor = dates
        else:
            format[:] = 'Q%q'
            format[year_start] = 'Q%q\n%F'
            if not has_level_label(year_start):
                if dates.size > 1:
                    idx = 1
                else:
                    idx = 0
                format[idx] = 'Q%q\n%F'
    #............................................
    elif span <= 11 * periodsperyear:
        if aslocator:
            major = dates[year_start]
            minor = dates
        else:
            format[year_start] = '%F'
    #............................................
    else:
        years = dates[year_start]//4 + 1
        nyears = span/periodsperyear
        (min_anndef, maj_anndef) = _get_default_annual_spacing(nyears)
        major_idx = year_start[(years % maj_anndef == 0)]
        if aslocator:
            major = dates[major_idx]
            minor = dates[year_start[(years % min_anndef == 0)]]
        else:
            format[major_idx] = '%F'
    #............................................
    if aslocator:
        return minor, major
    else:
        formatted = (format != '')
        return dict([(d,f) for (d,f) in zip(dates[formatted],format[formatted])])
#...............................................................................
def _annual_finder(vmin, vmax, freq, aslocator):
    if get_freq_group(freq) != _c.FR_ANN:
        raise ValueError("Unexpected frequency")
    (vmin, vmax) = (int(vmin), int(vmax+1))
    span = vmax - vmin + 1
    #............................................
    dates = numpy.arange(vmin, vmax+1)
    format = numpy.empty(span, dtype="|S8")
    format.flat = ''
    #............................................
    (min_anndef, maj_anndef) = _get_default_annual_spacing(span)
    major_idx = dates % maj_anndef == 0
    if aslocator:
        major = dates[major_idx]
        minor = dates[(dates % min_anndef == 0)]
    else:
        format[major_idx] = '%Y'
    #............................................
    if aslocator:
        return minor, major
    else:
        formatted = (format != '')
        return dict([(d,f) for (d,f) in zip(dates[formatted],format[formatted])])

#...............................................................................
class TimeSeries_DateLocator(Locator):
    "Locates the ticks along an axis controlled by a DateArray."

    def __init__(self, freq, minor_locator=False, dynamic_mode=True,
                 base=1, quarter=1, month=1, day=1):
        self.freq = freq
        self.base = base
        fgroup = get_freq_group(freq)
        (self.quarter, self.month, self.day) = (quarter, month, day)
        self.isminor = minor_locator
        self.isdynamic = dynamic_mode
        self.offset = 0
        #.....
        if fgroup == _c.FR_ANN:
            self.finder = _annual_finder
        elif fgroup == _c.FR_QTR:
            self.finder = _quarterly_finder
        elif freq == _c.FR_MTH:
            self.finder = _monthly_finder
        elif freq in (_c.FR_BUS, _c.FR_DAY) or fgroup == _c.FR_WK:
            self.finder = _daily_finder

    def asminor(self):
        "Returns the locator set to minor mode."
        self.isminor = True
        return self

    def asmajor(self):
        "Returns the locator set to major mode."
        self.isminor = False
        return self

    def _get_default_locs(self, vmin, vmax):
        "Returns the default locations of ticks."
        (minor, major) = self.finder(vmin, vmax, self.freq, True)
        if self.isminor:
            return minor
        return major

    def __call__(self):
        'Return the locations of the ticks.'
        self.verify_intervals()
        vmin, vmax = self.viewInterval.get_bounds()
        if vmax < vmin:
            vmin, vmax = vmax, vmin
        if self.isdynamic:
            locs = self._get_default_locs(vmin, vmax)
        else:
            base = self.base
            (d, m) = divmod(vmin, base)
            vmin = (d+1) * base
            locs = range(vmin, vmax+1, base)
        return locs

    def autoscale(self):
        """Sets the view limits to the nearest multiples of base that contain
    the data.
        """
        self.verify_intervals()
        dmin, dmax = self.dataInterval.get_bounds()
        locs = self._get_default_locs(dmin, dmax)
        (vmin, vmax) = locs[[0, -1]]
        if vmin == vmax:
            vmin -= 1
            vmax += 1
        return nonsingular(vmin, vmax)

#####---------------------------------------------------------------------------
#---- --- Formatter ---
#####---------------------------------------------------------------------------
class TimeSeries_DateFormatter(Formatter):
    """Formats the ticks along a DateArray axis."""

    def __init__(self, freq, minor_locator=False, dynamic_mode=True,):
        self.format = None
        self.freq = freq
        self.locs = []
        self.formatdict = {}
        self.isminor = minor_locator
        self.isdynamic = dynamic_mode
        self.offset = 0
        fgroup = get_freq_group(freq)
        #.....
        if fgroup == _c.FR_ANN:
            self.finder = _annual_finder
        elif fgroup == _c.FR_QTR:
            self.finder = _quarterly_finder
        elif freq == _c.FR_MTH:
            self.finder = _monthly_finder
        elif freq in (_c.FR_BUS, _c.FR_DAY) or fgroup == _c.FR_WK:
            self.finder = _daily_finder

    def asminor(self):
        "Returns the formatter set to minor mode."
        self.isminor = True
        return self

    def asmajor(self):
        "Returns the fromatter set to major mode."
        self.isminor = False
        return self

    def _set_default_format(self, vmin, vmax):
        "Returns the default ticks spacing."
        self.formatdict = self.finder(vmin, vmax, self.freq, False)
        return self.formatdict

    def set_locs(self, locs):
        'Sets the locations of the ticks'
        self.locs = locs
        if len(self.locs) > 0:
            self.verify_intervals()
            self._set_default_format(locs[0], locs[-1])
    #
    def __call__(self, x, pos=0):
        if self.isminor:
            fmt = self.formatdict.pop(x, '')
            if fmt is not '':
                retval = Date(self.freq, value=int(x)).strfmt(fmt)
            else:
                retval = ''
        else:
            retval = ''
        return retval




#####--------------------------------------------------------------------------
#---- --- TimeSeries plots ---
#####--------------------------------------------------------------------------
class TimeSeriesPlot(Subplot, object):
    """Defines a time series based subclass of Subplot."""
    def __init__(self, fig=None, *args, **kwargs):
        """
Accepts the same keywords as a standard subplot, plus a specific `series` keyword.

:Parameters:
    `fig` : Figure
        Base figure.

:Keywords:
    `series` : TimeSeries
        Data to plot

        """
        # Retrieve the series ...................
        _series = kwargs.pop('series',None)
        Subplot.__init__(self,fig,*args,**kwargs)
#        # Force fig to be defined .....
#        if fig is None:
#            fig = TSFigure(_series)
        # Process options .......................
        if _series is not None:
            assert hasattr(_series, "dates")
            self._series = _series.ravel()
            self.xdata = _series.dates
            self.freq = _series.dates.freq
            self.xaxis.set_major_locator

        else:
            self._series = None
            self.xdata = None
            self.freq = None
        self._austoscale = False
        # Get the data to plot
        self.legendsymbols = []
        self.legendlabels = []
    #......................................................
    def set_ydata(self, series=None):
        """Sets the base time series."""
        if self._series is not None:
            print "WARNING ! Base series is being changed."""
        self._series = series.ravel()
        if isinstance(series, TimeSeries):
            self.xdata = self.series.dates
    #....
    def get_ydata(self):
        """Gets the base time series."""
        return self._series
    ydata = property(fget=get_ydata, fset=set_ydata, doc='Time series')
    #......................................................
    def _check_plot_params(self,*args):
        """Defines the plot coordinates (and basic plotting arguments)."""
        remaining = list(args)
        # No args ? Use defaults, if any
        if len(args) == 0:
            if self.xdata is None:
                raise ValueError, "No date information available!"
            return (self.xdata, self.ydata)
        output = []
        while len(remaining) > 0:
            a = remaining.pop(0)
            # The argument is a format: use default dates/
            if isinstance(a,str):
                if self.xdata is None:
                    raise ValueError, "No date information available!"
                else:
                    output.extend([self.xdata, self.ydata, a])
            # The argument is a TimeSeries: use its dates for x
            elif isinstance(a, TimeSeries):
                (x,y) = (a._dates, a._series)
                if len(remaining) > 0 and isinstance(remaining[0], str):
                    b = remaining.pop(0)
                    output.extend([x,y,b])
                else:
                    output.extend([x,y])
            # The argument is a DateArray............
            elif isinstance(a, (Date, DateArray)):
                # Force to current freq
                if self.freq is not None:
                    if a.freq != self.freq:
                        a = a.asfreq(self.freq)
                # There's an argument after
                if len(remaining) > 0:
                    #...and it's a format string
                    if isinstance(remaining[0], str):
                        b = remaining.pop(0)
                        if self.ydata is None:
                            raise ValueError, "No data information available!"
                        else:
                            output.extend([a, self.ydata, b])
                    #... and it's another date: use the default
                    elif isinstance(remaining[0], DateArray):
                        if self.ydata is None:
                            raise ValueError, "No data information available!"
                        else:
                            output.extend([a, self.ydata])
                    #... and it must be some data
                    else:
                        b = remaining.pop(0)
                        if len(remaining) > 0:
                            if isinstance(remaining[0], str):
                                c = remaining.pop(0)
                                output.extend([a,b,c])
                            else:
                                output.extend([a,b])
                else:
                    if self.ydata is None:
                        raise ValueError, "No data information available!"
            # Otherwise..............................
            elif len(remaining) > 0 and isinstance(remaining[0], str):
                b = remaining.pop(0)
                if self.xdata is None:
                    raise ValueError, "No date information available!"
                else:
                    output.extend([self.xdata, a, b])
            elif self.xdata is None:
                raise ValueError, "No date information available!"
            else:
                output.extend([self.xdata, a])
        # Reinitialize the plot if needed ...........
        if self.xdata is None:
            self.xdata = output[0]
            self.freq = self.xdata.freq
        # Force the xdata to the current frequency
        elif output[0].freq != self.freq:
            output = list(output)
            output[0] = output[0].asfreq(self.freq)
        return output
    #......................................................
    def tsplot(self,*parms,**kwargs):
        """Plots the data parsed in argument.
This command accepts the same keywords as `matplotlib.plot`."""
#        parms = tuple(list(parms) + kwargs.pop('series',None))
#        print "Parameters: %s - %i" % (parms, len(parms))
#        print "OPtions: %s - %i" % (kwargs, len(kwargs))
        parms = self._check_plot_params(*parms)
        self.legendlabels.append(kwargs.get('label',None))
        Subplot.plot(self, *parms,**kwargs)
        self.format_dateaxis()
    #......................................................
    def format_dateaxis(self,maj_spacing=None, min_spacing=None,
                        strformat="%Y", rotate=True):
        """Pretty-formats the date axis (x-axis).

:Parameters:
    `major` : Integer *[5]*
        Major tick locator, in years (major tick every `major` years).
    `minor` : Integer *[12]*
        Minor tick locator, in months (minor ticks every `minor` months).
    `strformat` : String *['%Y']*
        String format for major ticks ("%Y").
        """
        # Get the locator class .................
        majlocator = TimeSeries_DateLocator(self.freq, dynamic_mode=True,
                                            minor_locator=False)
        minlocator = TimeSeries_DateLocator(self.freq, dynamic_mode=True,
                                            minor_locator=True)
        self.xaxis.set_major_locator(majlocator)
        self.xaxis.set_minor_locator(minlocator)
        # Get the formatter .....................
        majformatter = TimeSeries_DateFormatter(self.freq, dynamic_mode=True,
                                                minor_locator=False)
        minformatter = TimeSeries_DateFormatter(self.freq, dynamic_mode=True,
                                                minor_locator=True)
        self.xaxis.set_major_formatter(majformatter)
        self.xaxis.set_minor_formatter(minformatter)
        #........................................
#        if rcParams['backend'] == 'PS':
#            rotate = False
#            warnings.warn("dateplot: PS backend detected, rotate disabled")
#        if self.is_last_row():
#            if rotate:
#                setp(self.get_xticklabels(),rotation=45)
    #......................................................
    def set_datelimits(self, start_date=None, end_date=None):
        """Sets the date limits of the plot to start_date and end_date.
    The dates can be given as timeseries.Date objects, strings or integers.
        
:Inputs:
    start_date : var *[None]*
        Starting date of the plot. If None, the current left limit is used.
    end_date : var *[None]*
        Ending date of the plot. If None, the current right limit is used.
        """
        freq = self.freq
        if freq is None:
            raise ValueError("Undefined frequency! Date limits can't be fixed!")
        current_limits = self.get_xlim()
        #
        def get_datevalue(date, freq):
            if isinstance(date, timeseries.Date):
                return date.asfreq(freq).value
            elif isinstance(date, str):
                return timeseries.Date(freq, string=date).value
            elif isinstance(date, (int,float)) or \
                (isinstance(date, numpy.ndarray) and (date.size == 1)):
                return date
            raise ValueError("Unrecognizable date '%s'" % date)
        # Fix left limit ..............
        if start_date is None:
            xleft = current_limits[0]
        else:
            xleft = get_datevalue(start_date, freq)
        # Fix right limit .......
        if end_date is None:
            xright = current_limits[-1]
        else:
            xright = get_datevalue(end_date, freq)
        self.set_xlim(xleft, xright)
        return (xleft, xright)
    

TSPlot = TimeSeriesPlot

def add_yaxis(fsp=None, position='right', yscale=None, basey=10, subsy=None,
              **kwargs):
    """Adds a second y-axis to a plot.
 
:Parameters:
    `fsp` : Subplot *[None]*
        Subplot to which the secondary y-axis is added. If *None*, the current
        subplot is selected
    `position` : String in `('left','right')` *['right']*
        Position of the new axis.
    `yscale` : String, in `('log', 'linear')` *[None]*
        Scale of the new axis. If None, uses the same scale as the first y 
axis
    `basey` : Integer *[10]*
        Base of the logarithm for the new axis (if needed).
    `subsy` : sequence *[None]*
        Sequence of the location of the minor ticks;
        None defaults to autosubs, which depend on the number of decades in 
the plot.  
        Eg for base 10, subsy=(1,2,5) will  put minor ticks on 1,2,5,11,12,15, 
21, ....
        To turn off minor ticking, set subsy=[]
 
    """
    if fsp is None:
        fsp = pylab.gca()
    if not isinstance(fsp, TimeSeriesPlot):
        raise TypeError("The current plot is not a TimeSeriesPlot")
    fig = fsp.figure
    axisini = fsp.axis()
    fsp_alt_args = (fsp._rows, fsp._cols, fsp._num+1)
    fsp_alt = fig.add_tsplot(frameon=False, position=fsp.get_position(),
                             sharex=fsp, *fsp_alt_args)
    # Set position ....................
    if position == 'right':
        (inipos, newpos) = ('left', 'right')
    else:
        (inipos, newpos) = ('right','left')
    # Force scales tics to one side ...
    fsp.yaxis.set_ticks_position(inipos)
    fsp.yaxis.set_label_position(inipos)
    # Force 2nd ticks to the other side..
    fsp_alt.yaxis.set_ticks_position(newpos)
    fsp_alt.yaxis.set_label_position(newpos) 
    # Force period axis scale..........
    if yscale is None:
        yscale = fsp.get_yscale()
        try:
            basey = fsp.yaxis.get_major_locator()._base
        except AttributeError:
            basey = 10.
    fsp_alt.set_yscale(yscale, basey=basey, subsy=subsy)
    # Guess we're good ................
    fsp_alt.set_xticks('')
    fsp_alt.set_xticklabels('')
    
    pylab.draw_if_interactive()
    return fsp_alt

TimeSeriesPlot.add_yaxis = add_yaxis

#####--------------------------------------------------------------------------
#---- --- TimeSeries Figures ---
#####--------------------------------------------------------------------------
class TimeSeriesFigure(Figure):
    """Time Series Figure: all subplots share the same time series.
    """
    def __init__(self, series=None, **kwargs):
        self._series = series
        Figure.__init__(self,**kwargs)
        fspnum = kwargs.pop('fspnum',None)
        if fspnum is not None:
            self.add_tsplot(fspnum, series=series)
    #.........
    def add_tsplot(self, *args, **kwargs):
        """Adds a `TimeSeriesPlot` subplot to the figure."""
        kwargs.update(SubplotClass=TimeSeriesPlot)
        if self._series is not None:
            kwargs.update(series=self._series)
        return add_generic_subplot(self, *args, **kwargs)
    add_plot = add_tsplot
TSFigure = TimeSeriesFigure
#................................................
def tsfigure(series=None, **figargs):
    """Creates a new `TimeSeriesFigure` object.

:Parameters:
    `series` : TimeSeries object
        Input data.
    `figargs` : Dictionary
        Figure options [`figsize`, `dpi`, `facecolor`, `edgecolor`, `frameon`].
    """
    figargs.update(FigureClass=TSFigure)
    figargs.update(series=series)
    fig = pylab.figure(**figargs)
    return fig

def add_tsplot(axes, *args, **kwargs):
    kwargs.update(SubplotClass=TimeSeriesPlot)
    if 'series' not in kwargs.keys():
        if hasattr(axes, 'series'):
            kwargs['series'] = axes.series
        elif hasattr(axes, '_series'):
            kwargs['series'] = axes._series
        else:
            kwargs['series'] = None
    return add_generic_subplot(axes, *args, **kwargs)
Figure.add_tsplot = add_tsplot


def tsplot(series, *args, **kwargs):
    """Plots the series to the current TimeSeries subplot.
    If the current plot is not a TimeSeriesPlot, a new TimeSeriesFigure is created."""
    # allow callers to override the hold state by passing hold=True|False
    b = pylab.ishold()
    h = kwargs.pop('hold', None)
    if h is not None:
        pylab.hold(h)
    # Get the current figure, or create one
    figManager = _pylab_helpers.Gcf.get_active()
    if figManager is not None :
        fig = figManager.canvas.figure
        if not isinstance(fig, TimeSeriesFigure):
            fig = tsfigure(series=series)
    else:
        fig = tsfigure(series=series)
    # Get the current axe, or create one
    sub = fig._axstack()
    if sub is None:
        sub = fig.add_tsplot(111,series=series,**kwargs)
    try:
        ret = sub.tsplot(series, *args, **kwargs)
        pylab.draw_if_interactive()
    except:
        pylab.hold(b)
        raise
    pylab.hold(b)
    return ret

################################################################################
if __name__ == '__main__':

    da = date_array(start_date=Date(freq='B', year=2003, quarter=3, month=1, day=17),
                    length=10)
    ser = timeseries.time_series(MA.arange(len(da)), dates=da)
#    ser[4] = MA.masked
#    ser_2 = timeseries.time_series(MA.arange(len(da)), dates=da.asfreq('Q'))

    pylab.figure()
    pylab.gcf().add_tsplot(111)
    pylab.gca().tsplot(ser, 'ko-')
    pylab.gca().format_dateaxis()
#    pylab.gca().tsplot(ser_2, 'rs')
    pylab.show()