// -------------------------------------------------------------------------
// -----               PndRhoFitProbPresorter source file              -----
// -------------------------------------------------------------------------

#include <iostream>

#include "TMath.h"

#include "PndRhoFitProbPresorter.h"
#include "RhoCandidate.h"

#include "PndAnalysis.h"
//#include "PndKinFitter.h"
//#include "PndKalmanVtxFitter.h"


// define a cutoff for probability ratio calculation
// (otherwise TBrowser gets a problem of displaying)
#define CUTOFF_PROB 1.0e-10


PndRhoFitProbPresorter::PndRhoFitProbPresorter(RhoCandList * particleList, float m0) :
	fCandList(particleList), fm0(m0), fVtxProbLimit(0.01), fMassProbLimit(0.01),
	fvtx_bestFitProb(0), fmass_bestFitProb(0), fvtx_2ndBestFitProb(0), fmass_2ndBestFitProb(0)
{
}

PndRhoFitProbPresorter::~PndRhoFitProbPresorter()
{
}

void PndRhoFitProbPresorter::GenerateIndices()
{
    _FillSortedMaps();
    _GenerateVtxFitIndices();
    _GenerateMassFitIndices();
}


/* get the best fit index */
int PndRhoFitProbPresorter::GetVtxBestFitIndex(int id)
{
    std::map<int, int>::iterator it = fvtx_indexOfBestFit.find(id);
    if ( it != fvtx_indexOfBestFit.end() )
        return it->second;
    else
        return -1;
}

int PndRhoFitProbPresorter::GetMassBestFitIndex(int id)
{
    std::map<int, int>::iterator it = fmass_indexOfBestFit.find(id);
    if ( it != fmass_indexOfBestFit.end() )
        return it->second;
    else
        return -1;
}


/* get the probability index */
int PndRhoFitProbPresorter::GetVtxProbIndex(int id)
{
    std::map<int, int>::iterator it = fvtx_indexOfFitProb.find(id);
    if ( it != fvtx_indexOfFitProb.end() )
        return it->second;
    else
        return -1;
}

int PndRhoFitProbPresorter::GetMassProbIndex(int id)
{
    std::map<int, int>::iterator it = fmass_indexOfFitProb.find(id);
    if ( it != fmass_indexOfFitProb.end() )
        return it->second;
    else
        return -1;
}



/* get the ratio to the best probability */
float PndRhoFitProbPresorter::GetVtxRatioToBestProb(int id)
{
    float selectedProb, probRatio = -1;

    if ( 0x0 == GetVtxFit(id) )
        return probRatio;

    selectedProb = GetVtxFit(id)->GetProb();
    if ( selectedProb < CUTOFF_PROB ) selectedProb = 0;

    if ( !TMath::IsNaN(selectedProb) && selectedProb != 0 ) {
        probRatio = fvtx_bestFitProb / selectedProb - 1;
    }

    return probRatio;
}

float PndRhoFitProbPresorter::GetVtxRatio2ndToBestProb() {
    float probRatio = -1;

    if ( TMath::IsNaN(fvtx_2ndBestFitProb) || fvtx_2ndBestFitProb == 0.0 || fvtx_2ndBestFitProb < CUTOFF_PROB )
        return probRatio;

    probRatio = fvtx_bestFitProb / fvtx_2ndBestFitProb - 1;

    return probRatio;
}

float PndRhoFitProbPresorter::GetMassRatioToBestProb(int id)
{
    float selectedProb, probRatio = -1;

    if ( 0x0 == GetMassFit(id) )
        return probRatio;

    selectedProb = GetMassFit(id)->GetProb();
    if ( selectedProb < CUTOFF_PROB ) selectedProb = 0;

    if ( !TMath::IsNaN(selectedProb) && selectedProb != 0.0 ) {
        probRatio = fmass_bestFitProb / selectedProb - 1;
    }

    return probRatio;
}

float PndRhoFitProbPresorter::GetMassRatio2ndToBestProb() {
    float probRatio = -1;

    if ( TMath::IsNaN(fmass_2ndBestFitProb) || fmass_2ndBestFitProb == 0.0 || fmass_2ndBestFitProb < CUTOFF_PROB )
        return probRatio;

    probRatio = fmass_bestFitProb / fmass_2ndBestFitProb - 1;

    return probRatio;
}

// Accessor functions to the fit objects
PndKinVtxFitter* PndRhoFitProbPresorter::GetVtxFit(int id)
{
    std::map<int, PndKinVtxFitter*>::iterator it = fVtxFit.find(id);
    if ( it != fVtxFit.end() ) {
        return it->second;
    } else {
        return 0x0;
    }
}

PndKinFitter* PndRhoFitProbPresorter::GetMassFit(int id)
{
    std::map<int, PndKinFitter*>::iterator it = fMassFit.find(id);
    if ( it != fMassFit.end() ) {
        return it->second;
    } else {
        return 0x0;
    }
}


//
// private functions
//

void PndRhoFitProbPresorter::_FillSortedMaps()
{
    // loop through all particle candidates
    for (int j = 0; j < fCandList->GetLength(); j++) {
        RhoCandidate * curCand = fCandList->Get(j);

        // fill the lookup table
        fUidToCandId[curCand->Uid()] = j;

        // do the vertex fit
        PndKinVtxFitter * vertexFitter = new PndKinVtxFitter(fCandList->Get(j));
        fVtxFit[j] = vertexFitter;
        vertexFitter->Fit();

        // continue if the vertex fitter failed
        bool failVtxProb = TMath::IsNaN(vertexFitter->GetProb());
        bool failVtxChi2 = TMath::IsNaN(vertexFitter->GetChi2());
        if (failVtxChi2 || failVtxProb) continue;

        // store chi2/prob information of vertex fit
        fvtx_probForIndex[-vertexFitter->GetProb()] = j; // use negative prob to reverse sorting order
        if (vertexFitter->GetProb() > fVtxProbLimit) {
            fvtx_chi2ForIndex_good[vertexFitter->GetChi2()] = j;
        } else {    // <= 0.01 as default
            fvtx_chi2ForIndex_bad[vertexFitter->GetChi2()] = j;
        }

        // do the mass constraint fit
        PndKinFitter * massFitter = new PndKinFitter(curCand->GetFit());
        fMassFit[j] = massFitter;
        massFitter->AddMassConstraint(fm0);
        massFitter->Fit();

        // continue if the mass constraint fit failed
        bool failMassProb = TMath::IsNaN(massFitter->GetProb());
        bool failMassChi2 = TMath::IsNaN(massFitter->GetChi2());
        if (failMassProb || failMassChi2) continue;

        // store chi2/prob information of mass fit
        fmass_probForIndex[-massFitter->GetProb()] = j; // use negative prob to reverse sorting order
        if (massFitter->GetProb() > fMassProbLimit) {
            fmass_chi2ForIndex_good[massFitter->GetChi2()] = j;
        } else {    // <= 0.01 as default
            fmass_chi2ForIndex_bad[massFitter->GetChi2()] = j;
        }

    }

}

void PndRhoFitProbPresorter::_GenerateVtxFitIndices()
{
    std::map<double, int>::iterator it;
    int runningI = 0;
    for (it = fvtx_chi2ForIndex_good.begin(); it != fvtx_chi2ForIndex_good.end(); it++, runningI++) {
        fvtx_indexOfBestFit[it->second] = runningI + 1;
    }
    runningI = 0;
    for (it = fvtx_chi2ForIndex_bad.begin(); it != fvtx_chi2ForIndex_bad.end(); it++, runningI++) {
        fvtx_indexOfBestFit[it->second] = - (runningI + 1);
    }

    // the sorted probability index
    runningI = 0;
    fvtx_bestFitProb = 0;
    for (it = fvtx_probForIndex.begin(); it != fvtx_probForIndex.end(); it++, runningI++) {
        fvtx_indexOfFitProb[it->second] = runningI + 1;
        if ( runningI == 0 ) fvtx_bestFitProb = -it->first;
        if ( runningI == 1 ) fvtx_2ndBestFitProb = -it->first;
    }
}

void PndRhoFitProbPresorter::_GenerateMassFitIndices()
{
    std::map<double, int>::iterator it;
    int runningI = 0;
    for (it = fmass_chi2ForIndex_good.begin(); it != fmass_chi2ForIndex_good.end(); it++, runningI++) {
        fmass_indexOfBestFit[it->second] = runningI + 1;
    }
    runningI = 0;
    for (it = fmass_chi2ForIndex_bad.begin(); it != fmass_chi2ForIndex_bad.end(); it++, runningI++) {
        fmass_indexOfBestFit[it->second] = - (runningI + 1);
    }

    // the sorted probability index
    runningI = 0;
    fmass_bestFitProb = 0;
    for (it = fmass_probForIndex.begin(); it != fmass_probForIndex.end(); it++, runningI++) {
        fmass_indexOfFitProb[it->second] = runningI + 1;
        if ( runningI == 0 ) fmass_bestFitProb = -it->first;
        if ( runningI == 1 ) fmass_2ndBestFitProb = -it->first;
    }
}