#!/bin/env python

# aon_root.py
# Author: Andy O'Neill <oneill@phys.columbia.edu>
# 2006-07-13

# Common functions for plotting using pyROOT


# THE FILE FORMAT:
#    The first lines is the title
#    after that, very line should have three columns (t, V, phi)

# PYTHON IMPORTS
import os, sys    # for exit etc
import re     # regex support
from math import log10
from locale import atof, atoi
from time import gmtime, strftime, localtime,sleep # USING THIS TO CREATE UNIQUE
                                                   # FILENAMES
from ROOT import gROOT, TCanvas, gStyle 
import code  # for interactive console

from optparse import OptionParser
from array import array

save = False
batch = False
header = False
main_canvas = None
program_name = os.path.basename(sys.argv[0])

def parse_common_args(argv):
    global save, batch, header
    parser = OptionParser()
    parser.add_option("-s", "--save", dest="save", action="store_true",
                        help="Save the canvas to file")
    parser.add_option("-b", "--batch", dest="batch", action="store_true",
                        help="Don't prompt the user at all")

    parser.add_option("-H", "--header", dest="header", action="store_true",
                        help="File[s] contain[s] three lines of header info")

    (options, args) = parser.parse_args(argv)
    if options.save:
        save = True
    if options.batch:
        batch = True
    if options.header:
        header = True
    print "Header?", header, "\nBatch?", batch, "\nSave?", save
    return args

def Update(canvas=None):
    global batch
    if batch:
        return
    global main_canvas
    if canvas == None:
        canvas = main_canvas
    try:
        canvas.Update()
        canvas.Modified()
    except AttributeError:
        print "You have closed the canvas, so I'm exiting"
        sys.exit()

                                                                    

def initialize(canvas_name=""):
    global main_canvas, program_name
    print "Initializing ROOT canvas... ",
    gROOT.Reset()
    c1 = TCanvas(canvas_name, program_name + " : " + canvas_name, 1024, 768)
    c1.SetGridx()
    c1.SetGridy()
    c1.GetFrame().SetFillColor(21)
    c1.GetFrame().SetBorderSize(12)
    
    gStyle.SetTitleW(0.4)
    gStyle.SetTitleBorderSize(0)
    main_canvas = c1
    print "Done."
    return c1


def read_header(filename):
    if header:
        lines = open(filename).readlines()
        return (lines[0], lines[1], lines[3])
    else:
        return ("Filename", "X", "Y")

def read_columns(filename):
    """Read file as columns"""
    file = open(filename)
    columns = []
    if header:
        lines = file.readlines()[3:]
    else:
        lines = file.readlines()
    ncols = len(lines[0].split())
    print "Found", ncols, "columns"
    for i in range(0, ncols):
        columns.append(array('d'))
    for line in lines:
        split = line.split()
        for i in range(0, ncols):
            try:
                columns[i].append(atof(split[i]))
            except ValueError:
                print "Looks like there is a header in this file " + \
                        """- try using "--header" option"""
                sys.exit(-1)
    return columns
        
def max(array):
    """return maximum value in array"""
    max = -1e99
    for i in array:
        if i > max:
            max = i
    return max 

def min(array):
    """return minimum value in array"""
    min = 1e99
    for i in array:
        if i < min:
            min = i
    return min

def get_label(label_name, current_label=""):
    print "Enter new " + label_name + " or hit return to keep the old one"
    print ("Current " + label_name + ": " + current_label)
    new_label = raw_input("New " + label_name + ": ")
    if new_label == "":
        return current_label
    else:
        return new_label

def set_labels(hist=None, title="", x_label="", y_label=""):
    if hist == None:
        print "trying to set labels of non-existant histogram"
        return
    if not batch:
        while title == "":
            title = get_label("title")
        while x_label == "":
            x_label = get_label("X axis")
        while y_label == "":
            y_label = get_label("Y axis")
    hist.SetTitle(title)
    hist.SetXTitle(x_label)
    hist.SetYTitle(y_label)
    if not batch:
        hist.Draw()

# define a function to check user entry
def check_extension(filename):
    # get the file extension
    extension = filename.split('.')[-1]

    # do a regex match for good filetypes
    if not re.match ("eps|svg|gif|pdf|ps", extension):
        print "The filename you provide must have an extension of " \
              "eps, ps, pdf, gif or svg" 
        return ""         # will ask again
    else:
        return extension  # good to go

def set_line_properties(line = None, style=None, color=None, width=None,
        marker=None):
    """Set the line properties of any object that inherits TAttLine"""
    if line == None:
        return
    if marker != None:
        line.SetMarkerStyle(marker)
        
def save_canvas(canvas=None,filename=""):
    global main_canvas
    if not save:
        return
    if canvas == None:
        canvas = main_canvas
    if filename == "":
        filename = canvas.GetName()
    filename=filename+".ps"
    extension = check_extension(filename)
    if not batch:
        extension = ""
        # prompt the user to accept or enter a new one
        while (extension == ""):
            print "Enter filename or press enter for default ["+filename+"]"
            new_filename = raw_input ("New Name: ")
            if not new_filename == "":
                filename = new_filename  # only if they entered something
            # and make sure the extension is OK, if so we wil continue
            extension = check_extension(filename)
    Update()
    # now we should have a good filetype - use ROOT to save it
    canvas.Print(filename, "Landscape,"+extension)


def end(namespace):
    Update()
    if batch:
        sleep (2)
    else:
        input = raw_input ("Enter to quit, \"Y\" for a console: ")
        if input == "y" or input == "Y":
            code.interact(banner="Do your worst!", local=namespace)