#include <RAT/BerkeleyAlphaBeta.hh>
#include <RAT/DB.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DU/LightPathCalculator.hh>
#include <RAT/DU/Utility.hh>

#include <math.h>

using namespace RAT;
using namespace RAT::Classifiers;
using namespace RAT::DS;
using namespace RAT::DU;
using namespace std;

BerkeleyAlphaBeta::BerkeleyAlphaBeta() : fLastEventTime(NULL) {

}

BerkeleyAlphaBeta::~BerkeleyAlphaBeta() {
    if (fLastEventTime) delete fLastEventTime;
}

std::string BerkeleyAlphaBeta::GetName() const {
    return Name();
}

void BerkeleyAlphaBeta::Initialise(const std::string& param) {
    if (!DB::ParseTableName(param,fTBLName,fTBLIndex)) {
        fTBLName = "AB_PDF";
        try {
            try{
                DBLinkPtr inner_av = DB::Get()->GetLink("GEO","inner_av");
                fTBLIndex = inner_av->GetS("material");
            } catch(DBNotFoundError & ){
                try{
                    // If we are in a partial-fill geometry, use the upper
                    // material
                    DBLinkPtr inner_av = DB::Get()->GetLink("GEO","inner_av");
                    fTBLIndex = inner_av->GetS("material_top");
                }
                catch(DBNotFoundError & ) {
                    Log::Die("BerkeleyAlphaBeta: inner_av material and material_top "
                    "undefined. Try something like:\n"
                    "/rat/db/set GEO[inner_av] material "
                    "\"te_labppo_scintillator\" or "
                    "/rat/db/set GEO[inner_av] material_top "
                    "\"te_labppo_scintillator\"");
                }
            }
            info << fTBLIndex << endl;
            DBLinkPtr pdftable = DB::Get()->GetLink(fTBLName,fTBLIndex);
            pdftable->GetD("time_first"); //test to see if the table for the material exists
        } catch (...) {
            fTBLIndex = ""; //fallback tables
        }
    }
    info << "BerkeleyAlphaBeta: Using " << fTBLName << "[" << fTBLIndex << "] PDFs" << endl;
}

void BerkeleyAlphaBeta::BeginOfRun(DS::Run& /*ignored*/) {
    DBLinkPtr pdftable = DB::Get()->GetLink(fTBLName,fTBLIndex);

    fRetriggerCutoff = RAT::DS::UniversalTime(0.0,0.0,pdftable->GetD("retrig_cutoff"));

    fTimeFirst = pdftable->GetD("time_first");
    fTimeStep = pdftable->GetD("time_step");
    fScaleFactor = pdftable->GetD("scale_factor");

    fPDFAlphaLog = pdftable->GetDArray("pdf_alpha_prob");
    fPDFBetaLog = pdftable->GetDArray("pdf_beta_prob");
    Log::Assert((fNEntries = fPDFAlphaLog.size()) == fPDFBetaLog.size(),"BerkeleyAlphaBeta: Alpha PDF and Beta PDF should be same size");

    //Condition PDF and take logs ahead of time
    for (size_t i = 0; i < fNEntries; i++) {
        if (fPDFAlphaLog[i] <= 0) {
            fPDFAlphaLog[i] = 0.0;
        } else {
            fPDFAlphaLog[i] = log(fScaleFactor*fPDFAlphaLog[i]);
        }
        if (fPDFBetaLog[i] <= 0) {
            fPDFBetaLog[i] = 0.0;
        } else {
            fPDFBetaLog[i] = log(fScaleFactor*fPDFBetaLog[i]);
        }
    }
}


void BerkeleyAlphaBeta::EndOfRun(DS::Run& /*ignored*/) {

}

void BerkeleyAlphaBeta::DefaultSeed() {
    fEventPos = TVector3();
    fEventTime = 0.0;
}

void BerkeleyAlphaBeta::SetSeed(const DS::FitResult& seed) {
    //fEvent etc are updated before SetSeed is called
    RAT::DS::UniversalTime curEventTime = fEvent->GetUniversalTime();
    //fLastEventTime is NULL for first event
    if (fLastEventTime) {
        //if this event is within the fRetrigger cutoff of the last event (i.e. part of last event)
        fRetrigger = curEventTime - *fLastEventTime < fRetriggerCutoff;
        if (fRetrigger) {
            //offset event time by the time between the two triggers & use old position
            fEventTime -= (curEventTime - *fLastEventTime).GetNanoSeconds();
        } else {
            //use the seeded time and position
            fEventPos = seed.GetVertex(0).GetPosition();
            fEventTime = seed.GetVertex(0).GetTime();
        }
    } else {
        //first event init
        fRetrigger = false;
        fEventPos = seed.GetVertex(0).GetPosition();
        fEventTime = seed.GetVertex(0).GetTime();
        fLastEventTime = new UniversalTime();
    }
    *fLastEventTime = curEventTime;
}

DS::ClassifierResult BerkeleyAlphaBeta::GetClassification() {
    const PMTInfo &pmtInfo = Utility::Get()->GetPMTInfo();
    const EffectiveVelocity &effectiveVelocity = Utility::Get()->GetEffectiveVelocity();
    LightPathCalculator lightPathCalc = Utility::Get()->GetLightPathCalculator();

    //Reset from previous event
    fClassifierResult.Reset();

    //No PMTs -> do nothing
    if (fPMTData.empty())
        return fClassifierResult;

    //reset likelihood accumulators if this is not a fRetrigger
    if (!fRetrigger) {
        fLogAlphaProb = 0.0;
        fLogBetaProb = 0.0;
    }

    //Loop over selected PMTs
    for (vector<FitterPMT>::const_iterator pmt = fPMTData.begin(); pmt != fPMTData.end(); ++pmt) {
        // Calculate the hit time residuals
        const TVector3 pmtPos = pmtInfo.GetPosition(pmt->GetID());
        lightPathCalc.CalcByPosition(fEventPos,pmtPos);
        const double transitTime = effectiveVelocity.CalcByDistance(lightPathCalc.GetDistInInnerAV(),lightPathCalc.GetDistInAV(),lightPathCalc.GetDistInWater());
        const double timeResidual = pmt->GetTime() - transitTime - fEventTime;

        //Find the bin
        size_t pos = (size_t)((timeResidual-fTimeFirst)/fTimeStep);
        if (pos >= fNEntries || pos < 0) continue; //ignore points we have no data for

        //Add the log probabilities
        fLogAlphaProb += fPDFAlphaLog[pos];
        fLogBetaProb += fPDFBetaLog[pos];
    }

    //Store & return results
    fClassifierResult.SetClassification("likelihood", fLogAlphaProb-fLogBetaProb);
    fClassifierResult.SetClassification("retrigger", fRetrigger ? 1.0 : 0.0);
    fClassifierResult.SetValid(true);
    return fClassifierResult;
}