#!/usr/bin/env python

import sys
from array import array

import ROOT


if len(sys.argv) < 5:
  print(" Usage: LHE.py abspath_of_ExRootAnalysis input_file sampleType output_file")
  sys.exit(1)

pathExRootAnalysis = sys.argv[1]
print(pathExRootAnalysis)
ROOT.gSystem.Load(pathExRootAnalysis+"/libExRootAnalysis.so")

inputFile = sys.argv[2]

# Create chain of root trees
chain = ROOT.TChain("LHEF")
chain.Add(inputFile)

sampleType = sys.argv[3]
print("sampleType ",sampleType)

# Create object of class ExRootTreeReader
treeReader = ROOT.ExRootTreeReader(chain)
numberOfEntries = treeReader.GetEntries()

# Get pointers to branches used in this analysis
branchParticle = treeReader.UseBranch('Particle')
branchEvent    = treeReader.UseBranch("Event")

output_filename = sys.argv[4]
outputFile = ROOT.TFile(output_filename,"UPDATE")

# Book histograms
histInvMass = ROOT.TH1F("mass_eplus_eminus", "Invariant mass of e+ e-", 50, 0.0, 150.0)
histInvMass.Sumw2(True)

# Book TTree
treeSlimmed = ROOT.TTree("tree","slimmed tree")
mass = array('d', [0])
treeSlimmed.Branch('mass', mass, 'mass/D')
xsec = array('d',[0])
treeSlimmed.Branch('xsec', xsec, 'xsec/D')
N = array('d',[0])
treeSlimmed.Branch('N', N, 'N/D')
eventWeight = array('d',[0])
treeSlimmed.Branch('eventWeight', eventWeight, 'eventWeight/D')

# Loop over all events
for entry in range(0, numberOfEntries):
  # Load selected branches with data from specified event
  treeReader.ReadEntry(entry)

  ## main MC event weight
  weight =  branchEvent[0].Weight

  LVelectron = ROOT.TLorentzVector()
  LVpositron = ROOT.TLorentzVector()  

  found = 0
  for part in branchParticle:
    if part.PID == 11 :
      LVelectron.SetPxPyPzE( part.Px, part.Py, part.Pz, part.E )
      found += 1
    elif part.PID == -11 :
      LVpositron.SetPxPyPzE( part.Px, part.Py, part.Pz, part.E )
      found += 1
  '''
  # If event contains at least 1 jet
  if branchParticle.GetEntries() > 0:
    # Take first jet
    particle = branchParticle.At(0)
  '''  
  if found == 2 :
     invMass = (LVelectron+LVpositron).M() 
     histInvMass.Fill(invMass, weight/numberOfEntries)
     mass[0] = invMass
     eventWeight[0] = weight/numberOfEntries
     xsec[0] = weight
     N[0] = numberOfEntries
     treeSlimmed.Fill()

outputFile.mkdir(sampleType)
outputFile.cd(sampleType)
histInvMass.Write()
treeSlimmed.Write()

outputFile.Close()