from ROOT import TH1D, TH2D, TGraph, TCanvas, TF1, TProfile, gStyle
from numpy import loadtxt
from array import array
import sys, os, re

gStyle.SetOptFit()
gStyle.SetOptStat('nemroui')
gStyle.SetMarkerStyle(20)
c1 = TCanvas( 'c1', 'Example', 200, 10, 700, 500 )
func = TF1("myfit","[0]+(x>[1])*[2]*((x-[1])^2/(2*[3]^3))*exp(1-(x-[1])/[3])")
normp = TProfile("normalized","normalized",500,-100,150)
parhist = TH1D("parameters","parameters",100,5.0,15.0)
par2d = TH2D("par_vs_amplitude","par_vs_amplitude",100,5.0,15.0,100,0.0,2000.0)
rchisq = TH1D("rchisq","rchisq",100,0.0,5.0)
alldata={}
for arg in sys.argv[1:]:
	print arg
	filename = re.sub('\.','_',os.path.basename(arg))

	data = loadtxt(arg)

	data[:,0]*=1e9
	data[:,1]*=1e3
	alldata[filename]=data

	x = array("d",data[:,0])
	y = array("d",data[:,1])
	g = TGraph(len(x),x,y)
	g.GetXaxis().SetLimits(-100,150)
	g.Draw("AL")
	func.SetParameters(-1.0, -10.0, -300.0, 9.0)
	#func.FixParameter(0,0.001)
	g.Fit("myfit","Q","",-200.0,60.0)
	parhist.Fill(func.GetParameter(3))
	#g.Fit("myfit")
	c1.SaveAs(filename+'_fit.png')

	rchisq.Fill(func.GetChisquare()/func.GetNDF())

	par2d.Fill(func.GetParameter(3),-1*func.GetParameter(2))

	for i in range(data.shape[0]):
		normp.Fill((data[i,0]-func.GetParameter(1)),(data[i,1]-func.GetParameter(0))/func.GetParameter(2))

parhist.Draw()
c1.SaveAs('params.png')

rchisq.Draw()
c1.SaveAs('rchisq.png')

par2d.Draw()
c1.SaveAs('par_vs_amplitude.png')

func.SetParameters(0.0, 0.0, 1.0, 9.0)
normp.Fit("myfit","","",-100.0,50.0)
normp.Draw()
c1.SaveAs('norm_prof.png')
#tp = func.GetParameter(3)
tp = 6.95

for filename,data in alldata.items():
	print filename
	x = array("d",data[:,0])
	y = array("d",data[:,1])
	g = TGraph(len(x),x,y)
	g.GetXaxis().SetLimits(-100,150)
	func.SetParameters(-1.0, -10.0, -300.0, tp)
	func.FixParameter(3,tp)
	g.Fit("myfit","Q","",-200.0,60.0)
	#g.Fit("myfit")
	g.Draw("AL")
	c1.SaveAs(filename+'_fit2.png')