#!/usr/bin/env python

import ROOT
import numpy as n

var_list=["fqmueEratio:=fq1rtotmu[0][2]/fq1rtotmu[0][1]","fqmuemomratio:=fq1rmom[0][2]/fq1rmom[0][1]","fqmuenllratio:=fq1rnll[0][2]/fq1rnll[0][1]","fqpi0eEratio:=fqpi0totmu[0]/fq1rtotmu[0][1]","fqpi0emomratio:=fqpi0momtot[0]/fq1rmom[0][1]","fqpi0enllratio:=fqpi0nll[0]/fq1rnll[0][1]","fqpi0mass[0]","fqmunll:=fq1rnll[0][2]","fqmumom:=fq1rmom[0][2]","fqmuE:=fq1rtotmu[0][2]","fqenll:=fq1rnll[0][1]","fqemom:=fq1rmom[0][1]","fqeE:=fq1rtotmu[0][1]","costheta:=dir[0][0]*0.669764+dir[0][1]*-0.742179+dir[0][2]*0.024223","dir[0][0]","dir[0][1]","dir[0][2]"]

f_new = ROOT.TFile("PyROOT_newtree.root","RECREATE")

t_in = ROOT.TChain("h1")
t_in.Add("testtree.root")
print 'INPUT: No Events=' + str(t_in.GetEntries())
##t_in=t_in.CopyTree('','',1000)
t_in=t_in.CopyTree('')

##Map of new variables
var_map_in={}
for var in var_list:
	if var.find('=')>=0:
		print 'Composite variable with name ' + var.split(':')[0] + ' and formula ' + var.split('=')[1]
		var_map_in[ROOT.TTreeFormula(var.split(':')[0],var.split('=')[1],t_in)] = n.zeros(1, dtype=n.dtype('f4')) ## f4 is 32 bit float - C-iso standard floats
	else:
		print 'Standard variable with name and formula ' + var
		var_map_in[ROOT.TTreeFormula(var,var,t_in)] = n.zeros(1, dtype=n.dtype('f4')) ## f4 is 32 bit float - C-iso standard floats
var_keys_in=var_map_in.keys()
var_keys_in.sort(key=lambda x: x.GetName(), reverse=False)

print 'Check all TTreeFormulas'
for k in var_keys_in:
	print 'Name: ' + str(k.GetName()) + ', Formula: ' + str(k.GetExpFormula())

n_events = int(t_in.GetEntries())

## A select few histograms t fill on event by event basis
f_new.cd()
histos_event={}
histos_event['fqpi0enllratio']=ROOT.TH1D('he_fqpi0enllratio','he_fqpi0enllratio',300,0.5,1.5)
histos_event['fqmuenllratio']=ROOT.TH1D('he_fqmuenllratio','he_fqmuenllratio',600,-1.5,1.5)
histos_event['fqmunll']=ROOT.TH1D('he_fqmunll','he_fqmunll',400,0,40000)
histos_event['fqmumom']=ROOT.TH1D('he_fqmumom','he_fqmumom',300,0,1500)
histos_event['fqmuE']=ROOT.TH1D('he_fqmuE','he_fqmuE',600,0,12000)
histos_event['fqpi0mass[0]']=ROOT.TH1D('he_fqpi0mass[0]','he_fqpi0mass[0]',300,0.0,1200)
histos_event['costheta']=ROOT.TH1D('he_costheta','he_costheta',200,-1.0,1.0)
histos_event['dir[0][0]']=ROOT.TH1D('he_dir[0][0]','he_dir[0][0]',100,0.0,1.0)
histos_event['dir[0][1]']=ROOT.TH1D('he_dir[0][1]','he_dir[0][1]',100,0.0,1.0)
histos_event['dir[0][2]']=ROOT.TH1D('he_dir[0][2]','he_dir[0][2]',100,0.0,1.0)

##Fill new Tree with values
j=0
for event in t_in:
	for k,v in var_map_in.items():
		v[0]=k.EvalInstance()
		if v[0]==0:
			print k.GetName() + ' with equation ' + str(k.GetExpFormula()) + ' has a value of ' + str(k.EvalInstance())
		if k.GetName() in histos_event:
			histos_event[k.GetName()].Fill(v[0])
	j+=1


#Write the tree to the file

##Also get the histograms for the same variables through the TTree::Draw function
histos_draw={}
histos_draw['fqpi0enllratio']=ROOT.TH1D('hd_fqpi0enllratio','hd_fqpi0enllratio',300,0.5,1.5)
histos_draw['fqmuenllratio']=ROOT.TH1D('hd_fqmuenllratio','hd_fqmuenllratio',600,-1.5,1.5)
histos_draw['fqmunll']=ROOT.TH1D('hd_fqmunll','hd_fqmunll',400,0,40000)
histos_draw['fqmumom']=ROOT.TH1D('hd_fqmumom','hd_fqmumom',300,0,1500)
histos_draw['fqmuE']=ROOT.TH1D('hd_fqmuE','hd_fqmuE',600,0,12000)
histos_draw['fqpi0mass[0]']=ROOT.TH1D('hd_fqpi0mass[0]','hd_fqpi0mass[0]',300,0.0,1200)
histos_draw['costheta']=ROOT.TH1D('hd_costheta','hd_costheta',200,-1.0,1.0)
histos_draw['dir[0][0]']=ROOT.TH1D('hd_dir[0][0]','hd_dir[0][0]',100,0.0,1.0)
histos_draw['dir[0][1]']=ROOT.TH1D('hd_dir[0][1]','hd_dir[0][1]',100,0.0,1.0)
histos_draw['dir[0][2]']=ROOT.TH1D('hd_dir[0][2]','hd_dir[0][2]',100,0.0,1.0)

f_new.cd()
for k,v in var_map_in.items():
	if k.GetName() in histos_draw:
		t_in.Draw(str(k.GetExpFormula())+'>>hd_'+k.GetName())
		histos_draw[k.GetName()].Write()
		histos_event[k.GetName()].Write()


f_new.Close()