#!/usr/bin/python
"""Example usage:
>>> from fileplotter import FilePlotter
>>> tt = FilePlotter(title="Crosstalk probabilities",show_title=True,show_legend=True)
>>> tt.addFile("crosstalk_050.txt",title=True,axislabels=True,legend=True,xyplot=True)
x-Axis:  Overvoltage (V)
y-Axis:  Crosstalk probability
Labels:  ['', '050C']
>>> tt.addFile("crosstalk_100.txt",title=True,axislabels=True,legend=True,xyplot=True)
Labels:  ['', '100C']
>>> tt.show()
"""
"""Usage: python plotFile.py [-h|--help] [-x|--xyplot] filename
Options:
-h|--help: Show this help
-x|--xyplot: Use first columns as x-values
-t|--title: Use first three lines of the file as title, x-label, y-label
-l|--legend: Show legend from column titles
"""
import os,sys
import getopt
import numpy as np
from math import *
from ROOT import TCanvas, TMultiGraph, TGraph, TGraphErrors
from ROOT import gROOT, gStyle
from ROOT import TH2F

gROOT.SetStyle("Plain")

#*******************************************************************************
class FilePlotter:
	def __init__(self,show_legend=True,show_title=False,yrange=False,title=None,title_x=None,title_y=None,out_name=None,xmin=None,xmax=None,show_line=False):
		#Todo: doc-strings
		#ToDo: implement file calling (ie. "plotFile.py data.txt")
		#Todo: format size of the legend (self.canvLegend.GetNRows())
		#Todo: implement x-min/max > datapoints...
		#Todo: implement errors
		
		#-------------------------------------
		#Read constructor arguments
		#-------------------------------------
		self.show_legend = show_legend #Whether to show the legend, will be put to false if one of the files lacks legend information
		self.show_title = show_title
		self.show_line = show_line
		
		self.yrange = False
		self.ymin = np.nan
		self.ymax = np.nan
		
		self.title = title
		self.title_x = title_x
		self.title_y = title_y
		
		self.out_name = out_name
		
		self.xmin = xmin
		self.xmax = xmax
		if self.xmin is None and self.xmax is not None:
			self.xmin = self.xmax-1
		if self.xmin is not None and self.xmax is None:
			self.xmax = self.xmin+1
		
		#-------------------------------------
		#Define empty canvas
		#-------------------------------------
		self.canv = None
		#ToDo: hide if possible
		
		#-------------------------------------
		#Create graphs
		#-------------------------------------
		self.allGraphs = TMultiGraph("allgraphs","allGraphs")
		self.graphs = []
		
	def addFile(self,filename,title=False,axislabels=False,legend=True,xyplot=False,xyerror=False):
		#-------------------------------------
		#Define titles and legend
		#-------------------------------------
		skip_lines = 0
		
		file = open(filename)
		
		if title:
			skip_lines += 1
			tmp = file.readline().strip() #strip() to remove the \n at the end
			if self.title is None:
				self.title = tmp
				print "Title: ", self.title
		
		if axislabels:
			skip_lines += 2
			tmp = file.readline().strip()
			if self.title_x is None:
				self.title_x = tmp
				print "x-Axis: ", self.title_x
			tmp = file.readline().strip()
			if self.title_y is None:
				self.title_y = tmp
				print "y-Axis: ", self.title_y
		
		if legend:
			skip_lines += 1
			labels = file.readline().strip("\r\n").split("\t") #only strip \n, otherwise problem if no label is given for first column
			print "Labels: ", labels
		elif self.show_legend:
			print "Legend missing, deactivating show_legend."
			self.show_legend = False
		
		file.close()
		
		#-------------------------------------
		#Check output filename
		#-------------------------------------
		if self.out_name is None:
			self.out_name = filename
		
		#-------------------------------------
		#Read file
		#-------------------------------------
		data = np.genfromtxt(filename, skip_header=skip_lines, invalid_raise=False, unpack=True, usemask=True)
		#no delimiter usually does the right thing (eg. recognize multiple spaces)
		#invalid_raise=False to handle corrupt (eg. last) lines
		#print data
		
		if(xyerror):
			validValues = np.isfinite(data[0])*np.isfinite(data[1])*np.isfinite(data[2])*np.isfinite(data[3])
			baseX = data[0].compress(validValues)
			baseXerror = data[1].compress(validValues)
			values = data[2].compress(validValues)
			valuesError = data[3].compress(validValues)
			colmin = np.min(values)
			colmax = np.max(values)
			if(np.isnan(self.ymin) or colmin<self.ymin):
				self.ymin=colmin
			if(np.isnan(self.ymax) or colmin>self.ymax):
				self.ymax=colmax
			self.graphs.append(TGraphErrors(len(values),baseX,values,baseXerror,valuesError))
			
			#-------------------------------------
			#Graph properties
			#-------------------------------------
			if legend:
				self.graphs[-1].SetTitle(labels[2])
			self.graphs[-1].SetMarkerColor(len(self.graphs)) #TColor.GetColor("#000000")
			self.graphs[-1].SetMarkerStyle(20)
			self.graphs[-1].SetFillStyle(0)
			self.allGraphs.Add(self.graphs[-1])
			return
		
		if (xyplot):
			baseX = data[0]
			data = data[1:]
			if legend:
				labels = labels[1:]
		else:
			print len(data[0])
			baseX = np.arange(0,len(data[0]),1.0) #must be with 1.0, otherwise an integer array
		
		if legend:
			if self.show_legend and not len(data)==len(labels):
				print "Not the same amount of labels as columns, deactivating show_legend."
				self.show_legend = False
		
		#-------------------------------------
		#Add data to graphs
		#-------------------------------------
		for i,column in enumerate(data):
		#	column = np.nan_to_num(np.array(column)) #Sets all the nan to 0, replaced by the handling below (ommiting values)
			column = np.array(column)
			x = baseX
			x = x.compress(np.isfinite(column)*np.isfinite(baseX)) #only take entries where both x and y are valid
			column = column.compress(np.isfinite(column)*np.isfinite(baseX))
			colmin = np.min(column)
			colmax = np.max(column)
			if(np.isnan(self.ymin) or colmin<self.ymin):
				self.ymin=colmin
			if(np.isnan(self.ymax) or colmin>self.ymax):
				self.ymax=colmax
			self.graphs.append(TGraph(len(column),x,column))
			
			#-------------------------------------
			#Graph properties
			#-------------------------------------
			if legend:
				self.graphs[-1].SetTitle(labels[i])
			self.graphs[-1].SetMarkerColor(len(self.graphs)) #TColor.GetColor("#000000")
			self.graphs[-1].SetMarkerStyle(20)
			self.graphs[-1].SetFillStyle(0)
			self.allGraphs.Add(self.graphs[-1])
	
	def canvas(self):
		#Create a canvas
		if not isinstance(self.canv,TCanvas):
			print "Creating canvas..."
			self.canv = TCanvas("canv","canv",10,10,750,450)
	
	def show(self):
		#-------------------------------------
		#Create canvas if necessary (eg. if window closed)
		#-------------------------------------
		self.canvas()
		
		#-------------------------------------
		#Clear canvas (important if called multiple times)
		#-------------------------------------
		self.canv.Clear()
		
		#-------------------------------------
		#Set title
		#-------------------------------------
		if self.title is None:
			self.title = "Plot title"
		self.allGraphs.SetTitle(self.title)
		
		#-------------------------------------
		#Extend the x-axis range
		#-------------------------------------
		if self.xmin is not None and self.xmax is not None and not isinstance(self.tmp_histo,TH2F):
			self.tmp_histo = TH2F("tmp_histo","tmp_histo",1,self.xmin,self.xmax,1,self.ymin,self.ymax)
			self.tmp_histo.Draw()
		
		#-------------------------------------
		#Draw graphs
		#-------------------------------------
		if not self.show_title:
			gStyle.SetOptTitle(0)
		self.canv.cd()
		if(self.show_line):
			self.allGraphs.Draw("AL")
		else:
			self.allGraphs.Draw("AP")
		
		if self.title_x is None:
			self.title_x = "x-Axis"
		if self.title_y is None:
			self.title_y = "y-Axis"
		
		self.allGraphs.GetXaxis().SetTitle(self.title_x)
		self.allGraphs.GetYaxis().SetTitle(self.title_y)
		
		if(self.yrange):
			print "y-Axis range: ", self.ymin, self.ymax
			allGraphs.GetYaxis().SetRangeUser(self.ymin,self.ymax)

		#-------------------------------------
		#Legend
		#-------------------------------------
		if(self.show_legend):
			self.canvLegend = self.canv.BuildLegend(0.15,0.67,0.53,0.88,"") #Using standard size and position also does not account for the number of entries
			self.canvLegend.SetFillColor(0) #TColor.GetColor("#FFFFFF"))
		
		#-------------------------------------
		#Show the canvas
		#-------------------------------------
		self.canv.Modified()
		self.canv.Update()
		
	def save(self):
		#-------------------------------------
		#Save the plots
		#-------------------------------------
		filenameNoExt,tmp = os.path.splitext(os.path.basename(self.out_name))
		self.canv.Print(filenameNoExt+"_plot.root","root") #Change to "root,RECREATE" if necessary
		self.canv.Print(filenameNoExt+"_plot.pdf","pdf")
		self.canv.Print(filenameNoExt+"_plot.png","png")
	
##*******************************************************************************
#if __name__ == "__main__":
#	filename = r'/home/thomas/docs/Python/Examples/plotFile/72.TXT'
#	xyplot = False
#	legend = False
#	titles = False
#	
#	#Get the command line options
#	try:
#		opts, args = getopt.getopt(sys.argv[1:], "hlxt", ["help,legend,xyplot,titles"])
#	except getopt.error, msg:
#		print msg
#		print "For help use --help"
#		sys.exit(2)
#	
#	#Parse options
#	for option, argument in opts:
#		#Help
#		if option in ("-h", "--help"):
#			print __doc__
#			print plotFile.__doc__
#			sys.exit(0)
#		
#		#X-Y Plot (use first column as x)
#		elif option in ("-x", "--xyplot"):
#			print "xyplot = True (using first column as x values)"
#			xyplot = True
#		
#		#Show legend
#		elif option in ("-l", "--legend"):
#			print "legend = True (show legend)"
#			legend = True
#		
#		#Use first three lines for titles
#		elif option in ("-t", "--titles"):
#			print "titles = True (use titles)"
#			titles = True
#		
#	#Get project name
#	if len(args)>1:
#		print "More than one argument provided."
#		print "For help use --help"
#		sys.exit(2)
#	elif len(args)==1:
#		filename = os.path.abspath(args[0])
#		print "File name set to:", filename,"\n"
#	else:
#		print "Default file name used: ", filename,"\n"
#	
#	plotFile(filename,xyplot,legend,titles)