#include <RAT/Ext0NuCosTheta.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;


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

void Ext0NuCosTheta::Initialise(const std::string& param) {
    fParam = param;
}

void Ext0NuCosTheta::BeginOfRun(DS::Run&) {
    string tableName;
    string materialIndex;

    if (fParam.empty()) {
    tableName = "0NUEXT_COSTHETA_Tl208AV"; //default background
    try {
      materialIndex = DB::Get()->GetLink("GEO","inner_av")->GetS("material");
            DB::Get()->GetLink(tableName, materialIndex)->GetD("cos_theta_first");
    } catch (...) {
      materialIndex = ""; //fallback scintillator
    }
  }else {
    const size_t hyphenPos = fParam.find("-");
    if (hyphenPos != string::npos) {
      tableName = "0NUEXT_COSTHETA_"+ fParam.substr(0, hyphenPos);
      materialIndex = fParam.substr(hyphenPos+1);
    } else {
      tableName = "0NUEXT_COSTHETA_"+fParam;
      try {
        materialIndex = DB::Get()->GetLink("GEO","inner_av")->GetS("material");
                DB::Get()->GetLink(tableName, materialIndex)->GetD("cos_theta_first");
      } catch (...) {
        materialIndex = ""; //fallback scintillator
      }
    }
  }
    info << "Ext0NuCosTheta: Using " << tableName << "[" << materialIndex << "] weights" << endl;
    DBLinkPtr pdftable = DB::Get()->GetLink(tableName, materialIndex);

    fCosThetaFirst = pdftable->GetD("cos_theta_first");
    fCosThetaLast = pdftable->GetD("cos_theta_last");
    fTimeWindow = pdftable->GetDArray("time_window");
    fNormalisationFlag = pdftable->GetI("normalisation_flag");

    fWeights = pdftable->GetDArray("weights");
    fCosThetaStep = (fCosThetaLast-fCosThetaFirst)/fWeights.size();
}


void Ext0NuCosTheta::EndOfRun(DS::Run&) {
}

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

void Ext0NuCosTheta::SetSeed(const DS::FitResult& seed) {
    fEventPos = seed.GetVertex(0).GetPosition();
    fEventTime = seed.GetVertex(0).GetTime();
}

DS::ClassifierResult Ext0NuCosTheta::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 discriminant and hit count in timing window
    fDiscriminant = 0.0;
    size_t earlyHitCount = 0;

  //Loop over selected PMTs
    for (vector<FitterPMT>::const_iterator pmt = fPMTData.begin(); pmt != fPMTData.end(); ++pmt) {
        // Calculate cos theta
        const TVector3 pmtPos = pmtInfo.GetPosition(pmt->GetID());
        const double cosTheta = fEventPos.Dot(pmtPos)/(pmtPos.Mag()*fEventPos.Mag());

        lightPathCalc.CalcByPosition(fEventPos,pmtPos);
        const double transitTime = effectiveVelocity.CalcByDistance(lightPathCalc.GetDistInInnerAV(),lightPathCalc.GetDistInAV(),lightPathCalc.GetDistInWater());
        const double timeResidual = pmt->GetTime() - transitTime - fEventTime;

        //Check if in promt peak window and find the bin if so
        if ((timeResidual < fTimeWindow.at(0)) || (timeResidual > fTimeWindow.at(1))) continue;
        size_t bin = static_cast<size_t>(((cosTheta-fCosThetaFirst)/fCosThetaStep));
        earlyHitCount++;

        //Add the difference in logs
        fDiscriminant += fWeights[bin];
    }

    //Early time window empty -> exit
    if (!earlyHitCount)
        return fClassifierResult;

    //Normalise to one hit
    if (fNormalisationFlag)
        fDiscriminant /= static_cast<double>(earlyHitCount);

    //Store & return results
    fClassifierResult.SetClassification("discriminant", fDiscriminant);
    fClassifierResult.SetValid(true);
    return fClassifierResult;
}