#include <RAT/AlphaBetaClassifier.hh>
#include <RAT/DB.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DU/Utility.hh>
#include <math.h>
#include <RAT/Optimiser.hh>
using namespace RAT;
using namespace RAT::Classifiers;
using namespace RAT::DS;
using namespace std;

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

void
AlphaBetaClassifier::BeginOfRun( DS::Run& )
{
  string dbFileName;
  string materialIndex;
  DB *db= DB::Get();
  string defaultIndex = "te_diol_0p5_labppo_scintillator_bisMSB_Jan2016";

  if ( fIndex.empty() ){
    materialIndex = db->GetLink("GEO", "inner_av")->GetS("material");
    dbFileName = "LIKELIHOOD_Bi212_wPSD";
  }
  else
  {
    const size_t hyphenPos = fIndex.find("-");
    if( hyphenPos != string::npos )
    {
      dbFileName = "LIKELIHOOD_Bi"+fIndex.substr(0,hyphenPos);
      materialIndex = fIndex.substr(hyphenPos+1);
    }
    else{
      dbFileName = "LIKELIHOOD_Bi"+fIndex;
      materialIndex = db->GetLink("GEO", "inner_av")->GetS("material");
    }
  }

  DBLinkPtr dbPDFLink = db->GetLink(dbFileName,defaultIndex);

  DBLinkGroup grp = db->GetLinkGroup(dbFileName);
  DBLinkGroup::iterator it;
  for(it=grp.begin();it!=grp.end();++it)
  {
    if(materialIndex == it->first)
    {
      dbPDFLink = db->GetLink(dbFileName,materialIndex);
    }
  }

  fTimes = dbPDFLink->GetDArray("times");
  fBetaProbability = dbPDFLink->GetDArray("beta_probability");
  fAlphaProbability = dbPDFLink->GetDArray("alpha_probability");
  fTwoBetaProbability = dbPDFLink->GetDArray("two_beta_probability");
  fEnergies = dbPDFLink->GetDArray("energy_ratio");
  fEnergyProb = dbPDFLink->GetDArray("energy_probability");
  fParams.push_back(200.0);
  fPositiveErrors.push_back(200.0);
  fNegativeErrors.push_back(200.0);

  fTimeStep = fTimes[1] - fTimes[0];
  fEnergyStep = fEnergies[1] - fEnergies[0];

  fSummedAlphaProbability.resize( fTimes.size(), 0 );
  fSummedBetaProbability.resize( fTimes.size(), 0 );
  fSummedTwoBetaProbability.resize( fTimes.size(), 0 );

  fSummedBetaProbability[fTimes.size()-1] = fBetaProbability[fTimes.size()-1];
  fSummedTwoBetaProbability[fTimes.size()-1] = fTwoBetaProbability[fTimes.size()-1];
  fSummedAlphaProbability[fTimes.size()-1] = fAlphaProbability[fTimes.size()-1];

  for( int i = static_cast<int>(fTimes.size() - 2); i > -1; i-- )
    {
      fSummedBetaProbability[i] = fSummedBetaProbability[i+1] + fBetaProbability[i];
      fSummedTwoBetaProbability[i] = fSummedTwoBetaProbability[i+1] + fTwoBetaProbability[i];
      fSummedAlphaProbability[i] = fSummedAlphaProbability[i+1] + fAlphaProbability[i];
    }
  fMinValue = 1e10; // to ensure a minimum is found
  for( unsigned int i = 0; i < fTimes.size(); i++ )
    {
      if( fBetaProbability[i] < fMinValue && fBetaProbability[i] > 0 )
        fMinValue = fBetaProbability[i];
      if( fTwoBetaProbability[i] < fMinValue && fTwoBetaProbability[i] > 0 )
        fMinValue = fTwoBetaProbability[i];
    }
  // fMinValue is a global penalty term, used where PDF has value of 0 (to avoid log(0))
  // By dividing by 10 we ensure a smaller LogL for the penalty term.
  fMinValue /= 10.0;
  fLightPath = DU::Utility::Get()->GetLightPathCalculator();
}
void
AlphaBetaClassifier::DefaultSeed()
{
  fEventPos =TVector3();
  fEventTime =0.0;
}
void
AlphaBetaClassifier::SetSeed(const DS::FitResult& seed)
{
  FitVertex seedVertex;
  try {seedVertex = seed.GetVertex(0);}
  catch( FitResult::NoVertexError& error){return;}
  try {fEventPos = seedVertex.GetPosition();}
  catch( FitVertex::NoValueError& error) {return;}
  try{fEventTime=seedVertex.GetTime();}
  catch( FitVertex::NoValueError& error) {return;}
}

std::vector<double>
AlphaBetaClassifier::GetParams()
const {
  return fParams;
}
std::vector<double>
AlphaBetaClassifier::GetPositiveErrors()
const {
  return fPositiveErrors;
}

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

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

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

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

  fOptimiser->SetComponent( this );
  const DU::PMTInfo& pmtInfo = DU::Utility::Get()->GetPMTInfo();
  for ( vector<FitterPMT>::const_iterator iPMT = fPMTData.begin(); iPMT != fPMTData.end(); ++iPMT )
  {
    const TVector3 pmtPos = pmtInfo.GetPosition(iPMT->GetID());
    fLightPath.CalcByPosition(fEventPos,pmtPos);
    double distInInnerAV = fLightPath.GetDistInInnerAV();
    double distInAV = fLightPath.GetDistInAV();
    double distInWater = fLightPath.GetDistInWater();
    const double transitTime = DU::Utility::Get()->GetEffectiveVelocity().CalcByDistance(distInInnerAV,distInAV,distInWater);
    const double corTime = iPMT->GetTime() - transitTime - fEventTime;
    fTimeResiduals.push_back(corTime);
  }
  fParams[0] = 150;
  fFitHypothesis=false;
  const double logLTe =  fOptimiser->Maximise();
  fParams = fOptimiser->GetParams();
  for(unsigned int i =0; i< fTimeResiduals.size();++i)
  {
    fTimeResiduals[i]-=fParams[0];
  }
  fParams[0]=200.0;
  fParams.push_back(1.0);//Add energy parameters
  fPositiveErrors.push_back(1.0);
  fNegativeErrors.push_back(0.8);

  fFitHypothesis=true;
  const double logLBiPo = fOptimiser->Maximise();
  SetFOM("LogLBiPo",logLBiPo);
  SetFOM("LogLTe",logLTe);
  SetFOM("DecayTime",(fOptimiser->GetParams())[0]);
  fClassifierResult.SetClassification( "AlphaBetaClassifier",  logLTe-logLBiPo );

  fTimeResiduals.clear();//Ensures that the vector of times doesn't get larger with each new event
  fParams.pop_back();
  fPositiveErrors.pop_back();
  fNegativeErrors.pop_back();
  fClassifierResult.SetValid( true );
  return fClassifierResult;
}
double
AlphaBetaClassifier::operator() (const std::vector<double>& params)
{

  double fitLikelihood =0.0;

  // Which beta PDF is used depends on the hypothesis
  int nbins = fTimes.size();
  std::vector<double>* betaPDF;
  std::vector<double>* summedBetaPDF;
  // Use round here to ensure that the time closest to 0 is found
  int zeroIndex = round( ( 0 - fTimes[0] ) / fTimeStep );
  double normFactor=0;
  double globalOffset=0;
  double offset=0;
  double sizeRatio = 1.0; //Ratio of Beta peak size to Alpha peak size.

  if (fFitHypothesis)
  {
    betaPDF = &fBetaProbability;
    summedBetaPDF = &fSummedBetaProbability;
    offset = params[0];
    sizeRatio = params[1];
  }
  else
  {
    betaPDF = &fTwoBetaProbability;
    summedBetaPDF = &fSummedTwoBetaProbability;
    globalOffset = params[0];
  }

  // Truncate / floor here to act like bin finding (where bin spans fTimes[i] to fTimes[i+1])
  // The goal of the index offset variables is to represent how many bins a given change in time will be.
  int offsetIndex = static_cast<int>( ( offset - fTimes[0] ) / fTimeStep );
  int globalOffsetIndex = static_cast<int>( ( globalOffset - fTimes[0] ) / fTimeStep );
  globalOffsetIndex -= zeroIndex;
  offsetIndex -= zeroIndex;

  double lastTime = fTimes[ nbins - 1 ];
  double firstTime = fTimes[ 0 ];

  // Each hit time may have a likelihood from the PDF taken from 0
  // to max - offsetIndex.
  // Take a normalisation factor as the summed PDF over the same
  // time range.
  int lastTimeIndex1 = nbins - 1 - globalOffsetIndex;
  if( globalOffsetIndex < 0 )
    lastTimeIndex1 = nbins - 1;
  int lastTimeIndex2 = nbins-1 - globalOffsetIndex - offsetIndex;
  if( globalOffsetIndex + offsetIndex < 0 )
    lastTimeIndex2 = nbins - 1;

  // Normalisation factor for either hypothesis
  normFactor += sizeRatio * ( (*summedBetaPDF)[ 1 ] -
                              (*summedBetaPDF)[ lastTimeIndex1 ] );
  // Additional factor for the alpha-beta hypothesis only
  if( fFitHypothesis && ( lastTime - globalOffset - offset ) > firstTime )
    normFactor += sizeRatio * ( fSummedAlphaProbability[ 1 ] -
                                fSummedAlphaProbability[ lastTimeIndex2 ] );

  for(unsigned int i=0;i<fTimeResiduals.size();++i)
    {
      double prob = 0;
      int timeIndex = static_cast<int>( ( fTimeResiduals[i] - fTimes[0] ) / fTimeStep );

      // Probability for either hypothesis
      if( timeIndex > globalOffsetIndex && ( timeIndex - globalOffsetIndex ) < nbins )
        prob += ( sizeRatio * (*betaPDF)[ timeIndex - globalOffsetIndex ] );

      // Additional factor for the alpha-beta hypothesis only
      if( fFitHypothesis && timeIndex > ( globalOffsetIndex + offsetIndex )
          && ( timeIndex - globalOffsetIndex - offsetIndex ) < nbins )
        prob += ( sizeRatio * fAlphaProbability[ timeIndex - globalOffsetIndex - offsetIndex ] );
      prob /= normFactor;

      if (prob==0)
        {
          prob = fMinValue;
        }
      fitLikelihood += log(prob);
    }
  if(fFitHypothesis)
    {
      // Additional energy term in the likelihood for alpha-beta hypothesis
      double energyIndex = static_cast<int>( ( params[1] - fEnergies[0] ) / fEnergyStep );
      // Temporary hack to avoid crash  -- WJH   27 May 2015
//       fitLikelihood += log(fEnergyProb[ energyIndex ]);
      if( fEnergyProb[energyIndex] >0. ) {
        fitLikelihood += log(fEnergyProb[ energyIndex ]);
      }else{
        fitLikelihood += log(1.e-6);
      }
      // End of hack
    }

  return fitLikelihood;
}