// -------------------------------------------------------------------------
// -----          FairEvtFilterOnParticleDist source file              -----
// -------------------------------------------------------------------------

#include "FairEvtFilterOnParticleDist.h"
#include "TFile.h"
#include "TTree.h"
#include "FairRunSim.h"

#define MassProton 0.938272

std::ostream& operator <<(std::ostream& os, const std::pair< std::pair<Int_t, Int_t>, TF1 >& gd)
{
    std::pair< Int_t, Int_t > applyTo = gd.first;
    TF1 * applyFunc = (TF1*) &(gd.second);

    // pdgID
    os << "[< " << applyTo.first << ", ";

    // GeomFilter
    switch ( applyTo.second ) {
        case FairEvtFilterOnParticleDist::kTheta:
            os << "kTheta";
            break;

        case FairEvtFilterOnParticleDist::kCosTheta:
            os << "kCosTheta";
            break;

        case FairEvtFilterOnParticleDist::kPhi:
            os << "kPhi";
            break;

        case FairEvtFilterOnParticleDist::kVertexZ:
            os << "kVertexZ";
            break;

        case FairEvtFilterOnParticleDist::kVertexRho:
            os << "kVertexRho";
            break;

        case FairEvtFilterOnParticleDist::kVertexRadius:
            os << "kVertexRadius";
            break;

        default:
            os << "?";
    }

    // function
    os << " >, " << applyFunc->GetExpFormula("p") << "]";

    return os;
}


std::ostream& operator <<(std::ostream& os, const std::vector< std::pair< std::pair<Int_t, Int_t>, TF1 > >& gd)
{
    os << "(";
    Int_t lastIdx = gd.size()-1;
    if(lastIdx < 0){
        os << ")" ;
        return os;
    }

    for (UInt_t iVec = 0; iVec < lastIdx; ++iVec){
        os << gd[iVec] << ", ";
    }
    os << gd[lastIdx] << ")";
    return os;
}


void FairEvtFilterOnParticleDist::PrintVerboseOutputHeader(TString extraText = "")
{
    std::cout << std::endl << "[INFO   ] ------ " << this->GetTitle() << ": " << this->GetName() << " ------ " << std::endl;
    if ( extraText != "" ) {
        std::cout << "                  (" << extraText << ")" << std::endl;
    }
}


// -----   Default constructor   -------------------------------------------
FairEvtFilterOnParticleDist::FairEvtFilterOnParticleDist(const Float_t mom)
        : FairEvtFilter(" ", "FairEvtFilterOnParticleDist")
{
    InitParameters(mom);
}


// -----   Constructor with name and title   ------------------------------
FairEvtFilterOnParticleDist::FairEvtFilterOnParticleDist(const char* name, const Float_t mom, const char* title)
        : FairEvtFilter(name, title)
{
    InitParameters(mom);
}


void FairEvtFilterOnParticleDist::InitParameters(const Float_t mom)
{
    // particle system parameters
    Float_t energy = MassProton + TMath::Sqrt(MassProton*MassProton + mom*mom);
    fPbarpCMS = TLorentzVector();
    fPbarpCMS.SetPxPyPzE(0, 0, mom, energy);
    fPbarpCMSboost = -fPbarpCMS.BoostVector();

    // filter active default
    fDistFilterActive = kFALSE;

    // random number generator instance
    fRnd = TRandom3();

    // the statistics nTuple
    fFilterStats = new RhoTuple("EvtFilterDist", "Filtered events based on probability distribution");

    std::cout << this->GetTitle() << ": Initialize " << this->GetName() << std::endl;
    std::cout << "   fPbarpCMS: "; fPbarpCMS.Print();
    std::cout << "   fPbarpCMSboost: "; fPbarpCMSboost.Print();
    std::cout << std::endl;
}


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


// -----   Add filter method -----------------------------------------------
Bool_t FairEvtFilterOnParticleDist::AndGeomDist(TF1 * func, Int_t pdgCode, GeomFilter geom)
{
    if (fVerbose > 0) {
        PrintVerboseOutputHeader();
        std::cout << fGeomDist << std::endl << std::endl;
    }

    // build the object to store
    GeomDist tmpDist;
    tmpDist.first.first = pdgCode;
    tmpDist.first.second = geom;
    tmpDist.second = *(TF1*)(func->Clone("test"));

    // saving the distribution function in the vector
    fGeomDist.push_back(tmpDist);

    // turn filter on
    fDistFilterActive = kTRUE;

    if (fVerbose > 0) {
        PrintVerboseOutputHeader("after adding distribution filter");
        std::cout << fGeomDist << std::endl << std::endl;
    }

    return kTRUE;
}


Float_t FairEvtFilterOnParticleDist::GetGeomValue(TParticle * par, GeomFilter geom)
{
    // the return value
    Float_t geomValue = TMath::QuietNaN();

    // fill a lorentz vector with the particle's parameters
    TLorentzVector lv = TLorentzVector( par->Px(), par->Py(), par->Pz(), par->Energy() );

    switch ( geom ) {
        case kCosTheta:
            // get cos(theta) from the CMS
            lv.Boost(fPbarpCMSboost);
            geomValue = lv.CosTheta();
            break;

        default:
            break;
    }

    return geomValue;
}


Bool_t FairEvtFilterOnParticleDist::RejectFromDist(TParticle * particle, Int_t evtNr, Int_t trkNr)
{
    Bool_t rejectThis;
    Bool_t rejectEvt = kFALSE;

    // values from the filter
    GeomFilter filterGeom;
    Int_t      filterPdgCode;
    TF1 *      filterFunc;

    Float_t geomValue, probValue;

    // loop over all set filters
    for ( Int_t iFilter = 0; iFilter < fGeomDist.size(); iFilter ++ ) {

        rejectThis = kFALSE;

        if ( fVerbose > 5 ) {
            std::cout << "Filter " << iFilter << ": " << fGeomDist[iFilter] << std::endl;
        }

        filterPdgCode = fGeomDist[iFilter].first.first;
        filterGeom = (GeomFilter)fGeomDist[iFilter].first.second;
        filterFunc = &fGeomDist[iFilter].second;

        // skip if the filter doesn't apply to the current particle
        if ( filterPdgCode != particle->GetPdgCode() ) {
            if ( fVerbose > 10 ) {
                std::cout << "   skipped filter -- particle code (" << particle->GetPdgCode() << ") doesn't match filter's requirement." << std::endl;
            }
            continue;
        }

        // get the geometrical value the filter applies to
        geomValue = GetGeomValue(particle, filterGeom);

        // skip if the geometrical value is not a number
        if ( TMath::IsNaN(geomValue) ) {
            if ( fVerbose > 5 ) {
                std::cout << "   skipped filter -- geometrical value is not a number (NaN)." << std::endl;
            }
            continue;
        }

        // skip if the geometrical value lies outside of the function's range
        if ( geomValue < filterFunc->GetXmin() ||
             geomValue > filterFunc->GetXmax() ) {
            if ( fVerbose > 5 ) {
                std::cout << "   skipped filter -- geometrical value (" << geomValue << ") outside of function's range." << std::endl;
            }
            continue;
        }

        // reject if randomly generated value is above probability function
        probValue = filterFunc->Eval(geomValue);
        if ( probValue < fRnd.Rndm() ) {
            if ( fVerbose > 5 ) {
                std::cout << "   passed filter  -- REJECTED. Probability of passing: " << probValue << "." << std::endl;
            }
            rejectThis = kTRUE;
        } else {
            if ( fVerbose > 5 ) {
                std::cout << "   passed filter  -- ACCEPTED. Probability of passing: " << probValue << "." << std::endl;
            }
        }

        if ( rejectThis ) { rejectEvt = kTRUE; }

        // store some statistics
        fFilterStats->Column("evtNr",         (Int_t)   evtNr);
        fFilterStats->Column("trkNr",         (Int_t)   trkNr);
        fFilterStats->Column("filterPdgCode", (Int_t)   filterPdgCode);
        fFilterStats->Column("filterGeom",    (Int_t)   filterGeom);
        fFilterStats->Column("pdgCode",       (Int_t)   particle->GetPdgCode());
        fFilterStats->Column("geomValue",     (Float_t) geomValue);
        fFilterStats->Column("probability",   (Float_t) probValue);
        fFilterStats->Column("rejected",      (Bool_t)  rejectThis);
        fFilterStats->DumpData();
    }

    return rejectEvt;
}


Bool_t FairEvtFilterOnParticleDist::EventMatches(Int_t evtNr)
{
    if (fVerbose > 3) {
        std::cout << std::endl << std::endl;
        std::cout << "Generated event: " << evtNr << std::endl;
        PrintVerboseOutputHeader("Beginning of EventMatches");
        std::cout << "Nr. of simulated particles: " << fParticleList->GetEntries()  << std::endl;
        PrintAllTParticleInEvent();
        std::cout << "fGeomDist: " << fGeomDist << std::endl << std::endl;
    }

    Bool_t distOk = kTRUE;

    // loop over the produces particles and check for applicable distributions
    for ( Int_t iPart = 0; iPart < fParticleList->GetEntries(); iPart ++ ) {

        TParticle * particle = (TParticle*)fParticleList->At(iPart);
        if ( 0 == particle ) { continue; }


        // check whether the particle satisfies the probability distribution
        if ( fDistFilterActive ) {
            // if the particle doesn't fulfill the distribution filter, revoke
            // the okay flag
            if ( RejectFromDist(particle, evtNr, iPart) ) { distOk = kFALSE; }
        }
    } // end of the loop over all particles

    Bool_t evtOk = distOk;

    if ( fVerbose > 5 ){
        if ( kTRUE == evtOk ) {
            std::cout << "\n Event matches " << this->GetTitle() << ": " << this->GetName() << "\n\n";
        } else {
            std::cout << "\n Event does NOT match " << this->GetTitle() << ": " << this->GetName() << "\n\n";
        }
    }

    return evtOk;
}


void FairEvtFilterOnParticleDist::WriteEvtFilterStatsToRootFile()
{
    // get the output instance
    TFile * outputFile;
    outputFile = FairRunSim::Instance()->GetOutputFile();
    outputFile->cd();

    // write the statistics tuple
    fFilterStats->GetInternalTree()->Write();

    outputFile->cd();
}


ClassImp(FairEvtFilterOnParticleDist)