#include "AtlEvtWriterD3PDSgTopThinned.h"

#include <cstddef>

#include <TList.h>
#include <TTree.h>

#include <AtlEvent.h>


AtlEvtWriterD3PDSgTopThinned::AtlEvtWriterD3PDSgTopThinned()
    : AtlEvtWriterBase()
    , fJets(0)
    , fLeptons(0)
    , v_jet_pt(new Float_t(kMaxJets))
    , v_jet_eta(new Float_t(kMaxJets))
    , v_jet_phi(new Float_t(kMaxJets))
    , v_lep_pt(new Float_t(kMaxLeptons))
    , v_lep_eta(new Float_t(kMaxLeptons))
    , v_lep_phi(new Float_t(kMaxLeptons))
    , v_lep_pdgId(new Int_t(kMaxLeptons))
{
}

AtlEvtWriterD3PDSgTopThinned::~AtlEvtWriterD3PDSgTopThinned() {
    delete v_jet_pt;
    delete v_jet_eta;
    delete v_jet_phi;
    delete v_lep_pt;
    delete v_lep_eta;
    delete v_lep_phi;
    delete v_lep_pdgId;
}

void AtlEvtWriterD3PDSgTopThinned::SetJets(TList * jets) {
    fJets = jets;
}

void AtlEvtWriterD3PDSgTopThinned::SetLeptons(TList * leptons) {
    fLeptons = leptons;
}

void AtlEvtWriterD3PDSgTopThinned::BookTree(TTree * tree, AtlEvent * evt) {
    fEvent = evt;
    fTree->Branch("evt_weight", &v_evt_weight, "evt_weight/F");
    fTree->Branch("jet_n", &v_jet_n, "jet_n/I");
    fTree->Branch("jet_pt", v_jet_pt, "jet_pt[jet_n]/F");
    fTree->Branch("jet_eta", v_jet_eta, "jet_eta[jet_n]/F");
    fTree->Branch("jet_phi", v_jet_phi, "jet_phi[jet_n]/F");
    fTree->Branch("lep_n", &v_lep_n, "lep_n/I");
    fTree->Branch("lep_pt", v_lep_pt, "lep_pt[lep_n]/F");
    fTree->Branch("lep_eta", v_lep_eta, "lep_eta[lep_n]/F");
    fTree->Branch("lep_phi", v_lep_phi, "lep_phi[lep_n]/F");
    fTree->Branch("lep_pdgId", v_lep_pdgId, "lep_pdgId[lep_n]/I");
    fTree->Branch("met_pt", &v_met_pt, "met_pt/F");
}

TTree * AtlEvtWriterD3PDSgTopThinned::CreateTree(char const * name, char const * title, TTree * inputTree) {
    fTree = new TTree(name, title);
    fTree->SetAutoSave(0);
    return fTree;
}

void AtlEvtWriterD3PDSgTopThinned::WriteEvent() {

    /* event weight */
    AtlEventHeader const & eventHeader = *fEvent->GetEventHeader();
    if ( eventHeader.GetPileupWeight() == 0. )
        return;
    v_evt_weight = eventHeader.GetPreTagEvtWeight();

    /* jets */
    size_t jet_i = 0;
    for ( TIter next(fJets); HepJet * jet = static_cast<HepJet *>(next()); ) {
        if ( !(jet_i < kMaxJets) )
            break;
        v_jet_pt[jet_i] = jet->Pt();
        v_jet_eta[jet_i] = jet->Eta();
        v_jet_phi[jet_i] = jet->Phi();
        ++jet_i;
    }
    v_jet_n = jet_i;

    /* leptons */
    size_t lep_i = 0;
    for ( TIter next(fLeptons); HepParticle * lep = static_cast<HepParticle *>(next()); ) {
        if ( !(lep_i < kMaxLeptons) )
            break;
        v_lep_pt[lep_i] = lep->Pt();
        v_lep_eta[lep_i] = lep->Eta();
        v_lep_phi[lep_i] = lep->Phi();
        v_lep_pdgId[lep_i] = lep->GetPdgCode();
        ++lep_i;
    }
    v_lep_n = lep_i;

    /* save energy sums */
    v_met_pt = fEvent->GetEnergySum()->MissingEt_Mag();

    /* save! */
    fTree->Fill();
}