// -------------------------------------------------------------------------
// -----              PndTrackCandInfoTask source file                 -----
// -------------------------------------------------------------------------

#include <iostream>
#include "PndTrackCandInfoTask.h"
#include "PndTrackCand.h"
#include "PndMCTrack.h"


// -----   Constructor   ---------------------------------------------------
PndTrackCandInfoTask::PndTrackCandInfoTask() :
    fPersistence(kTRUE)
{
    fMvdPixelHitBranch = "MVDHitsPixel";
    fMvdStripHitBranch = "MVDHitsStrip";
    fSttHitBranch = "STTHit";
    fGemHitBranch = "GEMHit";
    fFtsHitBranch = "FTSHit";
}
// -------------------------------------------------------------------------

// -----   Destructor   ----------------------------------------------------
PndTrackCandInfoTask::~PndTrackCandInfoTask()
{
}
// -------------------------------------------------------------------------


// -----   Initialize   ----------------------------------------------------
InitStatus PndTrackCandInfoTask::Init() {

    FairRootManager * ioman = FairRootManager::Instance();
    if (!ioman) {
        std::cout << "-E- PndTrackCandInfoTask::Init: "
            << "RootManager not instantiated!" << std::endl;
        return kFATAL;
    }

    // load the track candidate and MC truth list
    fInputTrack = (TClonesArray*)ioman->GetObject( fTrackInputBranch.Data() );
    if ( !fInputTrack ) {
        std::cout << "-E- PndTrackCandInfoTask::Init: "
            << "TrackInputBranch coundn't be found."  << std::endl;
        return kERROR;
    }
    fMCTrack = (TClonesArray*)ioman->GetObject("MCTrack");
    if ( !fMCTrack ) {
        std::cout << "-E- PndTrackCandInfoTask::Init: "
            << "MCTrack coundn't be found. Make sure to add sim_params.root file."  << std::endl;
        return kERROR;
    }

    // register the output branch
    if ( !fTrackInfoOutputBranch ) {
        std::cout << "-E- PndTrackCandInfoTask::Init: "
            << "TrackInfoOutputBranch is not set."  << std::endl;
        return kERROR;
    }
    fHitCountsArray = new TClonesArray("PndTrackCandInfo");
    ioman->Register(
    	fTrackInfoOutputBranch, "TrackCandInfo", fHitCountsArray, fPersistence
    );

    // get detector hit IDs
    fHitIDMvdPixel = ioman->GetBranchId(fMvdPixelHitBranch);
    fHitIDMvdStrip = ioman->GetBranchId(fMvdStripHitBranch);
    fHitIDStt = ioman->GetBranchId(fSttHitBranch);
    fHitIDGem = ioman->GetBranchId(fGemHitBranch);
    fHitIDFts = ioman->GetBranchId(fFtsHitBranch);

    std::cout << "-I- PndTrackCandInfoTask::Init() finished" << std::endl;
    return kSUCCESS;
}
// -------------------------------------------------------------------------


// -----   Execute stuff for the event   -----------------------------------
void PndTrackCandInfoTask::Exec(Option_t * opt) {
    // Reset output array
    if ( ! fHitCountsArray ) {
        Fatal("Exec", "No fHitCountsArray initialized.");
    }
    fHitCountsArray->Delete();

    // variables used in the loop
    PndTrackCand * myTrackCand;
    PndMCTrack * myMcTrack;
    PndTrackCandInfo * myTrackCandInfo;
    TVector3 myStartVertex;
    TLorentzVector myMomentum;


    // go through the list of track candidates
    for ( int iTrkC = 0; iTrkC < fInputTrack->GetEntries(); iTrkC ++ ) {
        myTrackCand = (PndTrackCand*)fInputTrack->At(iTrkC);
        myMcTrack = (PndMCTrack*)fMCTrack->At(myTrackCand->getMcTrackId());

        // std::cout << "-I- PndTrackCandInfoTask::Exec:  " << std::endl;
        // std::cout << "   track candidate " << iTrkC << std::endl;
        // std::cout << "   PDG code: " << myMcTrack->GetPdgCode() << std::endl;
        // std::cout << "   hits by (MVDp,MVDs,STT,GEM) = ("
        //     << myTrackCand->GetNHitsDet(fHitIDMvdPixel) << ","
        //     << myTrackCand->GetNHitsDet(fHitIDMvdStrip) << ","
        //     << myTrackCand->GetNHitsDet(fHitIDStt) << ","
        //     << myTrackCand->GetNHitsDet(fHitIDGem) << ")" << std::endl;

        myTrackCandInfo = new PndTrackCandInfo();
        myTrackCandInfo->fHitsMvd = myTrackCand->GetNHitsDet(fHitIDMvdPixel) + myTrackCand->GetNHitsDet(fHitIDMvdStrip);
        myTrackCandInfo->fHitsStt = myTrackCand->GetNHitsDet(fHitIDStt);
        myTrackCandInfo->fHitsGem = myTrackCand->GetNHitsDet(fHitIDGem);
        myTrackCandInfo->fHitsFts = myTrackCand->GetNHitsDet(fHitIDFts);

        myTrackCandInfo->fMcPdgId = myMcTrack->GetPdgCode();
        myTrackCandInfo->fMotherId = myMcTrack->GetMotherID();

        myStartVertex = myMcTrack->GetStartVertex();
        myTrackCandInfo->fMcX = myStartVertex.X();
        myTrackCandInfo->fMcY = myStartVertex.Y();
        myTrackCandInfo->fMcZ = myStartVertex.Z();

        myMomentum = myMcTrack->Get4Momentum();
        myTrackCandInfo->fMcPx = myMomentum.Px();
        myTrackCandInfo->fMcPy = myMomentum.Py();
        myTrackCandInfo->fMcPz = myMomentum.Pz();
        myTrackCandInfo->fMcPt = myMomentum.Pt();
        myTrackCandInfo->fMcE = myMomentum.E();
        myTrackCandInfo->fMcTheta = myMomentum.Theta();
        myTrackCandInfo->fMcPhi = myMomentum.Phi();

        new((*fHitCountsArray)[fHitCountsArray->GetEntriesFast()])
            PndTrackCandInfo(*myTrackCandInfo);
    }
}
// -------------------------------------------------------------------------


// -----   Finish up   -----------------------------------------------------
void PndTrackCandInfoTask::Finish() {

}
// -------------------------------------------------------------------------


ClassImp(PndTrackCandInfoTask);