// -------------------------------------------------------------------------
// -----               PndTrackFilterTask source file                  -----
// -------------------------------------------------------------------------

#include <iostream>
#include "PndTrackFilterTask.h"
#include "PndTrackCand.h"
#include "PndTrackCandInfo.h"
#include "PndMCTrack.h"


// -----   Constructor   ---------------------------------------------------
PndTrackFilterTask::PndTrackFilterTask() :
    fPersistence(kTRUE), fVerbose(0), fPassCounter(0), fFailCounter(0)
{
    fMvdPixelHitBranch = "MVDHitsPixel";
    fMvdStripHitBranch = "MVDHitsStrip";
    fSttHitBranch = "STTHit";
    fGemHitBranch = "GEMHit";
    fFtsHitBranch = "FTSHit";
}
// -------------------------------------------------------------------------

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


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

    // load the track candidate and MC truth list
    fInputTrack = (TClonesArray*)ioman->GetObject( fTrackInBranchName.Data() );
    if ( !fInputTrack ) {
        std::cout << "-E- PndTrackFilterTask::Init: "
            << "TrackInputBranch coundn't be found."  << std::endl;
        return kERROR;
    }
    fInputTrackCandInfo = (TClonesArray*)ioman->GetObject( fTrackCandInfoInBranchName.Data() );
    if ( !fInputTrackCandInfo ) {
        std::cout << "-E- PndTrackFilterTask::Init: "
            << "TrackCandInfoInputBranch coundn't be found."  << std::endl;
        return kERROR;
    }

    // register the output branch
    if ( !fFilteredTrackOutBranchName ) {
        std::cout << "-E- PndTrackFilterTask::Init: "
            << "TrackOutputBranch is not set."  << std::endl;
        return kERROR;
    }
    fFilteredTrackArray = new TClonesArray("PndTrack");
    ioman->Register(
        fFilteredTrackOutBranchName, "FilteredTrack", fFilteredTrackArray, fPersistence
    );

    // register the output branch for track cand info
    if ( !fFilteredTrackCandInfoOutBranchName ) {
        std::cout << "-E- PndTrackFilterTask::Init: "
            << "TrackCandInfoOutputBranch is not set."  << std::endl;
        return kERROR;
    }
    fFilteredTrackCandInfoArray = new TClonesArray("PndTrackCandInfo");
    ioman->Register(
        fFilteredTrackCandInfoOutBranchName, "FilteredTrackCandInfo", fFilteredTrackCandInfoArray, 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- PndTrackFilterTask::Init() finished" << std::endl;
    return kSUCCESS;
}
// -------------------------------------------------------------------------


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

    // initialize variables
    PndTrack * myTrack;
    PndTrackCandInfo * myTrackCandInfo;

    // go through the list of track candidates
    for ( int iTrk = 0; iTrk < fInputTrack->GetEntries(); iTrk ++ ) {
        myTrack = (PndTrack*)fInputTrack->At(iTrk);
        myTrackCandInfo = (PndTrackCandInfo*)fInputTrackCandInfo->At(iTrk);

        // apply the filter for minimum hits
        if ( !TrackMatchesHitCountFilter(myTrack) )
            continue;

        new((*fFilteredTrackArray)[fFilteredTrackArray->GetEntriesFast()])
            PndTrack(*myTrack);
        new((*fFilteredTrackCandInfoArray)[fFilteredTrackCandInfoArray->GetEntriesFast()])
            PndTrackCandInfo(*myTrackCandInfo);
    }
}
// -------------------------------------------------------------------------


// -----   Finish up   -----------------------------------------------------
void PndTrackFilterTask::Finish()
{
    // print some statistics
    std::cout << "-I- PndTrackFilterTask::Finish:  Filter statistics" << std::endl;

    std::cout << "      Branch:  " << fTrackInBranchName << std::endl;

    std::cout << "      Filters: ";
    for ( std::vector<MinHits_t>::size_type i = 0; i != fMinHitsList.size(); i++ ) {
        if ( i > 0 ) std::cout << " or ";
        std::cout << fMinHitsList[i].GetPrint();
    }
    if ( fMinHitsList.size() > 0 && fMaxHitsList.size() > 0 ) {
        std::cout << std::endl << "               " << "  and  " << std::endl << "               ";
    }
    for ( std::vector<MaxHits_t>::size_type i = 0; i != fMaxHitsList.size(); i++ ) {
        if ( i > 0 ) std::cout << " or ";
        std::cout << fMaxHitsList[i].GetPrint();
    }
    std::cout << std::endl;

    std::cout << "      Tracks:" << std::endl;
    std::cout << "        total: " << fPassCounter+fFailCounter << std::endl;
    std::cout << "        pass:  " << fPassCounter << " (" << TString::Format("%.1f", (float)fPassCounter/(fPassCounter+fFailCounter)*100) << " %)" << std::endl;
    std::cout << "        fail:  " << fFailCounter << " (" << TString::Format("%.1f", (float)fFailCounter/(fPassCounter+fFailCounter)*100) << " %)" << std::endl;
    std::cout << std::endl;
}
// -------------------------------------------------------------------------


// -----   Print the contents of fMinHitsType in a nice format   -----------
TString PndTrackFilterTask::MinHits_t::GetPrint() {
    // collect the parts which contain something
    std::vector<TString> outParts;
    if ( Mvd > 0 ) outParts.push_back(TString::Format("MVD >= %d",Mvd));
    if ( Stt > 0 ) outParts.push_back(TString::Format("STT >= %d",Stt));
    if ( Gem > 0 ) outParts.push_back(TString::Format("GEM >= %d",Gem));
    if ( Fts > 0 ) outParts.push_back(TString::Format("FTS >= %d",Fts));
    if ( Tot > 0 ) outParts.push_back(TString::Format("total >= %d",Tot));

    // combine the parts together
    TString out = "[";
    for ( std::vector<TString>::size_type i = 0; i != outParts.size(); i++ ) {
        if ( i > 0 ) out.Append(" && ");
        out.Append(outParts[i]);
    }
    out.Append("]");
    return out;
}
// -------------------------------------------------------------------------


// -----   Print the contents of fMinHitsType in a nice format   -----------
TString PndTrackFilterTask::MaxHits_t::GetPrint() {
    // collect the parts which contain something
    std::vector<TString> outParts;
    if ( Mvd < 9999 ) outParts.push_back(TString::Format("MVD <= %d",Mvd));
    if ( Stt < 9999 ) outParts.push_back(TString::Format("STT <= %d",Stt));
    if ( Gem < 9999 ) outParts.push_back(TString::Format("GEM <= %d",Gem));
    if ( Fts < 9999 ) outParts.push_back(TString::Format("FTS <= %d",Fts));
    if ( Tot < 9999 ) outParts.push_back(TString::Format("total <= %d",Tot));

    // combine the parts together
    TString out = "[";
    for ( std::vector<TString>::size_type i = 0; i != outParts.size(); i++ ) {
        if ( i > 0 ) out.Append(" && ");
        out.Append(outParts[i]);
    }
    out.Append("]");
    return out;
}
// -------------------------------------------------------------------------


// -----   Add minimum number of hits filter   -----------------------------
void PndTrackFilterTask::AddMinHitsFilter(Int_t minMvd, Int_t minStt, Int_t minGem, Int_t minFts, Int_t minTot)
{
    MinHits_t minList;
    if ( minMvd > 0 ) minList.Mvd = minMvd;
    if ( minStt > 0 ) minList.Stt = minStt;
    if ( minGem > 0 ) minList.Gem = minGem;
    if ( minFts > 0 ) minList.Fts = minFts;
    if ( minTot > 0 ) minList.Tot = minTot;

    fMinHitsList.push_back(minList);

    if ( fVerbose > 0 ) {
        std::cout << "-I- PndTrackFilterTask::AddMinHitsFilter:" << std::endl;
        std::cout << "    Added " << minList.GetPrint() << " to the list." << std::endl;
    }
}
// -------------------------------------------------------------------------


// -----   Add maximum number of hits filter   -----------------------------
void PndTrackFilterTask::AddMaxHitsFilter(Int_t maxMvd, Int_t maxStt, Int_t maxGem, Int_t maxFts, Int_t maxTot)
{
    MaxHits_t maxList;
    if ( maxMvd < 9999 ) maxList.Mvd = maxMvd;
    if ( maxStt < 9999 ) maxList.Stt = maxStt;
    if ( maxGem < 9999 ) maxList.Gem = maxGem;
    if ( maxFts < 9999 ) maxList.Fts = maxFts;
    if ( maxTot < 9999 ) maxList.Tot = maxTot;

    fMaxHitsList.push_back(maxList);

    if ( fVerbose > 0 ) {
        std::cout << "-I- PndTrackFilterTask::AddMaxHitsFilter:" << std::endl;
        std::cout << "    Added " << maxList.GetPrint() << " to the list." << std::endl;
    }
}
// -------------------------------------------------------------------------


// -----   Check if the track matches the hit count filter   ---------------
Bool_t PndTrackFilterTask::TrackMatchesHitCountFilter(PndTrack * myTrack)
{
    // collect the hits in tracking detectors
    PndTrackCand myTrackCand = myTrack->GetTrackCand();
    Int_t hitsInMvd = myTrackCand.GetNHitsDet(fHitIDMvdPixel) + myTrackCand.GetNHitsDet(fHitIDMvdStrip);
    Int_t hitsInStt = myTrackCand.GetNHitsDet(fHitIDStt);
    Int_t hitsInGem = myTrackCand.GetNHitsDet(fHitIDGem);
    Int_t hitsInFts = myTrackCand.GetNHitsDet(fHitIDFts);
    Int_t hitsTotal = hitsInMvd + hitsInStt + hitsInGem + hitsInFts;

    if ( fVerbose > 9 ) {
        std::cout << "  McTrackId: " << myTrackCand.getMcTrackId() << std::endl;
        std::cout << "    Hits in MVD: " << hitsInMvd << std::endl;
        std::cout << "    Hits in STT: " << hitsInStt << std::endl;
        std::cout << "    Hits in GEM: " << hitsInGem << std::endl;
        std::cout << "    Hits in FTS: " << hitsInFts << std::endl;
        std::cout << "    Hits total:  " << hitsTotal << std::endl;
    }

    // loop through all hit count filters and see if one of them matches
    Bool_t match = false;
    for ( std::vector<MinHits_t>::iterator it = fMinHitsList.begin(); it != fMinHitsList.end(); ++it ) {
        MinHits_t minHitsIn = *it;

        if ( fVerbose > 9 ) {
            std::cout << "  Filter: " << minHitsIn.GetPrint() << std::endl;
        }

        if ( hitsInMvd >= minHitsIn.Mvd && hitsInStt >= minHitsIn.Stt &&
             hitsInGem >= minHitsIn.Gem && hitsInFts >= minHitsIn.Fts &&
             hitsTotal >= minHitsIn.Tot ) {
            match = true;
            break;
        }
    }

    // no minimum required == every track is fine
    if ( fMinHitsList.size() == 0 ) {
        match = true;
    }

    // check the maximum allowed hits
    for ( std::vector<MaxHits_t>::iterator it = fMaxHitsList.begin(); it != fMaxHitsList.end(); ++it ) {
        MaxHits_t maxHitsIn = *it;

        if ( fVerbose > 9 ) {
            std::cout << "  Filter: " << maxHitsIn.GetPrint() << std::endl;
        }

        if ( hitsInMvd > maxHitsIn.Mvd || hitsInStt > maxHitsIn.Stt ||
             hitsInGem > maxHitsIn.Gem || hitsInFts > maxHitsIn.Fts ||
             hitsTotal > maxHitsIn.Tot ) {
            match = false;
            break;
        }
    }

    if ( fVerbose > 5 ) {
        std::cout << "-I- PndTrackFilterTask::TrackMatchesHitCountFilter:  " <<
                ( match ? "keep" : "discard" ) << " track id " << myTrackCand.getMcTrackId() << std::endl;
    }

    // increase counters for statistics
    if ( match ) fPassCounter++;
    else         fFailCounter++;

    return match;
}
// -------------------------------------------------------------------------


ClassImp(PndTrackFilterTask);