# -*- coding: iso-8859-1 -*-
#
# quakepy/qpplot.py
# $Id: qpplot.py 156 2009-02-13 12:16:25Z fab $
#
# The QuakePy package
# http://www.quakepy.org
#

############################################################################
#    Copyright (C) 2007-2009 by Fabian Euchner and Danijel Schorlemmer     #
#    fabian@fabian-euchner.de                                              #
#                                                                          #
#    This program is free software; you can redistribute it and#or modify  #
#    it under the terms of the GNU General Public License as published by  #
#    the Free Software Foundation; either version 2 of the License, or     #
#    (at your option) any later version.                                   #
#                                                                          #
#    This program is distributed in the hope that it will be useful,       #
#    but WITHOUT ANY WARRANTY; without even the implied warranty of        #
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the         #
#    GNU General Public License for more details.                          #
#                                                                          #
#    You should have received a copy of the GNU General Public License     #
#    along with this program; if not, write to the                         #
#    Free Software Foundation, Inc.,                                       #
#    59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.             #
############################################################################

"""
The QuakePy package
http://www.quakepy.org
"""

__version__  = '$Id: qpplot.py 156 2009-02-13 12:16:25Z fab $'
__revision__ = '$Revision: 156 $'
__author__   = "Fabian Euchner <fabian@fabian-euchner.de>, Danijel Schorlemmer <ds@usc.edu>"
__license__  = "GPL"

import datetime
import matplotlib
from matplotlib import rcParams
from matplotlib.dates import MO
from matplotlib.dates import YearLocator, MonthLocator, DayLocator, WeekdayLocator, DateFormatter
matplotlib.use('PS')

from pylab import *

class QPPlot( object ):
    
    autoscale = False
    showgrid = True

    def __init__( self ):
        rcParams['figure.figsize'] = (8, 8)
    
    def plot( self, imgfile, abscissa_in, ordinate_in, **kwargs ):
        autoscale = False
        showgrid = True
    
        clf()
        
        #line_style = 'k-' # N-test - black solid line
        line_style = 'ks' # N-test - black solid line
        
        ax = subplot(111)
        
        if 'type' in kwargs.keys() and kwargs['type'] == 'fmd':
       
            semilogy( abscissa_in, ordinate_in, line_style )
        
            xlabel('Magnitude')
            ylabel('Number of Events')
        
        show()
        
        try:
            savefig( imgfile )
        except IOError, e:
            print "could not write image to file: ", e
            return None
            
    def plot_vs_date( self, imgfile, abscissa_in, ordinate_in, **kwargs  ):
        
        autoscale = False
        showgrid = True
    
        clf()
        line_style = 'r-' # N-test - black solid line
    
        # get values for abscissa, only date part (first 10 chars)
        abscissa = []
        for absc in abscissa_in:
            absc_parts = absc[0:10].split('-')
            abscissa_date = datetime.date( int(absc_parts[0]), int(absc_parts[1]), int(absc_parts[2]) )
            abscissa.append( date2num( abscissa_date ) )
            
        # abscissa (date axis) formatting / define tick locators and formatters
        # NOTE: we expect the dates in <testDate> to be in ascending order
            
        # get difference in years
        startyear = int( num2date(abscissa[0]).strftime('%Y') )
        endyear   = int( num2date(abscissa[len(abscissa)-1]).strftime('%Y') )
        yeardiff  = endyear - startyear
        
        if ( yeardiff >= 15 ):
            abscissaFmt  = DateFormatter( '%Y' )
            abscMajorLoc = YearLocator( 10, month=1, day=1 )
            abscMinorLoc = YearLocator( 2, month=1, day=1 )
        elif ( yeardiff >= 8 ):
            abscissaFmt  = DateFormatter( '%Y' )
            abscMajorLoc = YearLocator( 5, month=1, day=1 )
            abscMinorLoc = YearLocator( 1, month=1, day=1 )
        elif ( yeardiff >= 3 ):
            abscissaFmt  = DateFormatter( '%Y' )
            abscMajorLoc = YearLocator()
            abscMinorLoc = MonthLocator( (1, 4, 7, 10) )
        elif ( yeardiff >= 1 ):
            abscissaFmt  = DateFormatter( '%b %Y' )
            abscMajorLoc = MonthLocator( (1, 4, 7, 10) )
            abscMinorLoc = MonthLocator()
        else:
            # get month difference
            startmonth = int( num2date(abscissa[0]).strftime('%m') )
            endmonth   = int( num2date(abscissa[len(abscissa)-1]).strftime('%m') )
            monthdiff  = endmonth - startmonth
            if ( monthdiff >= 6 ):
                abscissaFmt  = DateFormatter( '%b %Y' )
                abscMajorLoc = MonthLocator( (1, 3, 5, 7, 9, 11) )
                abscMinorLoc = MonthLocator()
            elif ( monthdiff >= 2 ):
                abscissaFmt  = DateFormatter( '%b %Y' )
                abscMajorLoc = MonthLocator()
                abscMinorLoc = WeekdayLocator( MO )
            elif ( monthdiff >= 1 ):
                abscissaFmt  = DateFormatter( '%b %d' )
                abscMajorLoc = WeekdayLocator( MO )
                abscMinorLoc = DayLocator()
            else:
                # all dates in the same month
                # get difference of days
                startday = int( num2date(abscissa[0]).strftime('%d') )
                endday   = int( num2date(abscissa[len(abscissa)-1]).strftime('%d') )
                daydiff  = endday - startday
                if ( daydiff > 10 ):
                    abscissaFmt  = DateFormatter( '%b %d' )
                    abscMajorLoc = WeekdayLocator( MO )
                    abscMinorLoc = DayLocator()
                else:
                    # range is smaller than 10 days, major tick every day
                    abscissaFmt  = DateFormatter( '%b %d' )
                    abscMajorLoc = DayLocator()
                    abscMinorLoc = NullLocator()
                            
        # get ordinate
        ordinate       = map( float, ordinate_in )
            
        # check for equal array dimension
        if ( len(abscissa_in) != len(ordinate) ):
            raise IndexError, "error: data vector length mismatch"
            
        
        # use AUTO for ordinate formatting
        # ordinateFmt = FormatStrFormatter( '%0.3f' )
        # ax.yaxis.set_major_formatter( ordinateFmt )
                
        # determine y axis range
        ymin = 0.9 * min( ordinate )
        ymax = 1.1 * max( ordinate )
        
        ax = subplot(111)
        
        plot_date( abscissa, ordinate, line_style )
        # ylabel( 'Number of events' )
        ylabel( 'mean completeness' )
            
        # check if abscissa has only one data point: set x axis range explicitly
        # (plus / minus one day)
        if ( len( abscissa ) == 1 ):
            xlim( abscissa[0]-1, abscissa[0]+1 )
            
        # set y axis range - adjust to the '0.1'-bin
        ymin = 0.1 * floor( 10.0 * ymin )
        ymax = 0.1 * ceil( 10.0 * ymax )
        
        # avoid that ymin equals ymax
        # currently this is only possible if the whole dataset is zero
        # in this case, set range to -0.1 ... 0.1
        if ( ymin == ymax ):
            ymin -= 0.1
            ymax += 0.1
            
        ylim( ymin, ymax )
            
        # formatting of abscissa (date) axis
        ax.xaxis.set_major_formatter( abscissaFmt )
        ax.xaxis.set_major_locator( abscMajorLoc )
        ax.xaxis.set_minor_locator( abscMinorLoc )
            
        if ( autoscale == True ):
            ax.autoscale_view()
            
        labels = ax.get_xticklabels()
        setp( labels, 'rotation', 30, fontsize=14 )
        #grid( showgrid )
        
        show()
        
        try:
            savefig( imgfile )
        except IOError, e:
            print "could not write image to file: ", e
            return None
    
    
    