#include <RAT/DirectionLikelihood.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/Optimiser.hh>
#include <RAT/PMTSelector.hh>
#include <RAT/PDF.hh>
#include <RAT/DB.hh>
#include <RAT/Log.hh>
#include <iostream>
#include <cmath>

using namespace RAT;
using namespace RAT::Methods;
using namespace RAT::DS;

void
DirectionLikelihood::DefaultSeed()
{
  FitVertex vertex;
  //Initialise the FitResult
  vertex.SetPosition( TVector3( 0.0, 0.0, 0.0), false, true);
  vertex.SetDirection( TVector3( 1.0, 0.0, 0.0 ), false, true );
  vertex.SetDirectionErrors( TVector3( 1.0, 1.0, 1.0 ), false, true );
  fSeedResult.SetVertex( 0, vertex );
}

FitResult
DirectionLikelihood::GetBestFit()
{
  fFitResult.Reset();
  if( fPMTData.empty() )
    throw MethodFailError( "DirectionLikelihood: No hits to fit" );
  CopySeedToResult();

  SelectPMTData( fFitResult.GetVertex( 0 ) );
  fFitResult.SetFOM( "SelectedNHit", static_cast<double>( fSelectedPMTData.size() ) );
  if( fSelectedPMTData.empty() )
    throw MethodFailError( "DirectionLikelihood: No hits in selected PMT data" );

  fOptimiser->SetComponent( this );
  // Call the designated optimiser to fit and save the fom
  const double logL = fOptimiser->Maximise();
  fFitResult.SetFOM( "LogL", logL );
  // Now save the optimised values
  SetParams( fOptimiser->GetParams() );
  SetPositiveErrors( fOptimiser->GetPositiveErrors() );
  SetPositiveErrors( fOptimiser->GetNegativeErrors() );

  // The direction should be a unit vector and the errors
  // should be scaled accordingly.
  // This is not performed in SetParams as the optimiser adjusts and
  // uses via that function; the method should not force the optimisers decisions.
  FitVertex vertex = fFitResult.GetVertex( 0 );
  double scaling = vertex.GetDirection().Mag();
  vertex.SetDirection( vertex.GetDirection().Unit(), fOptimiser->GetValid() );
  vertex.SetPositiveDirectionError( TVector3( vertex.GetPositiveDirectionError().X() / scaling,
                                              vertex.GetPositiveDirectionError().Y() / scaling,
                                              vertex.GetPositiveDirectionError().Z() / scaling ),
                                    fOptimiser->GetValid() );
  vertex.SetNegativeDirectionError( TVector3( vertex.GetNegativeDirectionError().X() / scaling,
                                              vertex.GetNegativeDirectionError().Y() / scaling,
                                              vertex.GetNegativeDirectionError().Z() / scaling ),
                                    fOptimiser->GetValid() );
  fFitResult.SetVertex( 0, vertex );

  return fFitResult;
}

double
DirectionLikelihood::operator()( const std::vector<double>& params )
{
  SetParams( params );
  double logLike = 0.0;
  // loglike sum from pdf
  for( std::vector<FitterPMT>::const_iterator iPMT = fSelectedPMTData.begin(); iPMT != fSelectedPMTData.end(); ++iPMT )
    logLike += log( fPDF->GetProbability( *iPMT, fFitResult.GetVertex(0) ) );

  return logLike;
}

std::vector<double>
DirectionLikelihood::GetParams() const
{
  std::vector<double> params;
  FitVertex vertex = fFitResult.GetVertex(0);
  params.push_back( vertex.GetDirection().x() );
  params.push_back( vertex.GetDirection().y() );
  params.push_back( vertex.GetDirection().z() );
  return params;
}

std::vector<double>
DirectionLikelihood::GetPositiveErrors() const
{
  std::vector<double> errors;
  FitVertex vertex = fFitResult.GetVertex(0);
  errors.push_back( vertex.GetPositiveDirectionError().x() );
  errors.push_back( vertex.GetPositiveDirectionError().y() );
  errors.push_back( vertex.GetPositiveDirectionError().z() );
  return errors;
}

std::vector<double>
DirectionLikelihood::GetNegativeErrors() const
{
  std::vector<double> errors;
  FitVertex vertex = fFitResult.GetVertex(0);
  errors.push_back( vertex.GetNegativeDirectionError().x() );
  errors.push_back( vertex.GetNegativeDirectionError().y() );
  errors.push_back( vertex.GetNegativeDirectionError().z() );
  return errors;
}

void
DirectionLikelihood::SetParams( const std::vector<double>& params )
{
  FitVertex vertex = fFitResult.GetVertex(0);
  vertex.SetDirection( TVector3( params[0], params[1], params[2] ), fOptimiser->GetValid() );
  fFitResult.SetVertex( 0, vertex );
}

void
DirectionLikelihood::SetPositiveErrors( const std::vector<double>& errors )
{
  FitVertex vertex = fFitResult.GetVertex(0);
  vertex.SetPositiveDirectionError( TVector3( errors[0], errors[1], errors[2] ), fOptimiser->GetValid() );
  fFitResult.SetVertex( 0, vertex );
}

void
DirectionLikelihood::SetNegativeErrors( const std::vector<double>& errors )
{
  FitVertex vertex = fFitResult.GetVertex(0);
  vertex.SetNegativeDirectionError( TVector3( errors[0], errors[1], errors[2] ), fOptimiser->GetValid() );
  fFitResult.SetVertex( 0, vertex );
}