////////////////////////////////////////////////////////////////////////////////
//
// LepJetsAnalysis
// ---------------
//                                            Rob Harrington <harringt@fnal.gov>
//                                    Philipp Schieferdecker <schiefer@fnal.gov>
////////////////////////////////////////////////////////////////////////////////


#include "top_tree_lepjets/LepJetsAnalysis.hpp"

#include "top_tree_tools/DuplicateEventFilter.hpp"
#include "top_tree_tools/GoodRunSelector.hpp"
#include "top_tree_tools/GoodLBNSelector.hpp"
#include "top_tree_tools/SelectionTool.hpp"
#include "top_tree_tools/TopoTool.hpp"
#include "top_tree_tools/JesTool.hpp"
#include "top_tree_tools/JetPartonTool.hpp"
#include "top_tree_tools/TriggerTool.hpp"
#include "top_tree_tools/BTagProbTool.hpp"

#include "top_tree/TheEvent.hpp"
#include "top_tree/TheMissingEt.hpp"
#include "top_tree/TheProperty.hpp"

#include "mc_tree/MCTree.hpp"
#include "mc_tree/MCParticle.hpp"


#include "top_dq/analysis_dataquality.hpp"
#include "kinem_util/AnglesUtil.hpp"


//hitfit includes
#include "hitfit/Defaults_Text.hpp"
#include "hitfit/fourvec.hpp"
#include "hitfit/Lepjets_Event.hpp"
#include "hitfit/Lepjets_Event_Jet.hpp"
#include "hitfit/Lepjets_Event_Lep.hpp"
#include "hitfit/Fit_Results.hpp"
#include "hitfit/Fit_Result.hpp"
#include "hitfit/Top_Fit.hpp"
#include "hitfit/run1_ascii_io.hpp"



#include <TFile.h>
#include <TTree.h>
#include <TH1F.h>
#include <TH2F.h>
#include <TF2.h>
#include <TLorentzVector.h>
#include <TRandom.h>


#include <iostream>
#include <iomanip>
#include <cmath>


using namespace std;


////////////////////////////////////////////////////////////////////////////////
// declare local functions
////////////////////////////////////////////////////////////////////////////////

double TF2Lookup(TF2 *func, Float_t xval, Float_t yval);


////////////////////////////////////////////////////////////////////////////////
// construction/destruction
////////////////////////////////////////////////////////////////////////////////

//______________________________________________________________________________
LepJetsAnalysis::LepJetsAnalysis(const string& name,
				 const string& parameterFile)
    : TopTreeAnalysis(name, parameterFile) 
    , _dq_ejets(0)
    , _dq_mujets(0) 
    , _triggerTool_ejets(0)
    , _triggerTool_mujets(0)
    , _topoTool(0)
    , _jesTool(0)
    , _jetPartonTool(0)
    , _btagProbTool(0)
    , _muonID(0)
    , _nProcessed(0)
    , _nPassEventList(0)
    , _nPassPv(0)
    
{

    //
    // initialize jet multiplicity counters
    //
    for (int i=0;i<2;i++) {
	for (int j=0;j<5;j++) {
	    _nJetMult_loose[i][j] = 0;
	    _nJetMult_tight[i][j] = 0;
	}

	_nPassLepSel[i]  = 0;
	_nPassMEt[i]     = 0;
	_nPassTrigger[i] = 0;
	_nPassSLTVeto[i] = 0;
	_nPassJPMatch[i] = 0;
	_nPassTight[i]   = 0;
    }
	
    
    //
    // declare parameters
    //
    declareParameter("TreeFilePath",      _treeFilePath      =          "./tree");

    declareParameter("AssignRunEvt",      _assignRunEvt      =             false);
    declareParameter("RunNumber",         _runNumber         =               666);
    
    declareParameter("DeltaPvZMax",       _deltaPvZMax       =               5.0);
    declareParameter("LowMEtSelection",   _lowMEtSelection   =             false);
    declareParameter("QCDSelection",      _qcdSelection      =                 0);
    
    declareParameter("PvNTrksMin",        _pvNTrksMin        =                 3);
    declareParameter("PvZMax",            _pvZMax            =              60.0);
    
    declareParameter("ScndMuonVeto",      _scndMuonVeto      =              true);
    declareParameter("SoftMuonVeto",      _softMuonVeto      =             false);
    declareParameter("WriteMcMuon",       _writeMcMuon       =             false);
    
    declareParameter("ElecLhoodMin",      _elecLhoodMin      =              0.85);
    
    declareParameter("JetPtMin",          _jetPtMin          =              15.0);
    declareParameter("JetEtaMax",         _jetEtaMax         =               2.5);
    declareParameter("VetoBadJetEvents",  _vetoBadJetEvents  =             false);
    declareParameter("JesVariation",      _jesVariation      =                 0);

    declareParameter("METMin",            _metMin            =               20.);

    declareParameter("TriggerVariation",  _triggerVariation  =             false);
    declareParameter("TriggerVerMin",     _triggerVerMin     =                 8);
    declareParameter("TriggerVerMax",     _triggerVerMax     =                12);
    declareParameter("v13",               _v13               =             false);
    
    declareParameter("EventListFileName", _eventListFileName =                "");
    declareParameter("RejectEventList",   _rejectEventList   =             false);
    
    declareParameter("UseOwnMCParam",     _useOwnMCParam     =              true);
    declareParameter("Debug",             _debug             =             false);
    declareParameter("Sample",            _sample            =                "");

    declareParameter("SVXVertexType",     _SVXVertexType     =           "TIGHT");

    declareParameter("DuplEvtParamFile",  _duplEvtParamFile  =                "");
    declareParameter("GoodRunParamFile",  _goodRunParamFile  = "Disabled.params");
    declareParameter("GoodLBNParamFile",  _goodLBNParamFile  = "Disabled.params");

    

    //
    // compile parameter file and update parameter values
    //
    if (!setParameters()) exit(11);
    
    
    //
    // instantiate tools
    //
    
    // ejets/mujets selection tool
    _selectionTool = new SelectionTool();

    if (_qcdSelection==2) {
	_selectionTool->requireEMID(false);
	_selectionTool->requireEMTrkMatch(false);
	_selectionTool->requireEMDeltaZPv(false);
    }

    // topological variables
    _topoTool =  new TopoTool();
    _topoTool -> setDebug(_debug);

    // jes variations and jet-parton matching
    _jesTool       = new JesTool();
    _jetPartonTool = new JetPartonTool(0.5);
    
    
    // b-tagging tool setup 
    _btagProbTool  = new BTagProbTool();
    _btagProbTool -> setUseOwnParam(_useOwnMCParam);
    _btagProbTool -> setDebug(_debug); 
    
    if (_useOwnMCParam && _SVXVertexType != "TIGHT") {
	cout 
	    << "can only do TIGHT vertex selection with useOwnMCParam = true" 
	    << endl;
	exit(76);
    }

    // muon ID: data-to-MC efficiency scale
    string muonIDpath = getenv("SRT_LOCAL");
    muonIDpath += "/top_tree_tools/parametrisations/Muon_ID_SF_2D.root";
    TFile* f = new TFile(muonIDpath.c_str(), "READ"); 
    if (!f->IsOpen()) {
	cout << "ERROR: Could not open " << muonIDpath.c_str() << endl;
	exit(15);
    }
    _muonID = (TF2*)(f->Get("result")->Clone());

    // these aren't used yet, but can be for systematics
    //    _muonID_p = (TF2*)(f->Get("resultp")->Clone());
    //    _muonID_n = (TF2*)(f->Get("resultn")->Clone());
    
    f->Close();
    
    
    // data quality
    string dqpath = getenv("SRT_LOCAL");
    if (dqpath.empty()) exit(22);
    dqpath += "/top_dq_data";
    _dq_ejets  = new analysis_dataquality(dqpath, "EMQCD");
    _dq_mujets = new analysis_dataquality(dqpath, "MUQCD");
        
    // trigger efficiency
    _triggerTool_ejets  = new TriggerTool("ejets" ,_v13,false);
    _triggerTool_mujets = new TriggerTool("mujets",_v13,false);
    
    
    //
    // initialize random generator
    //
    gRandom->SetSeed(0);
    
    

    //
    // add event selectors
    //
    EventSelector* selector(0);
    selector=new DuplicateEventFilter(this,"DuplicateEventFilter",_duplEvtParamFile);
    addEventSelector(selector);
    selector=new GoodRunSelector(this,     "GoodRunSelector",     _goodRunParamFile);
    addEventSelector(selector);
    selector=new GoodLBNSelector(this,     "GoodLBNSelector",     _goodLBNParamFile);
    addEventSelector(selector);
    
    
    
    //
    // declare histograms
    //
    setSumw2();
    
    _hEventWeight_j       = declareTH1F("EventWeight_j",    "EventWeight (njt=1)",
					1000,  0.0,   1.0);
    _hEventWeight_jj      = declareTH1F("EventWeight_jj",   "EventWeight (njt=2)",
					1000,  0.0,   1.0);
    _hEventWeight_jjj     = declareTH1F("EventWeight_jjj",  "EventWeight (njt=3)",
					1000,  0.0,   1.0);
    _hEventWeight_jjjj    = declareTH1F("EventWeight_jjjj", "EventWeight (njt=4)",
					1000,  0.0,   1.0);
    
    _hJetMult             = declareTH1F("JetMult",          "JetMult",
					5, 0.5, 5.5);
    _hJetMult_loose       = declareTH1F("JetMult_loose",    "JetMult (loose)",
					5, 0.5, 5.5);
    _hJetMult_tight       = declareTH1F("JetMult_tight",    "JetMult (tight)",
					5, 0.5, 5.5);
    _hBadJetMult          = declareTH1F("BadJetMult",       "BadJetMult",
					5, 0.5, 5.5);
    
    _hMcChannel           = declareTH1F("McChannel",        "McChannel",
					5, 0.5, 5.5);
    _hNMatchedJets        = declareTH1F("NMatchedJets",     "NMatchedJets",
					5, -0.5, 4.5);
    
    //hitfit setup
  // read in defaults file
  char* hitfit_args[] = {"input.dat"};
  hitfit::Defaults_Text hitfit_defs( "./hitfit/run/defaults", 1, hitfit_args );
  
  // get the constrain on the leptonic W
  float lepw_mass = hitfit_defs.get_float( "lepw_mass" ); 

  // get the constrain on the hadronic W
  float hadw_mass = hitfit_defs.get_float( "hadw_mass" ); 

  // get the constrain on the top mass
  float top_mass = hitfit_defs.get_float( "top_mass" );   

  hitfit::Run1_Ascii_IO_Args hitfit_ioargs( hitfit_defs );
  hitfit::Top_Fit_Args hitfit_args_2( hitfit_defs );

  // create the HitFit object
  hitfit::Top_Fit hitfit_fitter = hitfit::Top_Fit( hitfit_args_2,lepw_mass, hadw_mass, top_mass );


    
    //
    // initialize output tree
    //
    initOutputTree(_treeFilePath + "/" + _outputFileName);
    
}


//______________________________________________________________________________
LepJetsAnalysis::~LepJetsAnalysis()
{
    
}


////////////////////////////////////////////////////////////////////////////////
// implementation of member functions
////////////////////////////////////////////////////////////////////////////////

//______________________________________________________________________________
void LepJetsAnalysis::analyzeEvent()
{
    _nProcessed++;
    
    _weight    =1.0;
    _muidweight=1.0;
    _trgweight =1.0;

    _tagprobTight=_tagprobMedium=_tagprobLoose = 0.0;
    _dbltagprobTight=_dbltagprobMedium=_dbltagprobLoose = 0.0;

    //
    // if an eventlist has been specified, evaluate
    //
    if (!_eventListFileName.empty() && _eventList.size()>0) {
	CollisionID collid(event()->GetRunNumber(), event()->GetEvtNumber());
	bool evtInList = 
	    binary_search(_eventList.begin(), _eventList.end(), collid);
	
	if (evtInList  &&  _rejectEventList) return;
	if (!evtInList && !_rejectEventList) return;
    }
    _nPassEventList++;
    
  // create the event holder object
    int evtnr = event()->GetEvtNumber();
    int runnr = event()->GetRunNumber();

    //
    // primary vertex selection
    //
    if (event()->GetNewPrimaryVertexTrks()        < _pvNTrksMin ||
	std::abs(event()->GetNewPrimaryVertexZ()) > _pvZMax ||
	std::abs(event()->GetPrimaryVertexZ() -
		 event()->GetNewPrimaryVertexZ()) > _deltaPvZMax) return;
    _nPassPv++;
    
    
    //
    // determine if this event is MC
    //
    bool isMC = (0 != mcInfo()->nParticles());

    // 
    // calorimeter-based event quality cuts
    //
    if (!isMC) {
	int twr_l1et_lt_2_em = 
	    event()->GetTottwr() - event()->GetTwr_l1et_gt_2_em();  
      
	bool passCN = event()->GetCoherentNoise()==0;
	bool passL1EvtConf = 
	    (double)event()->GetTwr_l1et_lt_n1_em()/event()->GetTottwr() > 0.01 || 
	    (double)event()->GetTwr_diff_gt_1_em() / twr_l1et_lt_2_em < 0.0030;
      
	if ( !passCN && !passL1EvtConf)
	    return;
    }
  
  
    //
    // dig out the kind of leptonic W decay at tree-level
    //    
    const MCParticle* mcLepton(0);
    if (isMC) {
	bool isWe(false);
	bool isWmu(false);
	bool isWtau(false);
	bool isWtaue(false);
	bool isWtaumu(false);
      
	MCPartIter_t it;
	for (it=mcInfo()->particles_begin();it!=mcInfo()->particles_end();++it) {
	    if ((*it)->mother() != 0 && (*it)->mother()->absPdgId() == 24) {
		
		int pdg = (*it)->absPdgId();
		
		// W -> e
		if      (pdg == 11) {
		    isWe = true;
		    mcLepton = (*it);
		    break;
		}
		// W -> mu
		else if (pdg == 13) {
		    isWmu = true;
		    mcLepton = (*it);
		    break;
		}
		// W -> tau
		else if (pdg == 15) {
		    isWtau  = true;
		    if (isWtau) {
			set<int> eids;
			eids.insert(11);
			eids.insert(-11);
			if ((*it)->hasDaughter(eids)) isWtaue = true;
			set<int> muids;
			muids.insert(13);
			muids.insert(-13);
			if ((*it)->hasDaughter(muids)) isWtaumu = true;
		    }
		    break;
		}
	    }
	}
	if (!isWe && !isWmu && !isWtau) {
	    cout << "No W->e/mu/tau decay found!" << endl;
	    //return;
	}
	if (isWe)     _hMcChannel->Fill(1);
	if (isWmu)    _hMcChannel->Fill(2);
	if (isWtau)   _hMcChannel->Fill(3);
	if (isWtaue)  _hMcChannel->Fill(4);
	if (isWtaumu) _hMcChannel->Fill(5);

	if (0 != mcLepton) {
	    _lepmce   = mcLepton->E();
	    _lepmcpt  = mcLepton->Pt();
	    _lepmceta = mcLepton->Eta();
	    _lepmcphi = mcLepton->Phi();
	}
	else {
	    _lepmce   = 0.0;
	    _lepmcpt  = 0.0;
	    _lepmceta = 0.0;
	    _lepmcphi = 0.0;
	}
    }
    
    
    //
    // retrieve all muons and jets
    //
    TClonesArray* elecs = objects()->GetElectrons();
    TClonesArray* muons = objects()->GetMuons();
    TClonesArray* jets  = objects()->GetJets();
    
    //
    // correct muon pT and missing Et for CFT-only tracks
    //
    // 01/11/2005 PS: top_analyze (Ipanema) does this now already
    //
    //setNewMuonPt();

    
    //
    // JES variation
    //
    if (_jesVariation != 0) _jesTool->applyJesVariation(this, _jesVariation);
    
    
    
    //
    // l o o s e   p r e s e l e c t i o n
    //
    TheElectronClass* isoEM   = _selectionTool->ejets_loose(this);
    TheMuonClass*     isoMuon = _selectionTool->mujets_loose(this);
    
    if (0 != isoMuon && 0 != isoEM) {
	cout << "ERROR: How the hell can this event be both, ejets AND mujets??"
	     << endl;
	return;
    }
    
    // set MEt tree variables
    _met    = missingEt()->GetMJMEt();
    _metx   = missingEt()->GetMJMEtx();
    _mety   = missingEt()->GetMJMEty();
    _metphi = missingEt()->GetMJMEtphi();
    _metedc = 1.0;
    

    // pointer to data quality
    analysis_dataquality* dq(0);
    

    //
    // e+jets
    //
    if (0 != isoEM) {

	_nPassLepSel[0]++;

	
	// missing Et selection
	if (_lowMEtSelection) {
	    if (_met > 10.0) return;
	}
	else {
	    if (_met < _metMin) return;
	    
	    // dphi(em, met) cut
	    double em_phi  = isoEM->Phi();
	    double dphi    = kinem::delta_phi(em_phi, _metphi);
	    //if (dphi < (1.7 - (1.7*_met/80.))) return;
	    if (dphi < (2.19911 - (2.19911*_met/48.86911))) return;
	}
	_nPassMEt[0]++;
	

	// handle MC
	if (isMC) {
	    
	    // calculate trigger weight
	    _trgweight=_triggerTool_ejets->eventWeight("ejets",elecs,muons,jets,_met);
	    _weight  *=_trgweight;
	    
	    // systematic trigger variation
	    if (_triggerVariation) {
		vector<float> wvec=_triggerTool_ejets->eventWeightSigma("ejets",elecs,muons,jets,_met);
		for (unsigned int i=0;i<18;i++) _trgweights[i]=wvec[i];
	    }
		
	    
	} // isMC
	else {
	    
	    // trigger
	    dq = _dq_ejets;
	}
        
	// set tree variables
	_leptype      = 1;
	_lepiso       = 1;
	_lepe         = isoEM->E();
	_leppt        = isoEM->pT();
	_leppx        = isoEM->pT();
	_leppy        = isoEM->pT();
	_leppz        = isoEM->pT();
	_lepeta       = isoEM->Eta();
	_lepphi       = isoEM->Phi();
	_lepdeta      = isoEM->CalDetectoreta();
	_lepdphimet   = kinem::delta_phi(_lepphi, _metphi);
	_lepnsmt      = isoEM->GetInfiducial(); // nsmt usefull for EMs?
		
	// calculate topological variables
	if (_leppt > 0.)
	    _topoTool->topology(this, isoEM);
	else return;
    }
    
    //
    // mu+jets
    //
    else if (0 != isoMuon) {
	
	_nPassLepSel[1]++;

	// MEt Selection
	if (_lowMEtSelection) {
	    if (_met >= 15.0 || missingEt()->JesMEt() >= 25.0) return;
	}
	else {
	    int          leadingJetIdx(0);
	    TheJetClass* leadingJet = (TheJetClass*)jets->At(leadingJetIdx);
	    while (0 != leadingJet && (leadingJet->pT() < _jetPtMin ||
				       std::abs(leadingJet->Eta()) > _jetEtaMax))
		leadingJet = (TheJetClass*)jets->At(++leadingJetIdx);
	    
	    if (0 == leadingJet) return;
	    
	    double pi          = TMath::Pi();
	    double met         = missingEt()->GetMJMEt();
	    
	    double phi_mu      = isoMuon->Phi();
	    double phi_met     = missingEt()->GetMJMEtphi();
	    double dphi_metmu  = kinem::delta_phi(phi_mu, phi_met);
	    
	    double phi_jet     = leadingJet->Phi();
	    double dphi_metjet = kinem::delta_phi(phi_jet, phi_met);
	

	    double p    = pi/10.;
	    double cut1 = _metMin;
	    double cut2 =  p;
	    double cut3 = 50.0;
	    double cut4 = 8.*p;
	    double cut5 = 30.0;

	    bool passMEtCut(false);
	    if (met         > cut1 &&
		dphi_metmu  > cut2 - met*cut2/cut3 &&
		dphi_metmu  < cut4 + met*(pi-cut4)/cut5) passMEtCut = true;
	    // removed dphi_metjet cut to match production presel
	    //dphi_metjet < cut6 + met*(pi-cut6)/cut7) passMEtCut = true;
	    
	    if (!passMEtCut) return;
	}
	_nPassMEt[1]++;
	

	// calculate tree variable for detector eta
	_lepdeta = _triggerTool_mujets->muonDetEta(isoMuon->Eta(),
						   isoMuon->Phi(),
						   isoMuon->Z0());
	
	if (isMC) {
	    
	    // calculate trigger weight
	    _trgweight=_triggerTool_mujets->eventWeight("mujets",elecs,muons,jets,_met);
	    _weight  *= _trgweight;
	    
	    // systematic trigger variation
	    if (_triggerVariation) {
		vector<float> wvec=_triggerTool_mujets->eventWeightSigma("mujets",elecs,muons,jets,_met);
		for (unsigned int i=0;i<18;i++) _trgweights[i]=wvec[i];
	    }
	    
	    // muonID efficiency
	    _muidweight=TF2Lookup(_muonID,_leppt,std::abs(_lepdeta));
	    _weight   *=_muidweight;
	    
	} // isMC
	else {
	    // trigger
	    dq = _dq_mujets;
	}
        
	// access track matched to muon
	TClonesArray* tracks = objects()->GetTracks();
	TheTrackClass* muTrk = (TheTrackClass*)tracks->At(isoMuon->TrackIndex());
	if (0 == muTrk) {
	    cout << "ERROR: isolated muon has invalid track index!"<<endl;
	    return;
	}
	
	// set tree variables
	_leptype    = 2;
	_lepiso     = 1;
	_lepe       = isoMuon->E();
	_leppt      = isoMuon->pT();
	_leppx      = isoMuon->pT();
	_leppy      = isoMuon->pT();
	_leppz      = isoMuon->pT();
	_lepeta     = isoMuon->Eta();
	_lepphi     = isoMuon->Phi();
	// _lepdeta already set above!
	_lepdphimet = kinem::delta_phi(_lepphi, _metphi);
	_lepnsmt    = static_cast<int>(muTrk->GetSMTHits());
	
	
	// calculate topological variables
	if (_leppt > 0.)
	    _topoTool->topology(this, isoMuon);
	else return;
	
    }
    else return;
    
    
    //
    // trigger
    //
    if (0 != dq) {
	int triggver = dq->getTriggerListVersion(event()->GetRunNumber());
	if (triggver < _triggerVerMin || triggver > _triggerVerMax) return;
	
	bool passTrigger(false);
	int lbn = event()->GetLumblk();
	const vector<string>& requiredTriggers = dq->getValidTriggers(lbn);
	vector<string>::const_iterator it;
	for (it = requiredTriggers.begin(); it != requiredTriggers.end(); ++it) {
	    if (triggerFired(*it)) passTrigger = true;
	}
	if (!passTrigger) return;
    }
    _nPassTrigger[_leptype-1]++;
    
    
    
    //
    // select jets
    //
    int nJets = jets->GetEntries();
    for (int i=0;i<nJets;i++) {
	TheJetClass* jet = (TheJetClass*)jets->At(i);
	if (jet->pT()            < _jetPtMin || 
	    std::abs(jet->Eta()) > _jetEtaMax) jets->Remove(jet);
    }
    jets->Compress();
    int nSelJets = jets->GetEntries();
    int jetMult = std::min(nSelJets, 5);
    _hJetMult->Fill(jetMult);
    

    //
    // require at least one jet
    //
    if (jetMult < 1) return;
    
    
    //
    // soft muon veto
    //
    if (_softMuonVeto) {
	for (int i=0;i<nSelJets;i++) {
	    TheJetClass* jet = (TheJetClass*)jets->At(i);
	    if (jet->GetMuoTag() != 0) return;
	}
    }
    _nPassSLTVeto[_leptype-1]++;
    
    
    //
    // reject events containing bad jets
    //
    TClonesArray* badJets = objects()->GetBadJets();
    int nBadJets = badJets->GetEntries();
    _hBadJetMult->Fill(nBadJets);
    if(_vetoBadJetEvents && nBadJets > 0) return;
    
    
    
    
    //
    // match jets to partons anc calculate btagging probability
    //
    if (isMC) {
      
	// setup btagging tool for this event
	_btagProbTool -> setObjects( (TClonesArray*)objects() );
	if (_debug) cout << "set btag objects" << endl;
	
	MCPartColl_t partons;
	
	// ttbar
	const MCParticle* t;	const MCParticle* tbar;
	const MCParticle* b;	const MCParticle* bbar;
	const MCParticle* lep;	const MCParticle* nu;
	const MCParticle* q1;	const MCParticle* q2;
	
	if (_sample == "" || _sample == "ttbar") {
	    if (mcInfo()->findTTbarLepJetsDecay(t, tbar, b, bbar, lep, nu, q1, q2)) {
		partons.push_back(b);
		partons.push_back(bbar);
		partons.push_back(q1);
		partons.push_back(q2);
		
		_hNMatchedJets->Fill(_jetPartonTool->match(jets, partons));
	    }
	}
	
	// W+jets
	if (_sample != "" && _sample != "ttbar") {
	    bool partonMatch = false;
	    if (_debug)
		cout << "EventTopology(" << _sample 
		     << ") = " << _btagProbTool->EventTopology(_sample)
		     << endl;
	    if (_btagProbTool->EventTopology(_sample) == _sample) partonMatch = true;
	    
	    if (!partonMatch) return;
	    
	    // how do you want to fill _hNMatchedJets???
	    // _hNMatchedJets->Fill(_jetPartonTool->match(jets, partons));
	}
	
	//
	// get btagging weight
	//
	if (_SVXVertexType == "ALL" || _SVXVertexType == "TIGHT") {
	    _btagProbTool -> setVertexType(TString("TIGHT"));
	    _tagprobTight    = _btagProbTool -> EvtTagProb(true);
	    _dbltagprobTight = _btagProbTool -> EvtTagProbDT(true);
	}
	
	if (_SVXVertexType == "ALL" || _SVXVertexType == "MEDIUM") {
	    _btagProbTool -> setVertexType(TString("MEDIUM"));
	    _tagprobMedium    = _btagProbTool -> EvtTagProb(true);
	    _dbltagprobMedium = _btagProbTool -> EvtTagProbDT(true);
	}
	
	if (_SVXVertexType == "ALL" || _SVXVertexType == "LOOSE") {
	    _btagProbTool -> setVertexType(TString("LOOSE"));
	    _tagprobLoose    = _btagProbTool -> EvtTagProb(true);
	    _dbltagprobLoose = _btagProbTool -> EvtTagProbDT(true);
	}

    }
    
    int numberTags = 0;
    
    
    //
    // calculate variables which are written to output tree
    //
    if (_assignRunEvt) {
	_run=_runNumber;
	_evt=_nProcessed;
    }
    else {
	_run=event()->GetRunNumber();
	_evt=event()->GetEvtNumber();
    }
    
    _type  =0;
    
    _wmt   =_topoTool->wmt();
    _wet   =_topoTool->wet();
    _weta  =_topoTool->weta();
    
    _ht20  =_topoTool->ht20();
    _ht25  =_topoTool->ht25();
    
    _apla  =_topoTool->aplanarity();
    _spher =_topoTool->spheresity();
    _ktminp=_topoTool->ktminp();
    _ht2p  =_topoTool->ht2p();
    
    _disc = _topoTool->disc();

    _lhood           = -1.0;
    _nnd             = -1.0;
    
    _rof         = event()->GetRingOfFire();
    _cnoise      = event()->GetCoherentNoise();
    _ecrate      = event()->GetEmptyCrate();

    
    _njt             = nSelJets;
    
    for (int i=0;i<_njt;i++) {
	TheJetClass* jet=(TheJetClass*)jets->At(i);
	_jte[i]       =jet->E();
	_jtpt[i]      =jet->pT();
	_jteta[i]     =jet->Eta();
	_jtphi[i]     =jet->Phi();
	_jtdeta[i]    =jet->detectorEta();
	_jtdphimet[i] =kinem::delta_phi(jet->Phi(), _metphi);
	_jtsvxtag[i]  =isSvxTagged(jet, 7.0);
	_jtmutag[i]   =jet->GetMuoTag();
	_jtjes[i]     =jet->Correction();
	_jtjesdata[i] =(_jtmutag[i])?
	    jet->Correction_data_hq():jet->Correction_data_lq();
	_jtjesmc[i]   =(_jtmutag[i])?
	    jet->Correction_mc_hq():jet->Correction_mc_lq();

	_jttaggable[i]=_btagProbTool->getTaggability(_jtpt[i], 
						     _jteta[i],
						     jet->Flavor(), 
						     0!=isoMuon);
	// per jet tagging efficiencies
	if (isMC) {
	    bool mujets = (0 != isoMuon) ? 1 : 0;
	    _jttagprobTight[i] =  
		_btagProbTool->getJetTagProb(jet,"TIGHT", mujets);
	    _jttagprobMedium[i] = 
		_btagProbTool->getJetTagProb(jet,"MEDIUM", mujets);
	    _jttagprobLoose[i] =  
		_btagProbTool->getJetTagProb(jet,"LOOSE", mujets);
	}
	else _jttagprobTight[i]=_jttagprobMedium[i]=_jttagprobLoose[i]=0.;
	
	
	_jtmcid[i] =  0;
	_jtmce[i]  =0.0;
	_jtmcpt[i] =0.0;
	_jtmceta[i]=0.0;
	_jtmcphi[i]=0.0;
	
	MCPartColl_t partons;
	if (1 == _jetPartonTool->matchedPartons(jet, partons)) {
	    const MCParticle* mcp = partons[0];
	    _jtmcid[i] =mcp->pdgId();
	    _jtmce[i]  =mcp->E();
	    _jtmcpt[i] =mcp->Pt();
	    _jtmceta[i]=mcp->Eta();
	    _jtmcphi[i]=mcp->Phi();
	}

	if ( isSvxTagged(jet, 7.) ) ++numberTags;
    }

    // for data, set _tagprob and _dbltagprob to 0 or 1
    if (!isMC) {
	if ( numberTags > 0 ) 
	    _tagprobTight=_tagprobMedium=_tagprobLoose = 1.;
	if ( numberTags > 1 ) 
	    _dbltagprobTight=_dbltagprobMedium=_dbltagprobLoose = 1.;
    } 
    
    hitfit::Lepjets_Event hitfit_event( evtnr, runnr ); // arguments are event & run number

    //isMC()?
    if(isMC())  hitfit_event.setMC(1);
    else  hitfit_event.setMC(0);
  
    // Add lepton:
    const hitfit::Fourvec lep_fourvec(_leppx,_leppy,_leppz,_lepe);// 4vec, based on CHLEP
    hitfit_event.add_lep( hitfit::Lepjets_Event_Lep( lep_fourvec, ltype, hitfit_ioargs.muo_res() ) ); //(4vector,type???,muon_resolution???used???)

  // Add missing Et:
    const hitfit::Fourvec met_fourvec(_metx,_mety,_met,_met);
    hitfit_event.met() = met_fourvec;
  
    // Add vertex:
    if(getDebug()) printf("After adding lepton, MET & vertex\n");


    // add the jets to hitfit
    for (int i=0;i<_njt;i++) {
	TheJetClass* jet=(TheJetClass*)jets->At(i);
	const hitfit::Fourvec jet_fourvec( jet->GetPx(), jet->GetPy(), jet->GetPz(), jet->E() );
	hitfit_event.add_jet( hitfit::Lepjets_Event_Jet( jet_fourvec, jtype, hitfit_ioargs.jet_res(),jet->GetSvxTags(), 0 ) ); //0=jetmuotag??????
    }
    //set kt_res
    hitfit_event.kt_res() = hitfit_ioargs.kt_res();
    
    if(getDebug()) printf("After setting kt_res()\n");

    //// Run hitfit: ////

    hitfit::Fit_Results hitfit_results = hitfit_fitter.fit( hitfit_event );
    hitfit::Fit_Result_Vec hitfit_result_vec_all = hitfit_results[hitfit::all_list];

    ///Get Resutls from hifit
    bool goodMassFit = true;
    //_nperm = properties()->GetNperm();
    int _nperm = hitfit_result_vec_all.size(); 
    if(_nperm==0){
      if (_debug) std::cout << "_nperm=" << _nperm << std::endl;
      goodMassFit = false;
    }
    if (goodMassFit) {
      // TopFitArray variables
      _nPassHitFit++;
      for (int i=0; i<_nperm; i++) {
        _perm[i]  = 
        _mt[i]    =  hitfit_result_vec_all[i].mt();
        _chisq[i] =  hitfit_result_vec_all[i].chisq();
        _sigmt[i] =  hitfit_result_vec_all[i].sigmt();
	hitfit::Lepjets_Event lepjet_event = hitfit_result_vec_all[i].ev();
	//Loop over the jets (to find the jet permutations...)
	int hf_njets = lepjet_event.njets();
	for(int j=0; j!=hf_njets; ++j) {
	  
	  //is this the bjet on the lepton side?
	  if(lepjet_event.jet(j).type() == hitfit::lepb_label) {  
	    //do something here
	    
	  }
	}

      }
    }
    /**
    // get hitfit info
    bool goodMassFit = true;
    _nperm = properties()->GetNperm();
    if(_nperm==0){
      if (_debug) std::cout << "_nperm=" << _nperm << std::endl;
      goodMassFit = false;
    }
    if (goodMassFit) {
     // TopFitArray variables
      _nPassHitFit++;
      TClonesArray* TopFitArray = properties()->GetFits();
      for (int i=0; i<_nperm; i++) {
        TheTopFitClass* topfit = (TheTopFitClass*)TopFitArray->At(i);
        _perm[i]  = topfit->Getperm();
        _mt[i]    = topfit->GetMt();
        _chisq[i] = topfit->GetChisq();
        _sigmt[i] = topfit->GetSigmt();
      }
    }
    **/

    //
    // jet multiplicity counter
    //
    _nJetMult_loose[_leptype-1][jetMult-1]++;
    _hJetMult_loose->Fill(jetMult, _weight);
    

        
    //
    // t i g h t   s e l e c t i o n
    //
    if (0 != isoEM&&_qcdSelection!=2) {
	if (_qcdSelection==1) {
	    if (isoEM->Lhood() > _elecLhoodMin) {
		_outputTree->Fill();
		return;
	    }
	}
	else if (isoEM->Lhood() < _elecLhoodMin) {
	    _outputTree->Fill();
	    return;
	}
    }
    else if (0 != isoMuon) {
	if (_qcdSelection==1) {
	    if (isoMuon->Isolated()) {
		_outputTree->Fill();
		return;
	    }
	}
	else if (!isoMuon->Isolated()) {
	    _outputTree->Fill();
	    return;
	}
    }

    _nPassTight[_leptype-1]++;
    
    
    _lepiso = 2;
    _outputTree->Fill();
    
    _nJetMult_tight[_leptype-1][jetMult-1]++;
    _hJetMult_tight->Fill(jetMult, _weight);
    
    
    
    return;
}


//______________________________________________________________________________
void LepJetsAnalysis::endOfJob()
{
    //
    // write output tree
    //
    cout << "Writing "     << _outputTree->GetEntries() 
	 << " entries to " << _outputTreeFile->GetName() << ".\n" 
	 << endl;
    _outputTreeFile->cd();    
    _outputTree->Write();
    _outputTreeFile->Close();
    
    
    //
    // print summary
    //
    cout << "%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%\n"
	 << "LepJetsAnalysis Summary\n" << "----------------------\n" 
	 << endl;
    
    // calculate inclusive numbers
    int nLoose_incl[2][5];
    int nTight_incl[2][5];
    for (int i=0;i<2;i++) {
	for (int j=4;j>=0;j--) {
	    nLoose_incl[i][j] = 0;
	    nTight_incl[i][j] = 0;
	    for (int k=j;k<5;k++) {
		nLoose_incl[i][j] += _nJetMult_loose[i][k];
		nTight_incl[i][j] += _nJetMult_tight[i][k];
	    }
	}
    }
    
    // print pass counts per jet multiplicity
    vector<string> channel;
    channel.push_back("e+jets");
    channel.push_back("mu+jets");
    
    
    for (int i=0;i<2;i++) {

	cout << endl << channel[i] << ":\n--------\n" << endl;
    
	cout << setw(7)  << "JetMult" << " "
	     << setw(15) << "Nloose"  << " "
	     << setw(15) << "Ntight"  << " "
	     << setw(15) << "Eff"
	     << endl;
	cout << "--------------------------------------------------------" 
	     << endl;
	for (int j=0;j<5;j++) {
	    cout << setw(7) << j+1                   << " " 
		 << setw(7) << _nJetMult_loose[i][j] << " " 
		 << setw(7) << nLoose_incl[i][j]     << " " 
		 << setw(7) << _nJetMult_tight[i][j] << " "
		 << setw(7) << nTight_incl[i][j]     << " " 
		 << setprecision(3)
		 << setw(7) 
		 << _nJetMult_tight[i][j]/(double)_nJetMult_loose[i][j] << " "
		 << setw(7) << nTight_incl[i][j]/(double)nLoose_incl[i][j]
		 << endl;
	}
    }

    
    //
    // cut flow
    //
    vector<string> names;
    vector<int>    npass;
    
    names.push_back("EventList");
    names.push_back("PrimaryVertex");
    names.push_back("IsolatedLepton");
    names.push_back("MissingEt");
    names.push_back("Trigger");
    names.push_back("SLTVeto");
    names.push_back("TightIsolation");
    names.push_back("HitFit");

    
    cout << "\nCUT FLOW:" << "\n_________\n"<< endl;
    
    for (int i=0;i<2;i++) {

	cout << channel[i] << ":\n" << endl;
	
	npass.clear();
	npass.push_back(_nPassEventList);
	npass.push_back(_nPassPv);
	npass.push_back(_nPassLepSel[i]);
	npass.push_back(_nPassMEt[i]);
	npass.push_back(_nPassTrigger[i]);
	npass.push_back(_nPassSLTVeto[i]);
	npass.push_back(_nPassTight[i]);
        npass.push_back(_nPassHitFit);

	for (int j=0;j<names.size();j++) {
	    cout << resetiosflags(ios::adjustfield) 
		 << setiosflags(ios::left)
		 << setw(20)
		 << names[j]
		 << resetiosflags(ios::adjustfield) 
		 << setiosflags(ios::right)
		 << setw(6)
		 << npass[j]
         	 << "  ";
	    
	    if (j>0) {
		cout << setw(6)
		     << setprecision(3)
		     << npass[j]/(double)npass[j-1]
		     << "  "
		     << setw(6)
		     << setprecision(3)
		     << npass[j]/(double)npass[0]
		     << endl;
	    }
	    else cout << endl;
	}
	
	cout << endl;
    }
    
    cout << _selectionTool->printCutFlow("ejets")  << endl;
    cout << _selectionTool->printCutFlow("mujets") << endl;

    return;
}


//______________________________________________________________________________
void LepJetsAnalysis::initOutputTree(const string& fileName)
{
    _outputTreeFile = new TFile(fileName.c_str(), "RECREATE");
    _outputTree     = new TTree("METree", "Top Mass Matrix Element Method Tree");
    _outputTree->SetMaxVirtualSize(32000000);
    
    
    _outputTree->Branch("run",          &_run,          "run/I");
    _outputTree->Branch("evt",          &_evt,          "evt/I");
    _outputTree->Branch("type",         &_type,         "type/I");    

    _outputTree->Branch("weight",       &_weight,       "weight/D");
    _outputTree->Branch("muidweight",   &_muidweight,   "muidweight/D");
    _outputTree->Branch("trgweight",    &_trgweight,    "trgweight/D");

    if (_triggerVariation)
	_outputTree->Branch("trgweights",    &_trgweights,    "trgweights[18]/D");

    _outputTree->Branch("tagprobTight", &_tagprobTight, "tagprobTight/D");
    _outputTree->Branch("tagprobMedium",&_tagprobMedium,"tagprobMedium/D");
    _outputTree->Branch("tagprobLoose", &_tagprobLoose, "tagprobLoose/D");

    _outputTree->Branch("dbltagprobTight", &_dbltagprobTight, "dbltagprobTight/D");
    _outputTree->Branch("dbltagprobMedium",&_dbltagprobMedium,"dbltagprobMedium/D");
    _outputTree->Branch("dbltagprobLoose", &_dbltagprobLoose, "dbltagprobLoose/D");
    
    _outputTree->Branch("rof",          &_rof,          "rof/I");
    _outputTree->Branch("cnoise",       &_cnoise,       "cnoise/I");
    _outputTree->Branch("ecrate",       &_ecrate,       "ecrate/I");
    
    _outputTree->Branch("met",          &_met,          "met/D");
    _outputTree->Branch("metx",         &_metx,         "metx/D");
    _outputTree->Branch("mety",         &_mety,         "mety/D");
    _outputTree->Branch("metphi",       &_metphi,       "metphi/D");
    _outputTree->Branch("metedc",       &_metedc,       "metedc/D");
    
    _outputTree->Branch("wmt",          &_wmt,          "wmt/D");
    _outputTree->Branch("wet",          &_wet,          "wet/D");
    _outputTree->Branch("weta",         &_weta,         "weta/D");

    _outputTree->Branch("ht20",         &_ht20,         "ht20/D");
    _outputTree->Branch("ht25",         &_ht25,         "ht25/D");
    _outputTree->Branch("apla",         &_apla,         "apla/D");
    _outputTree->Branch("spher",        &_spher,        "spher/D");
    _outputTree->Branch("ktminp",       &_ktminp,       "ktminp/D");
    _outputTree->Branch("ht2p",         &_ht2p,         "ht2p/D");
    _outputTree->Branch("lhood",        &_lhood,        "lhood/D");
    _outputTree->Branch("nnd",          &_nnd,          "nnd/D");
    _outputTree->Branch("disc",         &_disc,          "disc/D");
    
    _outputTree->Branch("leptype",      &_leptype,      "leptype/I");
    _outputTree->Branch("lepiso",       &_lepiso,       "lepiso/I");
    _outputTree->Branch("lepe",         &_lepe,         "lepe/D");
    _outputTree->Branch("leppt",        &_leppt,        "leppt/D");
    _outputTree->Branch("lepeta",       &_lepeta,       "lepeta/D");
    _outputTree->Branch("lepphi",       &_lepphi,       "lepphi/D");
    _outputTree->Branch("lepdeta",      &_lepdeta,      "lepdeta/D");
    _outputTree->Branch("lepdphimet",   &_lepdphimet,   "lepdphimet/D");
    _outputTree->Branch("lepnsmt",      &_lepnsmt,      "lepnsmt/I");
    
    _outputTree->Branch("lepmce",       &_lepmce,       "lepmce/D");
    _outputTree->Branch("lepmcpt",      &_lepmcpt,      "lepmcpt/D");
    _outputTree->Branch("lepmceta",     &_lepmceta,     "lepmceta/D");
    _outputTree->Branch("lepmcphi",     &_lepmcphi,     "lepmcphi/D");
    
    _outputTree->Branch("njt",          &_njt,          "njt/I");
    _outputTree->Branch("jte",          &_jte,          "jte[njt]/D");
    _outputTree->Branch("jtpt",         &_jtpt,         "jtpt[njt]/D");
    _outputTree->Branch("jteta",        &_jteta,        "jteta[njt]/D");
    _outputTree->Branch("jtphi",        &_jtphi,        "jtphi[njt]/D");
    _outputTree->Branch("jtdeta",       &_jtdeta,       "jtdeta[njt]/D");
    _outputTree->Branch("jtdphimet",    &_jtdphimet,    "jtdphimet[njt]/D");
    _outputTree->Branch("jtsvxtag",     &_jtsvxtag,     "jtsvxtag[njt]/I");
    _outputTree->Branch("jtmutag",      &_jtmutag,      "jtmutag[njt]/I");
    _outputTree->Branch("jtjes",        &_jtjes,        "jtjes[njt]/D");
    _outputTree->Branch("jtjesdata",    &_jtjesdata,    "jtjesdata[njt]/D");
    _outputTree->Branch("jtjesmc",      &_jtjesmc,      "jtjesmc[njt]/D");

    _outputTree->Branch("jttaggable",    &_jttaggable,    "jttaggable[njt]/D");
    _outputTree->Branch("jttagprobTight",&_jttagprobTight,"jttagprobTight[njt]/D");
    _outputTree->Branch("jttagprobMedium",&_jttagprobMedium,"jttagprobMedium[njt]/D");
    _outputTree->Branch("jttagprobLoose",&_jttagprobLoose,"jttagprobLoose[njt]/D");

    _outputTree->Branch("jtmcid",       &_jtmcid,       "jtmcid[njt]/I");
    _outputTree->Branch("jtmce",        &_jtmce,        "jtmce[njt]/D");
    _outputTree->Branch("jtmcpt",       &_jtmcpt,       "jtmcpt[njt]/D");
    _outputTree->Branch("jtmceta",      &_jtmceta,      "jtmceta[njt]/D");
    _outputTree->Branch("jtmcphi",      &_jtmcphi,      "jtmcphi[njt]/D");
    
    _outputTree->Branch("nperm",        &_nperm,        "nperm/I");
    _outputTree->Branch("perm",         &_perm,         "perm[nperm]/I");
    _outputTree->Branch("mt",           &_mt,           "mt[nperm]/D");
    _outputTree->Branch("chisq",        &_chisq,        "chisq[nperm]/D");        
    _outputTree->Branch("sigmt",        &_sigmt,        "sigmt[nperm]/D");        
    _outputTree->Branch("icorrect",     &_icorrect,     "icorrect/I");

    return;
    
}


//________________________________________________________________________________
void LepJetsAnalysis::setNewMuonPt()
{
    float metx = missingEt()->GetMJMEtx();
    float mety = missingEt()->GetMJMEty();   
  
    TClonesArray* muons = objects()->GetMuons();
    for (unsigned int i=0; i != muons->GetEntries(); ++i) {
	TheMuonClass* muon = (TheMuonClass*)muons->At(i);
	if(muon -> SMTHits() != 0) continue;
	if(muon -> HasCentral() == 0) continue;
    
	float pt_old = muon -> pT();
	float px_old = muon -> GetPx();
	float py_old = muon -> GetPy();
	int idxtrk   = muon -> TrackIndex();
	float dca_mu = muon -> DCA(); 

	TClonesArray* trks = objects()->GetTracks();
	TheTrackClass* trk = (TheTrackClass*)trks->At(idxtrk);
	float qopt = trk -> GetParameter(4);   // qopt
	float err_r_qopt = trk -> GetError(4);   // r qopt
	float err_r_r = trk -> GetError(0);   // r r


	// calculate new pt
	float pt_new = qopt - dca_mu * (err_r_qopt / err_r_r);
	pt_new = std::abs(1./pt_new);

	// set new muon momentum
	float SF = pt_new/pt_old;
	muon -> SetPt(pt_new);

	muon -> SetPx(SF*px_old);
	muon -> SetPy(SF*py_old);

	// propagate to met
	metx -= SF*px_old - px_old;
	mety -= SF*py_old - py_old;
    }
    // set new met
    missingEt() -> SetMJMEt ( metx , mety );
}


////////////////////////////////////////////////////////////////////////////////
// implementation of local functions
////////////////////////////////////////////////////////////////////////////////

//________________________________________________________________________
double TF2Lookup(TF2 *func, Float_t xval, Float_t yval)
{
    if(func != 0) {
	double xval_max = func->GetXmax();
	double xval_min = func->GetXmin();
	double yval_max = func->GetYmax();
	double yval_min = func->GetYmin();
	if (xval > xval_max) xval = xval_max;
	if (xval < xval_min) xval = xval_min;
	if (yval > yval_max) yval = yval_max;
	if (yval < yval_min) yval = yval_min;
	
	double val = func->Eval(xval,yval);

	if (val==0.0) cout<<"val=0! xval="<<xval<<" yval="<<yval<<endl;
	
	return val;
    }

    cout<<"WARNING! This Function does not exist."<<endl;
    return 0;
}