"""
deco.py

Routines and Matplotlib Settings for Various Useful Plot Formats

Lev, 8 Apr 2009
"""

# Changelog
# 31 Jan 2009
# add 'subtext', 'istex', modify 'footer'
#
# 1 Feb 2009
# move 'plotdate' here from 'flib' to speed up loading of 'flib'
#
# 7 Feb 2009
# move 'poly2str' here from 'flib'
# add 'float2tex'
#
# 8 Apr 2009
# add 'eqlims' function
#
# 26 Jun 2009
# add 'setnumxticks()', 'setnumyticks()' functions
#
# 27 Jun 2009
# add 'savefigs()', 'legendline()' functions
#
# 21 Oct 2009
# add 'letter' optional argument to 'portraitA4' and 'landscapeA4'
#
# 14 Jul 2010
# modify 'texoff()' to set sans-serif fonts

import pylab, matplotlib.dates, numpy, matplotlib, matplotlib.ticker
import flib, time, aver
import os, tempfile, math, shutil

def texon(): "Turn TeX rendering on"; pylab.rc('text', usetex=True)

def texoff():
    "Turn TeX rendering off. Set default sans serif fonts"
    pylab.rcParams.update({
        'text.usetex': False,
        'font.family': 'sans-serif',
        'font.fantasy': ['Comic Sans MS', 'Chicago', 'Charcoal', 'ImpactWestern', 'fantasy'],
        'font.monospace': ['Bitstream Vera Sans Mono', 'DejaVu Sans Mono', 'Andale Mono', 'Nimbus Mono L', 'Courier New', 'Courier', 'Fixed', 'Terminal', 'monospace'],
        'font.sans-serif': ['Bitstream Vera Sans', 'DejaVu Sans', 'Lucida Grande', 'Verdana', 'Geneva', 'Lucid', 'Arial', 'Helvetica', 'Avant Garde', 'sans-serif'],
        'font.serif': ['Bitstream Vera Serif', 'DejaVu Serif', 'New Century Schoolbook', 'Century Schoolbook L', 'Utopia', 'ITC Bookman', 'Bookman', 'Nimbus Roman No9 L', 'Times New Roman', 'Times', 'Palatino', 'Charter', 'serif']
    })

def istex(): "Check if TeX rendering is on. Return True/False"; return pylab.rcParams['text.usetex']

def dumprc():
    "Print matplotlib settings in alphabetical order"
    for k in sorted(pylab.rcParams.keys()):
        print '%s: %s' % (k, pylab.rcParams[k])

def datestamp(tm = None, **vargs):
    """
    Add a date stamp to the top right corner of the plot. Use today if 'tm'
    is not specified. Extra keywords are passed to 'pylab.figtext'
    """
    if tm is None: tm = time.time()
    pylab.figtext(0.98, 0.98, 'Plotted ' + flib.dt2str(tm), va='top', ha='right', **vargs)

def pagestamp(text, **vargs):
    """
    Add a page number stamp to the top left corner of the plot. Extra keywords are passed to 'pylab.figtext'
    """
    pylab.figtext(0.02, 0.98, text, va='top', ha='left', **vargs)

def footer(text, **vargs):
    """
    Add a comment to the bottom of the page
    """
    pars = pylab.gcf().subplotpars
    if istex():
        width_inches = pylab.gcf().get_size_inches()[0]
        text = r'\parbox{%.2fin}{%s}' % (width_inches * (pars.right - pars.left), text.replace('\n', r'\\'))
    pylab.figtext(pars.left, 0.01, text, va='bottom', ha='left', **vargs)

def subtext(subplot, text, **vargs):
    """
    Add a text in place of a subplot
    """

    # some code from matplotlib/axes.py
    s = str(subplot)
    if len(s) != 3:
        raise ValueError('Subplot index must be 3 digits long')
    rows, cols, num = map(int, s)

    pars = pylab.gcf().subplotpars
    totWidth = pars.right-pars.left
    totHeight = pars.top-pars.bottom

    figH = totHeight/(rows + pars.hspace*(rows-1))
    sepH = pars.hspace*figH

    figW = totWidth/(cols + pars.wspace*(cols-1))
    sepW = pars.wspace*figW

    rowNum, colNum =  divmod(num-1, cols)

    figBottom = pars.top - (rowNum+1)*figH - rowNum*sepH
    figLeft = pars.left + colNum*(figW + sepW)

    # now my code
    figTop = figBottom + figH
    
    if istex():
        width_inches = pylab.gcf().get_size_inches()[0]
        text = r'\parbox{%.2fin}{%s}' % (width_inches * figW, text.replace('\n', r'\\'))

    pylab.figtext(figLeft, figTop, text, va='top', ha='left', **vargs)

def portraitA4(**vargs):
    """
    Create a figure for an A4 plot in portrait orientation. TeX rendering is not altered.
    
    optional arguments:
    'override' - a dictionary with RC parameters to override
    'letter' - alters the page size from 8 x 11 to 8 x 10.8 inches,
    suitable for american Letter page size
    """

    pylab.rcParams.update({
        'xtick.major.pad': 8,
        'font.size': 12,
        'xtick.labelsize': 'medium',
        'ytick.labelsize': 'medium',
        'legend.fontsize': 'medium',
        'font.family': 'serif' if istex() else 'sans-serif',
#        'font.family': 'sans-serif',
        'font.sans-serif': ['cm', 'sans-serif'],
        'font.serif': ['cm', 'serif'],
        'font.monospace': ['cm', 'monospace'],
        'legend.numpoints': 1,
        'figure.subplot.left': 0.150,
        'figure.subplot.bottom': 0.075,
        'figure.subplot.top': 0.9,
        'figure.subplot.right': 0.9,
        'figure.subplot.wspace': 0.2,
        'figure.subplot.hspace': 0.2,
        'lines.markeredgewidth': 0.5,
        'lines.linewidth': 1.0,
        'axes.linewidth': 1.0
    })

    if 'override' in vargs:
        pylab.rcParams.update(vargs['override'])
        del vargs['override']

    if 'letter' in vargs:
        del vargs['letter']
        vargs['figsize']=(8,10.5)
    else:
        vargs['figsize']=(8,11)
    
    fig = pylab.figure(**vargs)
    datestamp()
    return fig

def landscapeA4(**vargs):
    """
    Create a figure for an A4 plot in landscape orientation. TeX rendering is not altered.
    
    optional arguments:
    'override' - a dictionary with RC parameters to override
    'letter' - alters the page size from 11 x 8 to 10.8 x 8 inches,
    suitable for american Letter page size
    """

    pylab.rcParams.update({
        'xtick.major.pad': 8,
        'font.size': 12,
        'xtick.labelsize': 'medium',
        'ytick.labelsize': 'medium',
        'legend.fontsize': 'medium',
        'font.family': 'serif' if istex() else 'sans-serif',
#        'font.family': 'sans-serif',
#        'font.sans-serif': ['cm'],
#        'font.serif': ['cm'],
#        'font.monospace': ['cm'],
        'font.sans-serif': ['cm', 'sans-serif'],
        'font.serif': ['cm', 'serif'],
        'font.monospace': ['cm', 'monospace'],
        'legend.numpoints': 1,
        'figure.subplot.left': 0.1,
        'figure.subplot.bottom': 0.12,
        'figure.subplot.top': 0.87,
        'figure.subplot.right': 0.95,
        'figure.subplot.hspace': 0.2,
        'figure.subplot.wspace': 0.3,
        'lines.markeredgewidth': 0.5,
        'lines.linewidth': 1.0,
        'axes.linewidth': 1.0
    })

    if 'override' in vargs:
        pylab.rcParams.update(vargs['override'])
        del vargs['override']
    
    if 'letter' in vargs:
        del vargs['letter']
        vargs['figsize']=(10.5,8)
    else:
        vargs['figsize']=(11,8)
    
    fig = pylab.figure(**vargs)
    datestamp()
    return fig

def pubplot(**vargs):
    """
    Create a figure for a publication plot. Default size is 3x1.9 inch, a
    landscape plot with a golden aspect ratio that fits into a column of PRL.
    TeX rendering is turned on.
    """
    
    texon()

    pylab.rcParams.update({
        'xtick.major.pad': 5,
        'font.size': 10,
        'xtick.labelsize': 'small',
        'ytick.labelsize': 'small',
        'legend.fontsize': 'small',
        'font.family': 'serif',
#        'font.sans-serif': ['cm'],
#        'font.serif': ['cm'],
#        'font.monospace': ['cm'],
        'font.sans-serif': ['cm', 'sans-serif'],
        'font.serif': ['cm', 'serif'],
        'font.monospace': ['cm', 'monospace'],
        'legend.numpoints': 1,
        'text.usetex': True,
        'figure.subplot.left': 0.150,
        'figure.subplot.bottom': 0.20,
        'figure.subplot.top': 0.95,
        'figure.subplot.right': 0.95,
        'figure.subplot.wspace': 0.2,
        'figure.subplot.hspace': 0.2,
        'lines.markeredgewidth': 0.3,
        'lines.linewidth': 0.5,
        'axes.linewidth': 1.0
    })

    if 'override' in vargs:
        pylab.rcParams.update(vargs['override'])
        del vargs['override']

    if 'figsize' not in vargs.keys():
        vargs['figsize']=(3, 1.9)
    fig = pylab.figure(**vargs)
    fig.set_size_inches(*vargs['figsize'])
    return fig

def tex_pre_helvetica():
    pylab.rcParams.update({'text.latex.preamble': [
        r'\usepackage[helvet]{sfmath}',
        r'\usepackage[T1]{fontenc}',
        r'\usepackage[scaled]{helvet}',
        r'\renewcommand*\familydefault{\sfdefault}']})

def tex_pre_sfmath():
    pylab.rcParams.update({'text.latex.preamble': [r'\usepackage{sfmath}']})

def tex_pre_empty():
    pylab.rcParams.update({'text.latex.preamble': ['']})

def pubplot_nature(**vargs):
    """
    Create a figure for a publication in Nature. Default size is 3.5x2.2 inch, a
    landscape plot with a golden aspect ratio that fits into a column of text.
    TeX rendering is turned on.
    """
    
    texon()

    pylab.rcParams.update({
        'xtick.major.pad': 4,
        'xtick.major.size': 2,
        'ytick.major.pad': 2,
        'ytick.major.size': 2,
        'font.size': 7,
        'xtick.labelsize': 'small',
        'ytick.labelsize': 'small',
        'legend.fontsize': 'medium',
        'font.family': 'sans-serif',
#        'font.sans-serif': ['cm'],
#        'font.serif': ['cm'],
#        'font.monospace': ['cm'],
#        'font.sans-serif': ['cm', 'sans-serif'],
        'font.serif': ['cm', 'serif'],
#        'font.monospace': ['cm', 'monospace'],
        'legend.numpoints': 1,
        'text.usetex': True,
        'figure.subplot.left': 0.150,
        'figure.subplot.bottom': 0.20,
        'figure.subplot.top': 0.95,
        'figure.subplot.right': 0.95,
        'figure.subplot.wspace': 0.2,
        'figure.subplot.hspace': 0.2,
        'lines.markeredgewidth': 0.3,
        'lines.markersize': 3,
        'lines.linewidth': 0.5,
        'axes.linewidth': 1.0,
        'text.latex.preamble': [r'\usepackage{sfmath}'],
    })

    if 'override' in vargs:
        pylab.rcParams.update(vargs['override'])
        del vargs['override']

    if 'figsize' not in vargs.keys():
        vargs['figsize']=(3.5, 2.5)
    fig = pylab.figure(**vargs)
    fig.set_size_inches(*vargs['figsize'])
    return fig

def overheadplot(**vargs):
    """
    Create a figure for a overhead presentation plot. Default size is 8x6 inch, a
    landscape plot with a golden aspect ratio that fits into a column of PRL.
    TeX rendering is turned on.
    """
    
    texon()

    pylab.rcParams.update({
        'xtick.major.pad': 11,
        'font.size': 20,
        'xtick.labelsize': 'small',
        'ytick.labelsize': 'small',
        'legend.fontsize': 'small',
        'font.family': 'sans-serif',
        'font.sans-serif': ['cm'],
        'font.serif': ['cm'],
        'font.monospace': ['cm'],
        'legend.numpoints': 1,
        'text.usetex': True,
        'figure.subplot.left': 0.10,
        'figure.subplot.bottom': 0.15,
        'figure.subplot.top': 0.95,
        'figure.subplot.right': 0.95,
        'figure.subplot.wspace': 0.2,
        'figure.subplot.hspace': 0.2,
        'savefig.dpi': 300,
        'lines.markeredgewidth': 1.0,
        'lines.linewidth': 1.0,
        'axes.linewidth': 1.5
    })

    if 'override' in vargs:
        pylab.rcParams.update(vargs['override'])
        del vargs['override']

    if 'figsize' not in vargs.keys():
        vargs['figsize']=(8, 6)
    fig = pylab.figure(**vargs)
    fig.set_size_inches(*vargs['figsize'])
    return fig
    
def escapetex(str):
    """
    Escape TeX special characters &, $, %, #, _, {, }, \ in the string
    """
    str = str.replace('\\', '\\textbackslash ')
    for ch in '&$%#_{}':
        str = str.replace(ch, '\\' + ch)
    return str

def traceinfo2tex(info, plm=None, mct=None, paro=None, maxextratime = 300, xmits=[2,4,8,16,32], gains=[10,20], caption=None, show_plm_M0=True):
    """
    Convert trace info to TeX. If 'plm'/'mct'/'paro' is specified,
    temperature/pressure during taking this log is returned.
    'maxextratime' is the maximum extention of time interval around the trace aquisition time
    in a search for thermometry/pressure points.
    'xmits', 'gains' are PLM4 XMIT pulses and GAINs to be considered
    """    
    path = info.path
    
    for loc in ['/data/archive/phpc329/exp', '/mnt/phpc329/exp']:
        if path.startswith(loc):
            path = '...' + path[len(loc):]
            break
        
    path = path.replace('*', '#')
        
    tex = r'\begin{flushleft}'
    
    if caption is not None:
        tex += caption + ': ' #r'\\'
    
    tex += r'\texttt{%s}\\' % (escapetex(path).replace('\\\\', '\\\\\\-').replace('/', '/\\-').replace(' ', '~'))
    if info.Nav > 0:
        tex += '%d averages, ' % info.Nav
    tex += 'captured: %s' % info.timespan()
    if plm is not None:
        for xmit in xmits:
            for gain in gains:
                p = plm.at(xmit = xmit, gain = gain)
                p = p.period(info.starttime - maxextratime, info.finaltime + maxextratime)
                p = p.around(info.starttime, info.finaltime)
                if len(p) < 1: continue
                
                M0 = aver.average(p.M0())
                T = aver.average(p.T())
                tex += r'\\PLM4 xmit %d gain %d: $%s$ [mK]' % (xmit, gain, T.totex(varname='T', fmt='%.3f'))
                if show_plm_M0:
                    tex += r', $%s$' % (M0.totex(varname='M_0', fmt='%.1f'))

    if mct is not None:
        mct = mct.period(info.starttime - maxextratime, info.finaltime + maxextratime)
        mct = mct.around(info.starttime, info.finaltime)
        if len(mct) > 0:
            tex += r'\\MCTNS $%s$ [mK,Greywall] / $%s$ [mK,PLTS high]' % (
                aver.average(mct.T_grwl()).totex(varname='T', fmt='%.2f'),
                aver.average(mct.T_pltsH()).totex(varname='T', fmt='%.1f'))

    if paro is not None:
        paro = paro.period(info.starttime - maxextratime, info.finaltime + maxextratime)
        paro = paro.around(info.starttime, info.finaltime)
        if len(paro) > 0:
            tex += r'\\Paroscientific: $%s$ [mbar]' % (aver.average(paro.P()).totex(varname='P', fmt='%.1f'))
            
    tex += r'\end{flushleft}'
    return tex

def plotdate(x, y, fmt='.', tz=None, xdate=True, ydate=False, datefmt='%d/%m/%y %H:%M:%S', axes=None, rotation=90, **kwargs):
    """
    A version of pylab.plot_date()
    Use a matplotlib.dates.DateFormatter with datefmt format.
    Call autofmt_xdate(rotation=45, ha='center') for a current figure.
    """
    if axes is None:
        axes = pylab.gca()
#        pylab.axes(axes)

    p = axes.plot_date(x, y, fmt, tz, xdate, ydate, **kwargs)
    fmtr = matplotlib.dates.DateFormatter(datefmt, tz)

    if xdate:
        axes.xaxis.set_major_formatter(fmtr)
        axes.autoscale_view(True, True, False)
        axes.get_figure().autofmt_xdate(rotation=rotation, ha='center')
    
    if ydate:
        axes.yaxis.set_major_formatter(fmtr)
        axes.autoscale_view(True, False, True)

    pylab.draw_if_interactive()

    return p

def poly2str(p, argname='x', format='%+g'):
    """
    Convert a list of polynomial coefficients 'p' into a TeX-style string using
    'argname' for the argument. 'p' is assumed to be in power-decending order,
    as returned by 'numpy.polyfit' routine.
    
    When overriding coefficient format 'format', please remember to use '+' in
    the format abbreviation that puts '+' signes in front of positive
    coefficients. No extra '+' are added between terms in the polynomial.
    
    '+' is automatically deleted in front of the first argument, '0' is
    returned for a zeroth-order polynomail (empty array 'p')
    The string returned is not surrounded by math-mode delimiters (e.g. '$').
    """
    N = len(p)
    s = ""
    if N > 2: # higher powers
        for i in range(0, N-2):
            s += format%(p[i]) + argname +  '^%d'%(N-i-1)

    if N > 1: s += format%(p[-2]) + argname # linear term

    if N > 0: s += format%(p[-1]) # constant
    
    if N == 0: s = '0' # zero-order polynomial case

    if s[0] == '+': s = s[1:] # drop leading '+' in front of p[0]

    return s

def float2tex(n, fmt='%.3f', min_order=-1, max_order=2):
    """
    Convert a number to scientific representation in TeX 
    """
    power10 = numpy.floor(numpy.log10(abs(n)))    
    if power10 < min_order or power10 > max_order:
        n /= 10 ** power10
        return (fmt % n) + r'\times 10^{%d}' % power10
    else:
        return fmt % n
    
def eqlims(axes_and_figures, x = False, y = False, x_inverse = False, y_inverse = False):
    """
    Equalise limits of all given axes and all given figures to show all the data.
    'axes_and_figures' - an array of axes and figures to modify, can be a figure, then all subplots are adjusted
    'x','y' - determine whenever to adjust x/y range
    'x_inverse', 'y_inverse' - determine whenever the range should be from small to large or
    the other way round.
    """
    
    axes = []

    if isinstance(axes_and_figures, pylab.Figure):
        axes_and_figures = [axes_and_figures]
    for a in axes_and_figures:
        if isinstance(a, matplotlib.axes.Axes):
            axes.append(a)
        elif isinstance(a, matplotlib.figure.Figure):
            axes += a.axes
        else:
            raise ValueError("'axes_and_figures' contains an element that is neither a set of axes nor a figure")
    
    if len(axes) < 1:
        return
    
    xlim = numpy.concatenate([a.get_xlim() for a in axes])
    ylim = numpy.concatenate([a.get_ylim() for a in axes])

    xlim = (max(xlim), min(xlim)) if x_inverse else (min(xlim), max(xlim))
    ylim = (max(ylim), min(ylim)) if y_inverse else (min(ylim), max(ylim))

    for a in axes:
        if x:
            a.set_xlim(xlim)
        if y:
            a.set_ylim(ylim)
#        a.figure.show()

    return numpy.vstack([xlim, ylim])

def setnumxticks(n):
    """
    Set a LinearLocator with 'n' ticks as the current x-axis locator
    """
    pylab.gca().xaxis.set_major_locator(matplotlib.ticker.LinearLocator(numticks=n))
    pylab.draw()

def setnumyticks(n):
    """
    Set a LinearLocator with 'n' ticks as the current y-axis locator
    """
    pylab.gca().yaxis.set_major_locator(matplotlib.ticker.LinearLocator(numticks=n))
    pylab.draw()

def savefigs(filename, *args, **vargs):
    """
    Save the current figure as PDF and PNG
    the 'filename' is augmented by '.pdf'/'.png' extension.
    All other arguments are passed to 'pylab.savefig()'
    """
    pylab.rcParams['savefig.dpi'] = 600
    
    pylab.savefig(filename + '.pdf', *args, **vargs)
    pylab.savefig(filename + '.png', *args, **vargs)

def legendline(text):
    """
    Plot a void line in order to add a line to the legend.
    The legend must autogenerated for this line to be present.
    """
    return pylab.plot([], [], ls='None', marker='None', label=text)

def set_colormap(cmap):
    '''
    Set the default colormap to a given value and apply to current image if any.
    '''
    pylab.rc('image', cmap=cmap.name)
    im = pylab.gci()

    if im is not None:
        im.set_cmap(cmap)
    pylab.draw_if_interactive()

def save_pdf_book(outfilename, figures):
    """
    Save an array of figures as a multipage PDF. Requires "pdfjoin/pdfjam" tool.
    """
    curdir = os.getcwd()
    try:
        tmpdir = tempfile.mkdtemp()
        try:
            page = 1
            fmt = 'page%%0%dd.pdf' % int(math.ceil(math.log10(len(figures))) + 1)
            for f in figures:
                f.savefig(os.path.join(tmpdir, fmt % page))
                page += 1
            
            os.chdir(tmpdir)
            os.system('pdfjoin --outfile book.pdf page*.pdf')
            os.chdir(curdir)
            shutil.move(os.path.join(tmpdir, 'book.pdf'), outfilename)
        finally:
            try:
                os.rmdir(tmpdir)
            except:
                pass
    finally:
        os.chdir(curdir)

def aver_plot(x, y, Nav, *args, **vargs):
    """
    Plot x vs y using pylab.plot averaging sets of Nav points in both arrays using flib.average
    """
    return pylab.plot(flib.average(x, Nav), flib.average(y, Nav), *args, **vargs)