#!/usr/bin/env python

import ROOT
import numpy as n
import os

import optparse

def enurec_str_old():
	V  = 27              ## MeV, binding potential
	mn = 939.56536      ## neutron mass
	mp = 938.27203      ## proton mass
	ml = 0.510999     ## electron mass
	pl = 'amome[0]'                      ## momentum of lepton
	el = '(TMath::Sqrt(TMath::Power('+pl+',2) + TMath::Power('+str(ml)+' ,2)))' ## energy of lepton
	
	costh = '(dir[0][0]*0.669764+dir[0][1]*-0.742179+dir[0][2]*0.024223)'
	enu = '(('+str(mn)+'-'+str(V)+')*'+el+' - '+str(ml)+'*'+str(ml)+'/2.0 + '+str(mn)+'*'+str(V)+' - '+str(V)+'*'+str(V)+'/2.0 + ('+str(mp)+'*'+str(mp)+' - '+str(mn)+'*'+str(mn)+')/2.0 ) / ('+str(mn)+' - '+str(V)+' - '+el+' + '+pl+'*'+costh+' )' ##neutrino energy
	return enu.replace(' ','')

def enurec_str():
	V  = 27              ## MeV, binding potential
	mn = 939.565346      ## neutron mass
	mp = 938.272013      ## proton mass
	ml = 0.510998910     ## electron mass
	pl = 'amome[0]'                      ## momentum of lepton
	el = '(TMath::Sqrt(TMath::Power('+pl+',2) + TMath::Power('+str(ml)+' ,2)))' ## energy of lepton
	
	costh = '(dir[0][0]*0.669764+dir[0][1]*-0.742179+dir[0][2]*0.024223)'
	enu = '(('+str(mn)+'-'+str(V)+')*'+el+' - '+str(ml)+'*'+str(ml)+'/2.0 + '+str(mn)+'*'+str(V)+' - '+str(V)+'*'+str(V)+'/2.0 + ('+str(mp)+'*'+str(mp)+' - '+str(mn)+'*'+str(mn)+')/2.0 ) / ('+str(mn)+' - '+str(V)+' - '+el+' + '+pl+'*'+costh+' )' ##neutrino energy
	return enu.replace(' ','')

def fqenurec_str():
	V  = 27              ## MeV, binding potential
	mn = 939.565346      ## neutron mass
	mp = 938.272013      ## proton mass
	ml = 0.510998910     ## electron mass
	pl = 'fq1rmom[0][1]'                      ## momentum of lepton
	el = '(TMath::Sqrt(TMath::Power('+pl+',2) + TMath::Power('+str(ml)+' ,2)))' ## energy of lepton
	
	costh = '(fq1rdir[0][1][0]*0.669764+fq1rdir[0][1][1]*-0.742179+fq1rdir[0][1][2]*0.024223)'
	enu = '(('+str(mn)+'-'+str(V)+')*'+el+' - '+str(ml)+'*'+str(ml)+'/2.0 + '+str(mn)+'*'+str(V)+' - '+str(V)+'*'+str(V)+'/2.0 + ('+str(mp)+'*'+str(mp)+' - '+str(mn)+'*'+str(mn)+')/2.0 ) / ('+str(mn)+' - '+str(V)+' - '+el+' + '+pl+'*'+costh+' )' ##neutrino energy
	return enu.replace(' ','')

##############################################################################
# Parser Options

parser = optparse.OptionParser()
parser.add_option("-o","--outfile",dest="outfile_name",type="string",default="PyROOT_newtree.root",help="A name for the output file\ndefault=PyROOT_newtree.root")
parser.add_option("-i","--input",dest="input",type="string",default="testree.root",help="A comma delimited string of input files.")
parser.add_option("-v","--vars",dest="varlist",type="string",
		  default="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]",
		  help="A comma delimited string of variables to use for MVA - must be the same as in training weights file.")
parser.add_option("-u","--specvars",dest="specvarlist",type="string",
		  default="ipnu[0],costheta:=dir[0][0]*0.669764+dir[0][1]*-0.742179+dir[0][2]*0.024223,amome[0],wall,evis,nhitac,nmue,nring,fmode:=mode*1.0",help="A comma delimited string of spectator variables to carry across")

(options,args) = parser.parse_args()
##############################################################################

input_list=options.input.split(',')
outfile_name=options.outfile_name

var_list=options.varlist.split(",")
var_list.sort()
print 'var_list'
print var_list

specvar_list=options.specvarlist.split(",")
specvar_list.sort()

##############################################################################
###### Create the signal and background chains
##scratch_dir = '/scratch/' ##Some nodes do not seem to have enough space on them for a number of few GB files.
f_new = ROOT.TFile(outfile_name,"RECREATE")
t_in = ROOT.TChain("h1")
for s in input_list:
	t_in.Add(s.replace('\n',''))
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 'We have an = in the var name'
		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:
		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)

specvar_map_in={}
for specvar in specvar_list:
	if specvar.find('=')>=0:
		print 'We have an = in the specvar name'
		specvar_map_in[ROOT.TTreeFormula(specvar.split(':')[0],specvar.split('=')[1],t_in)] = n.zeros(1, dtype=n.dtype('f4')) ## f4 is 32 bit float - C-iso standard floats
	else:
		specvar_map_in[ROOT.TTreeFormula(specvar,specvar,t_in)] = n.zeros(1, dtype=n.dtype('f4')) ## f4 is 32 bit float - C-iso standard floats

Enurec_in = (ROOT.TTreeFormula('Enurec',enurec_str_old(),t_in),n.zeros(1, dtype=n.dtype('f4')))
		
specvar_keys_in=specvar_map_in.keys()
specvar_keys_in.sort(key=lambda x: x.GetName(), reverse=False)

print 'var_keys_in'
for k in var_keys_in:
	print k.GetExpFormula()

n_events = int(t_in.GetEntries())


f_new.cd()
histos={}
histos['fqpi0enllratio']=ROOT.TH1D('h_fqpi0enllratio','h_fqpi0enllratio',300,0.5,1.5)
histos['fqmuenllratio']=ROOT.TH1D('h_fqmuenllratio','h_fqmuenllratio',600,-1.5,1.5)
histos['fqmunll']=ROOT.TH1D('h_fqmunll','h_fqmunll',400,0,40000)
histos['fqmumom']=ROOT.TH1D('h_fqmumom','h_fqmumom',300,0,1500)
histos['fqmumom']=ROOT.TH1D('h_fqmumom','h_fqmumom',300,0,1500)
histos['fqmuE']=ROOT.TH1D('h_fqmuE','h_fqmuE',600,0,12000)
histos['fqpi0mass[0]']=ROOT.TH1D('h_fqpi0mass[0]','h_fqpi0mass[0]',300,0.0,1200)

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


#Write the tree to the file

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

f_new.cd()
##t_in.Write()
f_new.Close()