#!/usr/bin/env python

# vale cofer-shabica, 2014
# [first name]@brown.edu

# a python program with matplotlib to quickly generate publication quality probability densities
# try: ./hist data.dat --out test.eps
# where data.dat has a single collumn of values.
# open the generated file, test.eps

# ./hist --help should also be instructive

import matplotlib
import numpy as np
import argparse
import sys

# Specify renderer
matplotlib.use('Agg')

# Boiler-plate settings for producing pub-quality figures
# 1 point = 1/72 inch
# all these settings could be stored in:
# $HOME/.config/matplotlib/matplotlibrc
# which render this section unnecessary
matplotlib.rcParams.update({'figure.figsize'  : (8,5)    # inches
                            ,'font.size'       : 22      # points
                            ,'legend.fontsize' : 16      # points
                            ,'lines.linewidth' : 2       # points
                            ,'axes.linewidth'  : 2       # points
                            ,'text.usetex'     : True    # Use LaTeX to layout text
                            ,'font.family'     : "serif" # Use serifed fonts
                            ,'xtick.major.size'  : 6     # length, points
                            ,'xtick.major.width' : 2     # points
                            ,'xtick.minor.size'  : 3     # length, points
                            ,'xtick.minor.width' : 1     # points
                            ,'ytick.major.size'  : 6     # length, points
                            ,'ytick.major.width' : 2     # points
                            ,'ytick.minor.size'  : 3     # length, points
                            ,'ytick.minor.width' : 1     # points
                            # Define a list of fonts to use for each style; no need to edit these
                            ,'font.serif'      : ("Times"
                                                  ,"Palatino"
                                                  ,"Computer Modern Roman"
                                                  ,"New Century Schoolbook"
                                                  ,"Bookman")
                            ,'font.sans-serif' : ("Helvetica"
                                                  ,"Avant Garde"
                                                  ,"Computer Modern Sans serif")
                            ,'font.monospace'  : ("Courier"
                                                  ,"Computer Modern Typewriter")
                            ,'font.cursive'    : "Zapf Chancery"
})

# these imports must come *after* the call to matplotlib.rcParams.update()
import matplotlib.pyplot as plt
import matplotlib.ticker as T

# if you're playing with the options for tic formatting, try the following import
#from matplotlib.ticker import ScalarFormatter, FormatStrFormatter

# main method, if you're only interested in matplotlib, skip to plotADensity()
def main():
    parser = argparse.ArgumentParser("Generate a quick, dirty, and publication quality probability density")
    parser.add_argument("file", nargs='+', type=np.loadtxt, help="data source(s)")
    parser.add_argument("--bins",  type=int, help="specify number of bins", default=25, metavar="N")
    parser.add_argument("--title",  help="figure title", default="title")
    parser.add_argument("--alabel",  help="abscissa label", default="abscissa")
    parser.add_argument("--units",  help="units for abscissa and ordinate", default=None)
    parser.add_argument("--out",  help="filename to save", default="hist.eps", metavar="FILENAME")
    parser.add_argument("--nozero", action='store_true', help="don't set origin to 0", default=False)
    parser.add_argument("--legend", action='store_true', help="include a legend (requires --labels)", default=False)
    parser.add_argument("--labels", nargs='+', type=str,  help="labels for each data set", metavar="label")
    parser.add_argument("--styles", nargs='+', type=str,  help="styles for each data set", metavar="style")

    opt = parser.parse_args()
    
    if opt.legend:
        if opt.labels == None or (len(opt.file) != len(opt.labels)):
            parser.print_help()
            sys.exit()
    else:
        opt.labels=['' for x in range(len(opt.file))]

    if opt.styles != None:
        if len(opt.styles) != len(opt.file):
            parser.print_help()
            sys.exit()        

    plotADensity(
        opt.file
        , opt.bins
        , opt.title
        , opt.alabel
        , opt.units
        , opt.out
        , opt.legend
        , opt.labels
        , zeroing = not opt.nozero
        , styles = opt.styles
    )
    print("Wrote " + opt.out)

def centers(e):
    """given a list of edges, returns a list of centers of the bins"""
    return [ np.mean(p) for p in zip(e[1:], e[:-1]) ] 

# see here for matplotlib usage
def plotADensity(sets, nbins, title, xlabel, units, fname, legend, labels, styles=None, zeroing=True):
    plt.ioff() # disable interactive mode
    plt.clf()  # clear
    plt.cla()  # really, clear it all!
    ax=plt.gca() # add axes

    # this is a bit ugly, re: styles
    if styles == None:
        # plot each datum
        for (datum, legLabel) in zip(sets,labels):
            n, e = np.histogram(datum[:], bins=nbins, density = True)
            c = centers(e)
            plt.plot(c, n, label=legLabel)
    else:
        for (datum, legLabel, dStyle) in zip(sets,labels, styles):
            n, e = np.histogram(datum[:], bins=nbins, density = True)
            c = centers(e)
            plt.plot(c, n, dStyle, label=legLabel)

    plt.title(title)

    if units != None:
        ylabel_fmted = "Probability Density $\cdot$ " + units
        xlabel_fmted = xlabel + " / " + units
    else:
        ylabel_fmted = "Probability Density"
        xlabel_fmted = xlabel
    
    plt.ylabel(ylabel_fmted)
    plt.xlabel(xlabel_fmted)
    
    if legend:
        plt.legend()

    if zeroing:
        x_min, x_max = plt.xlim()
        plt.xlim(xmin=0, xmax=(x_max*1.15))
        plt.ylim(ymin=0)


    # Ticks at nice locations
    ax.yaxis.set_major_locator(T.MaxNLocator(5))
    ax.xaxis.set_major_locator(T.MaxNLocator(7))

    # Format the ticks sanely
    #formatter = FixedOrderFormatter(order)
    #formatter.set_scientific(True)
    #formatter.set_powerlimits((-2,3))
    #ax.yaxis.set_major_formatter(ScalarFormatter())
    #ax.tick_param(width=3)
    
    plt.tight_layout()
    plt.savefig(fname)

main()