"""
A module for plotting ASCII or AIPS++/CASA tables using CASA.

News:
     1/22/2008: simpleplot()'s output is much more customizable now.
                It passes its keyword arguments to plotfunc, and takes a
                list of symbols for plotting each y column.  Column labels
                are now filtered in case you are using LaTeX.

Use: From casapy,
         # Preferred
             execfile 'tabplotter.py'

    example:
       # Set btb to be a tb tool filled with the contents of
       # ./cfgres_10to20.csv.  The contents will be
       # fetched from ./cfgres_10to20.tb if it already
       # exists, and ./cfgres_10to20.csv otherwise.
       btb = filetotable('cfgres_10to20.csv')

       # Normal tb commands work.
       btb.browse()               # You may need btb.clearlocks()
       btb.colnames()

       # As does help.
       help simpleplot
       
       # Plot everything in btb except the first column vs. the first column,
       # with a logarithmic y axis.
       simpleplot(btb, plotfunc=pl.semilogy)
       
       # Plot Nom.rms.res vs. Actual_res.
       simpleplot(btb, 'Nom.rms.res', 'Actual_res')

       # Plot Nom.rms.res and Actual_res vs. Nom.avg.res
       simpleplot(btb, ['Nom.rms.res', 'Actual_res'], 'Nom.avg.res')

       # timeaxis() is just a helper function now, but Gene DuVall might
       # extend it to use minutes, hours, etc., and/or to supply a time axis
       # for customized plots.

       # tablename() returns the table's filename without any directory or
       # .csv. 
"""
__version__ = 0.03
__author__  = "Rob Reid, rreid@nrao.edu, starting from an example by D. Shepherd"

import os, re

def uniquify_strlist(duplist):
    """Goes through the list of strings duplist, and appends a count number
    to any duplicates.  i.e. given
    
      ['a',   'b', 'c',   'a',   'c',   'd', 'a',   'e'],
      
      it returns
      
      ['a_0', 'b', 'c_0', 'a_1', 'c_1', 'd', 'a_2', 'e'].

      Note that

      ['a',   'a', 'a_0']

      would go to

      ['a_0', 'a_1', 'a_0'].

      If you have a list like that, you could run uniquify_strlist on the
      result.
    """
    strlist = duplist[:]           # Make a copy and leave duplist alone.
    strdict = {}
    for i in range(len(strlist)):
        s = strlist[i]
        if s in strdict:
            strdict[s].append(i)
        else:
            strdict[s] = [i]
    for s in strdict:
        if len(strdict[s]) > 1:
            for j in range(len(strdict[s])):
                strlist[strdict[s][j]] += "_%d" % j
            strdict[s] = []
    return strlist
    
def filetotable(fn, overwrite=False):
    """
    Returns a CASA table (tb) tool as read from the comma separated value (CSV)
    file fn.  Built around the fromascii line by Debra Shepherd.

    The first row in the CSV is assumed to be, and used as, column names,
    not as data values.  If you want to use the tb for plotting, all other
    rows should be numbers. 

    I usually produce my tables with whitespace separation for readability, but
    they can be easily converted to CSV format with this bash alias:

    # Use: csvize < table.readable > table.csv
    alias csvize='perl -wp -e "s/[ \t]+/,/g"'
    
    csvize will silently FAIL if there are any commas in table.readable!
    """

    tbname = fn.replace('.csv', '')  # Do it this way instead of one
    tbname += '.tb'                  # replace in case fn doesn't end
                                     # in .txt.

    mytb = tbtool.create()           # tb is persistently global!

    if os.path.exists(tbname) and not overwrite:
        mytb.open(tbname)
        return mytb

    ##  Fill (load) the data into CASA.
    # autoheader=True instructs mytb to create its own headers
    # (called Column1, Column2, etc) because the headers cannot be
    # interpreted.
    mytb.fromascii(tablename = tbname, asciifile = fn,
                   sep=',', autoheader = True, nomodify = false)
    mytb.close()                        # Probably NOT modifiable now.
    mytb.open(tbname, nomodify=false)   # Yes, I want to modify it!

    # The first row is column names.  Move it elsewhere.
    boringcns = mytb.colnames()              # Column1, Column2, ...

    # CASA doesn't like column names with spaces.
    colnames = [mytb.getcell(bcn, 0).replace(' ',
                                             '_') for bcn in boringcns]
    # This isn't just paranoia.
    colnames = uniquify_strlist(colnames)

    mytb.removerows([0])

    for (bcn, ncn) in zip(boringcns, colnames):
        mytb.renamecol(bcn, ncn)

    mytb.flush()
    return mytb

def timeaxis(datatb):
    """
    Attempts to return a reasonable time axis and label for it from datatb.
    """
    sec = datatb.getcol('Seconds').astype(float)

    t = sec           # REMOVE when things get smarter!
    tlabel = "Time (s)"

    tint = sec[-1] - sec[0]
    if tint < 300.0:
        t = sec                        # Leave it in seconds
        tlabel = "Time (s)"

    return t, tlabel

def tablename(datatb):
    """
    Returns a slightly cleaned up name for datatb.
    """
    tablab = datatb.name()
    tablab = re.sub(r"^.*/", '', tablab)
    tablab = re.sub(r"\.tb", '', tablab)
    return tablab

def simpleplot(datatb, ycolnames=[], xcolname='', colnames=[],
               plotfunc=pl.plot, symblist=['b-'], **kwargs):
    """
    Graph datatb[ycolname] vs. datatb[xcolname] for each ycolname in ycolnames,
    on the same plot.

    xcolname defaults to the first column in colnames or datatb.colnames().

    ycolnames defaults to all but xcolname in colnames or
    datatb.colnames().

    symblist is a list of symbols to plot each y column with, in plotfunc's
    format.

    **kwargs is passed on to plotfunc.
    """

    # clear plotter
    pl.clf()

    # turn interactive mode ON (allows automatic updates of plots)
    pl.ion()

    # Handle defaults.
    if len(colnames) == 0:
        colnames = datatb.colnames()
    if len(xcolname) == 0:
        xcolname = colnames[0]
    if len(ycolnames) == 0:
        ycolnames = colnames[:]
        ycolnames.remove(xcolname)

    x = datatb.getcol(xcolname)
    
    if isinstance(ycolnames, str):
        ycns = [ycolnames]
    else:
        ycns = ycolnames

    if isinstance(symblist, str):
        symbs = [symblist]
    else:
        symbs = symblist

    ylabs = map(raw_to_latex, ycns)
    
    # map(None, l1, l2, ...) is like zip with padding by None instead of
    # truncation.
    for ycolname, ylab, symb in map(None, ycns, ylabs, symbs):
        y = datatb.getcol(ycolname)
        if not symb:
            symb = symbs[0]
        plotfunc(x, y, symb, label = ylab, **kwargs)
        
    if len(ylabs) == 1:
        pl.ylabel(ylabs[0])
    else:
        pl.legend(loc='best')

    pl.xlabel(raw_to_latex(xcolname))
    pl.title(tablename(datatb))

def raw_to_latex(a):
    """
    Returns a copy of s that has been converted from plain text to valid LaTeX
    input.  Do not run on strings that already have LaTeX markup!
    """
    s = a.replace('_', r'$\_$')
    s = s.replace('#', r'\#')
    return s