#!/bin/env python

# xy_plot.py
# Author: Andy O'Neill (AON) oneill@phys.columbia.edu
# Begun: 2004-11-15
# A Python/ROOT script to plot a simple xy graph from two columns in a text file
# http://www.nevis.columbia.edu/~oneill/research


# THE FILE FORMAT:
#    The first three lines are the title, x-axis and y-axis titles
#    after that, very line should have two columns (x and y)

# PYTHON IMPORTS
import os, sys    # for exit etc
from math import log10
import root
from ROOT import gROOT, TCanvas, TPad, TH1F, TImage, TGraph    # ETC ETC
from optparse import OptionParser
########################################################################
##                     PARSE ARGUMENTS
########################################################################

do_mean = False
args = root.parse_common_args(sys.argv[1:])

parser = OptionParser(usage=sys.argv[0]+ " [args] filename[s]")
parser.add_option("-m", "--mean", dest="mean", action="store_true", 
                    help="Calculate & print mean of data points")
(options, files) = parser.parse_args(args)

if options.mean:
    do_mean = True

print "Do Mean?", do_mean

def xy_plot(x_array, y_array, title="XY Plot", x_label="X", y_label="Y"):
    if len(x_array) != len(y_array):
        print "ERROR: X and Y data sets must be the same size"
        sys.exit(-1)
    graph = TGraph(len(x_array), x_array, y_array)  # done!
    hist = graph.GetHistogram()  # get this explicity to access the axes
    root.set_labels(hist, title, x_label, y_label)
    #graph.Draw("P")  # allow user to set draw option
    graph.SetMarkerStyle(6)
    graph.SetMarkerSize(1)

    if do_mean and len(y_array) > 0:
        mean = 0
        for y in y_array:
            mean += y
        mean = mean / len(yArray)
        print "Mean of y-axis data is:", mean
    return graph


#############################################################################
##                     PLOT THE GRAPHS
#############################################################################

if __name__ == '__main__':
    for filename in files:
        print "File:", filename
        (title, x_label, y_label) = root.read_header(filename)
        if not root.batch or root.save:
            canvas = root.initialize(filename)
        columns = root.read_columns(filename)
        graph = xy_plot(columns[0], columns[1], x_label, y_label)
        graph.Draw("P")
        root.Update()
        root.save_canvas(filename=filename)
        root.end(globals().copy())