// -------------------------------------------------- //
// EnergyFunctional.cc
// Contact Person: Krish Majumdar <Krishanu.Majumdar@physics.ox.ac.uk>
// See EnergyFunctional.hh and DocDB-2782 for more details
// -------------------------------------------------- //

#include <RAT/EnergyFunctional.hh>
#include <RAT/DB.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/Log.hh>
using namespace RAT;
using namespace RAT::Methods;
using namespace RAT::DS;

#include <string>
#include <vector>
#include <math.h>
using namespace std;

#include <TVector3.h>


void EnergyFunctional::Initialise(const std::string& param)
{
  fIndex = param;
}

void EnergyFunctional::BeginOfRun(DS::Run&)
{
  DB * const db = DB::Get();

  fMaterial = fIndex;
  if (fIndex.empty())
  {
    try
    {
      fMaterial = db->GetLink("GEO", "inner_av")->GetS("material");
    }
    catch( DBNotFoundError& e )
    {
      try
      {
        // If we are in a partial-fill geometry, use the upper material
        fMaterial = db->GetLink("GEO", "inner_av")->GetS("material_top");
      }
      catch ( DBNotFoundError & )
      {
        Log::Die("EnergyFunctional::BeginOfRun: inner_av material and material_top undefined.  Try something like: \n"
        "/rat/db/set GEO[inner_av] material \"labppo_scintillator\"  or \n"
        "/rat/db/set GEO[inner_av] material_top \"labppo_scintillator\"");
      }
    }
  }

  // Default parameters are those of "labppo_scintillator"
  DBLinkPtr dbNAVLink = db->GetLink("FIT_ENERGY_FUNCTIONAL", "labppo_scintillator");
  DBLinkGroup grp = db->GetLinkGroup("FIT_ENERGY_FUNCTIONAL");

  // Check for the material-specific entry and overwrite the default values only if found
  for (DBLinkGroup::iterator it = grp.begin(); it != grp.end(); ++it)
  {
    if (fMaterial == it->first)
    {
      dbNAVLink = db->GetLink("FIT_ENERGY_FUNCTIONAL", fMaterial);
      break;
    }
  }

  fEnergyCoefficients = dbNAVLink->GetDArray("energyCoeffs");
  if (fEnergyCoefficients.size() != 4)
  {
    Log::Die("EnergyFunctional::BeginOfRun: the \"energyCoeffs\" vector does not contain 4 parameters - please check the .ratdb file");
  }
  fRadiusCoefficients_Internal_Low = dbNAVLink->GetDArray("radiusCoeffsInternalLow");
  fRadiusCutoff_Internal_Mid = dbNAVLink->GetD("internalMidRadiusCut");
  fRadiusCoefficients_Internal_Mid = dbNAVLink->GetDArray("radiusCoeffsInternalMid");
  fRadiusCutoff_Internal_High = dbNAVLink->GetD("internalHighRadiusCut");
  fRadiusCoefficients_Internal_High = dbNAVLink->GetDArray("radiusCoeffsInternalHigh");
  fRadiusCutoff_External_Low = dbNAVLink->GetD("externalLowRadiusCut");
  fRadiusCoefficients_External_Low = dbNAVLink->GetDArray("radiusCoeffsExternalLow");
  fRadiusCutoff_External_High = dbNAVLink->GetD("externalHighRadiusCut");
  fRadiusCoefficients_External_High = dbNAVLink->GetDArray("radiusCoeffsExternalHigh");
  fZCoefficients = dbNAVLink->GetDArray("zCoeffs");

  fAVRadius = db->GetLink("SOLID", "acrylic_vessel_outer")->GetD("r_sphere");
}

void EnergyFunctional::EndOfRun(DS::Run&)
{
  fEnergyCoefficients.clear();
  fRadiusCoefficients_Internal_Low.clear();
  fRadiusCoefficients_Internal_Mid.clear();
  fRadiusCoefficients_Internal_High.clear();
  fRadiusCoefficients_External_Low.clear();
  fRadiusCoefficients_External_High.clear();
  fZCoefficients.clear();
}

void EnergyFunctional::DefaultSeed()
{
  FitVertex seedVertex;
  // In the following line, the booleans indicate: (1) the validity of the position, (2) if the position was calculated here (false) or passed from a seed (true)
  seedVertex.SetPosition( TVector3(0.0, 0.0, 0.0), false, true);
  fSeedResult.SetVertex( 0, seedVertex );
}

DS::FitResult EnergyFunctional::GetBestFit()
{
  fFitResult.Reset();
  if( fPMTData.empty() )
    return fFitResult;
  CopySeedToResult();

  FitVertex fitVertex = fSeedResult.GetVertex( 0 );
  double eventRadius = fitVertex.GetPosition().Mag();
  double eventZ = fitVertex.GetPosition().Z();
  const int eventNhits = fPMTData.size();

  // Get the Segmentor information:
  // 1) a vector of integers denoting in which (theta, phi) segment each raw PMT is in
  // 2) a vector of integers denoting the raw PMT population of each segment
  // 3) a vector of integers denoting the triggered PMT population of each segment
  DU::Segmentor segmentor = DU::Utility::Get()->GetSegmentor();
  vector<unsigned int> rawPMTSegmentIDs = segmentor.GetSegmentIDs();
  vector<unsigned int> rawPMTSegmentPopulations = segmentor.GetSegmentPopulations();
  vector<unsigned int> hitPMTSegmentPopulations (rawPMTSegmentPopulations.size(), 0);
  for (int j = 0; j < eventNhits; j++)
  {
    hitPMTSegmentPopulations[rawPMTSegmentIDs[fPMTData[j].GetID()]]++;
  }

  double hParameter = 0.0;
  for (unsigned int s = 0; s < rawPMTSegmentPopulations.size(); s++)
  {
    double numberOfRawPMTs = static_cast<double>(rawPMTSegmentPopulations[s]);
    double numberOfHitPMTs = static_cast<double>(hitPMTSegmentPopulations[s]);

    if (numberOfHitPMTs > numberOfRawPMTs)
    {
      Log::Die("EnergyFunctional::GetBestFit: (fatal error) encountered segment with more hit PMTs than actual PMTs - signifies a deeper problem ... exiting energy fitter");
    }
    else if (numberOfHitPMTs == numberOfRawPMTs)
    {
      warn << "EnergyFunctional::GetBestFit: (warning) encountered saturated segment with the same number of hit PMTs as actual PMTs ... correcting and continuing \n";
      numberOfHitPMTs = numberOfRawPMTs - 1;
    }

    hParameter -= (numberOfRawPMTs * log(1.0 - (numberOfHitPMTs / numberOfRawPMTs)));
  }

  // Calculate the value of the radius function, choosing the specific set of function coefficients based on the reconstructed event radius
  double functionOfRadius;
  if (eventRadius < fRadiusCutoff_Internal_Mid)
  {
    functionOfRadius = EvaluatePolynomial(fRadiusCoefficients_Internal_Low, eventRadius);
  }
  else if ((eventRadius >= fRadiusCutoff_Internal_Mid) && (eventRadius < fRadiusCutoff_Internal_High))
  {
    functionOfRadius = EvaluatePolynomial(fRadiusCoefficients_Internal_Mid, eventRadius);
  }
  else if ((eventRadius >= fRadiusCutoff_Internal_High) && (eventRadius < fRadiusCutoff_External_Low))
  {
    functionOfRadius = EvaluatePolynomial(fRadiusCoefficients_Internal_High, eventRadius);
  }
  else if ((eventRadius >= fRadiusCutoff_External_Low) && (eventRadius < fRadiusCutoff_External_High))
  {
    functionOfRadius = EvaluatePolynomial(fRadiusCoefficients_External_Low, eventRadius);
  }
  else
  {
    functionOfRadius = EvaluatePolynomial(fRadiusCoefficients_External_High, eventRadius);
  }

  // Calculate the value of the z function
  double functionOfZ = EvaluatePolynomial(fZCoefficients, eventZ);

  // Set the energy coefficients, based on whether the event's radius puts it as an internal or external event
  double energyCoeff0, energyCoeff1;
  if (eventRadius < fRadiusCutoff_External_Low)
  {
    energyCoeff0 = fEnergyCoefficients[0];
    energyCoeff1 = fEnergyCoefficients[1];
  }
  else
  {
    energyCoeff0 = fEnergyCoefficients[2];
    energyCoeff1 = fEnergyCoefficients[3];
  }

  // Calculate the final fitted energy
  double energy = (1.0 / energyCoeff1) * ((hParameter / (functionOfRadius * functionOfZ)) - energyCoeff0);

  // If the event has reconstructed outside the AV and the material being used is NOT lightwater_sno, mark the fit validity as false but set the fit value anyway
  if ((eventRadius > fAVRadius) && (fMaterial != "lightwater_sno")) {fitVertex.SetEnergy(energy, false);}
  else {fitVertex.SetEnergy(energy, true);}
  fitVertex.SetEnergyErrors(1.0);  // Need to explore this energy error a bit more ...
  fFitResult.SetVertex(0, fitVertex);
  return fFitResult;
}


double EnergyFunctional::EvaluatePolynomial(std::vector<double> coefficients, double parameter)
{
  double value = 0.0;

  for (unsigned int p = 0; p < coefficients.size(); p++)
  {
    value += (coefficients[p] * std::pow(parameter, (int)p));
  }

  return value;
}