#include <RAT/BiPoLikelihoodDiff.hh>
#include <RAT/DB.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DU/LightPathCalculator.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/Optimiser.hh>
#include <RAT/Log.hh>
using namespace RAT;
using namespace RAT::Classifiers;
using namespace RAT::DS;

#include <TH1.h>
using namespace ROOT;

#include <string>
#include <vector>
#include <algorithm>
using namespace std;

#include <G4PhysicsOrderedFreeVector.hh>


void BiPoLikelihoodDiff::Initialise(const std::string& param)
{
  // This index should be of the form <likelihood>-<material>
  // where the second portion (material) is optional; if not used
  // then the current material in the inner AV will be used if possible
  // or a default material if there is no PDF for the current material
  fIndex = param;
}

void BiPoLikelihoodDiff::BeginOfRun(DS::Run&)
{
  // Get entry for the given material/isotope combination index
  DB* db = DB::Get();
  string materialIndex;
  string usedIndex;
  string defaultIndex = "te_diol_dda_0p5_labppo_scintillator_bisMSB_Jan2018_214BiPo";
  string biPoIndex;
  if( fIndex.empty() == true )
    {
      try{
        materialIndex = db->GetLink("GEO", "inner_av")->GetS("material");
      } catch(DBNotFoundError & ) {
        // If we are in a partial-fill geometry, use the upper
        // material
        try{
          materialIndex = db->GetLink("GEO", "inner_av")->GetS("material_top");
        } catch(DBNotFoundError &){
          Log::Die("BiPoLikelihoodDiff: 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\"");
        }
      }
      biPoIndex = "214";
    }
  else
    {
      // If there is a hyphen, then get both material and pdf type (212 or 214)
      // if there is no hyphen then get only the pdf type
      const size_t hyphenPos = fIndex.find("-");
      if( hyphenPos != string::npos )
        {
          materialIndex = fIndex.substr( 0, hyphenPos );
          biPoIndex = fIndex.substr( hyphenPos+1 );
        }
      else
        {
          try{
            materialIndex = db->GetLink("GEO", "inner_av")->GetS("material");
          } catch(DBNotFoundError &){
            // If we are in a partial-fill geometry, use the upper
            // material
            try{
              materialIndex = db->GetLink("GEO", "inner_av")->GetS("material_top");
            } catch(DBNotFoundError &){
              Log::Die("BiPoLikelihoodDiff: 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\"");
            }
          }
          biPoIndex = fIndex;
        }
    }
  // Now construct the table index from BiPoType
  string index = materialIndex + "_" + biPoIndex + "BiPo";

  DBLinkPtr dbLink = db->GetLink("CLASSIFIER_BIPO_LIKELIHOODDIFF", defaultIndex);
  DBLinkGroup grp = db->GetLinkGroup("CLASSIFIER_BIPO_LIKELIHOODDIFF");
  DBLinkGroup::iterator it;
  for(it=grp.begin();it!=grp.end();++it)
    {
      if(index == it->first)
        {
          dbLink = db->GetLink("CLASSIFIER_BIPO_LIKELIHOODDIFF", index);
          usedIndex = index;
        }
    }
  if( usedIndex != index )
    warn << "BiPoLikelihoodDiff::No PDF available for " << index <<  ", using the default entry ("
         << defaultIndex << ")" << newline;

  fTimes = dbLink->GetDArray("times");
  fProbabilityTe = dbLink->GetDArray("pdf_Te");
  fProbabilityBi = dbLink->GetDArray("pdf_Bi");
  fProbabilityPo = dbLink->GetDArray("pdf_Po");
  fMeanPoAlphaNhits = dbLink->GetD("meanPoAlphaNhits");
  fTriggerWindowStart = dbLink->GetD("triggerWindowStartTime");
  fTriggerWindowEnd = dbLink->GetD("triggerWindowEndTime");

  fBasePDFTe = new G4PhysicsOrderedFreeVector( &(fTimes[0]), &(fProbabilityTe[0]), fTimes.size() );
  fBasePDFBi = new G4PhysicsOrderedFreeVector( &(fTimes[0]), &(fProbabilityBi[0]), fTimes.size() );
  fBasePDFPo = new G4PhysicsOrderedFreeVector( &(fTimes[0]), &(fProbabilityPo[0]), fTimes.size() );

  fLightPath = DU::Utility::Get()->GetLightPathCalculator();
}

void BiPoLikelihoodDiff::EndOfRun(DS::Run&)
{
  // Clear all vectors and delete all pointers
  fTimes.clear();
  fProbabilityTe.clear();
  fProbabilityBi.clear();
  fProbabilityPo.clear();
  delete fBasePDFTe;
  delete fBasePDFBi;
  delete fBasePDFPo;
}

void BiPoLikelihoodDiff::DefaultSeed()
{
  fEventPosition = TVector3();
  fEventTime = 0.0;
}

void BiPoLikelihoodDiff::SetSeed(const DS::FitResult& seed)
{
  FitVertex seedVertex;

  try {seedVertex = seed.GetVertex(0);}
  catch( FitResult::NoVertexError& error ) {return;}

  try {fEventPosition = seedVertex.GetPosition();}
  catch( FitVertex::NoValueError& error ) {return;}
  try {fEventTime = seedVertex.GetTime();}
  catch( FitVertex::NoValueError& error ) {return;}
}

std::vector<double> BiPoLikelihoodDiff::GetParams() const
{
  return fParams;
}

std::vector<double> BiPoLikelihoodDiff::GetPositiveErrors() const
{
  return fPositiveErrors;
}

std::vector<double> BiPoLikelihoodDiff::GetNegativeErrors() const
{
  return fNegativeErrors;
}

void BiPoLikelihoodDiff::SetParams(const std::vector<double>& params)
{
  fParams = params;
}

void BiPoLikelihoodDiff::SetPositiveErrors(const std::vector<double>& errors)
{
  fPositiveErrors = errors;
}

void BiPoLikelihoodDiff::SetNegativeErrors(const std::vector<double>& errors)
{
  fNegativeErrors = errors;
}

double BiPoLikelihoodDiff::operator() (const std::vector<double>& params)
{
  // For the current trial value of deltaTime ...
  double deltaTime = params[0];

  //  ... find the expected fraction of Po-Nhits up to this value of deltaTime ...
  double inWindowPoAlphaNhitsFraction = 0.0;
  for (unsigned int b = 0; b <= fProbabilityPo.size(); b++)
  {
      if ((fTriggerWindowEnd - deltaTime) > fProbabilityPo[b]) {inWindowPoAlphaNhitsFraction += fProbabilityPo[b];}
  }
  double inWindowPoAlphaNhits = fMeanPoAlphaNhits * inWindowPoAlphaNhitsFraction;

  double currentLogLikelihood = 0.0;

  // ... and for each offset Time Residual ...
  for (unsigned int t = 0; t < fTimeResiduals.size(); t++)
  {
      // ... find the probability of getting that Time Residual value from the Bi PDF ...
      double probBi = fBasePDFBi->Value(fTimeResiduals[t]);
      if (probBi == 0.0) {continue;}

      // ... find the probability of getting that Time Residual value offset by this value of deltaTime from the Po PDF ...
      double probPo = fBasePDFPo->Value(fTimeResiduals[t] - deltaTime);
      if (probPo == 0.0) {continue;}

      // ... find the corresponding single PMT log-likelihood value and add to the summation
      double singlePMTValue = log(((fNHitsWithinWindow - inWindowPoAlphaNhits) * probBi) + (inWindowPoAlphaNhits * probPo));
      currentLogLikelihood += singlePMTValue;
  }

  return currentLogLikelihood;
}

double BiPoLikelihoodDiff::logLikelihoodTe()
{
  double logLikelihood = 0.0;

  // For each offset Time Residual ...
  for (unsigned int t = 0; t < fTimeResiduals.size(); t++)
  {
      // ... find the probability of getting that Time Residual value from the 130Te PDF ...
      double probTe = fBasePDFTe->Value(fTimeResiduals[t]);
      if (probTe == 0.0) {continue;}

      // ... find the corresponding single PMT log-likelihood value and add to the summation
      double singlePMTValue = log(fNHitsWithinWindow * probTe);
      logLikelihood += singlePMTValue;
  }

  return logLikelihood;
}

DS::ClassifierResult BiPoLikelihoodDiff::GetClassification()
{
  fClassifierResult.Reset();
  if( fPMTData.empty() )
    return fClassifierResult;

  // Set the starting value and the range over which to optimise: between (fParams[i] - fNegativeErrors[i]) and (fParams[i] + fPositiveErrors[i]) for each dimension i to be optimised
  // The number of divisions in this range is set in the macro by using:  /rat/procset optimiser "grid-X"  ... where X refers to X divisions in the range between fNegativeErrors[i] and fPositiveErrors[i]
  fParams.push_back(0.0);
  fPositiveErrors.push_back(fTriggerWindowEnd);
  fNegativeErrors.push_back(0.0);

  fOptimiser->SetComponent(this);

  // Store the Time Residuals for later use, and count the number of PMTs in the event with Time Residual within the Trigger Window
  const int eventNhits = fPMTData.size();

  fNHitsWithinWindow = 0.0;
  for (int j = 0; j < eventNhits; j++)
  {
      const TVector3 pmtPosition = DU::Utility::Get()->GetPMTInfo().GetPosition(fPMTData[j].GetID());
      const double pmtTime = fPMTData[j].GetTime();

      DU::LightPathCalculator lightPath =  DU::Utility::Get()->GetLightPathCalculator();
      lightPath.CalcByPosition(fEventPosition, pmtPosition);
      double distInInnerAV = lightPath.GetDistInInnerAV();
      double distInAV = lightPath.GetDistInAV();
      double distInWater = lightPath.GetDistInWater();

      const double flightTime = DU::Utility::Get()->GetEffectiveVelocity().CalcByDistance(distInInnerAV, distInAV, distInWater);
      const double timeResid = pmtTime - flightTime - fEventTime;

      fTimeResiduals.push_back(timeResid);
      if ((timeResid >= fTriggerWindowStart) && (timeResid < fTriggerWindowEnd)) {fNHitsWithinWindow += 1.0;}
  }

  // Find the 10th earliest Time Residual and offset all the Time Residuals by this 10th value - do this now to save having to do it multiple times later on
  sort(fTimeResiduals.begin(), fTimeResiduals.end());
  double offsetTime = fTimeResiduals[9];
  for (unsigned int t = 0; t < fTimeResiduals.size(); t++) { fTimeResiduals[t] -= offsetTime; }

  // Calculate the Log-Likelihood compared to the Te-PDF
  double lglklhdTe = logLikelihoodTe();
  SetFOM("lglklhdTe", lglklhdTe);

  // Calculate the Maximum BiPo Log-Likelihood value using the Optimiser
  const double lglklhdBiPo = fOptimiser->Maximise();
  SetFOM("lglklhdBiPo", lglklhdBiPo);

  // Save the optimised value of deltaTime and the optimiser errors, and then clear the optimiser parameters ready for the next event
  SetParams(fOptimiser->GetParams());
  SetFOM("optmDelTime", fParams[0]);
  SetPositiveErrors(fOptimiser->GetPositiveErrors());
  SetNegativeErrors(fOptimiser->GetNegativeErrors());

  fParams.clear();
  fPositiveErrors.clear();
  fNegativeErrors.clear();
  fTimeResiduals.clear();

  // The Classifier Result is the Difference in the Log-Likelihoods
  fClassifierResult.SetClassification("lglklhdDiff", (lglklhdBiPo - lglklhdTe));
  fClassifierResult.SetValid(true);

  return fClassifierResult;
}