#include <G4UnitsTable.hh>
#include <G4PhysicalConstants.hh>

#include <RAT/SOCPositionTimeChiSquared.hh>
#include <RAT/Optimiser.hh>
#include <RAT/PMTSelector.hh>
#include <RAT/DB.hh>
#include <RAT/Processor.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/SOC.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.hh>
#include <RAT/DU/LightPathCalculator.hh>
#include <RAT/DU/GroupVelocity.hh>
using namespace RAT;
using namespace RAT::Methods;
using namespace RAT::DS;

#include <CLHEP/Units/SystemOfUnits.h>
#include <CLHEP/Units/PhysicalConstants.h>
#include <cmath>
#include <TMath.h>
using namespace CLHEP;
using namespace std;

void
SOCPositionTimeChiSquared::Initialise( const std::string& )
{
  fPhotonEnergyOverride = false;
}

void
SOCPositionTimeChiSquared::BeginOfRun( DS::Run& )
{
  DB* db = DB::Get();
  fAVRadius = db->GetLink("SOLID", "acrylic_vessel_outer")->GetD("r_sphere");
  fLightPath = DU::Utility::Get()->GetLightPathCalculator();
}

void
SOCPositionTimeChiSquared::SetD( const std::string& param,
                                 const double value )
{
  if( param == string( "wavelength" ) ){
    fPhotonEnergy = h_Planck * c_light / ( value * nm );
    fPhotonEnergyOverride = true;
    info << "SOCPositionTimeChiSquared::SetD: Setting wavelength to: " << value << " nm.\n";
  }
  else{ throw Processor::ParamUnknown( param ); }
}

void
SOCPositionTimeChiSquared::DefaultSeed()
{
  FitVertex vertex;
  //Initialise the SeedResult
  vertex.SetPosition( TVector3( 0.0, 0.0, 0.0 ), false, true );
  vertex.SetPositionErrors( ROOT::Math::XYZVectorF(6000.0, 6000.0, 6000.0), false, true );
  vertex.SetTime( 230.0, false, true );
  vertex.SetTimeErrors( 100.0, false, true );
  fSeedResult.SetVertex( 0, vertex );
  if ( !fPhotonEnergyOverride ){
    info << "SOCPositionTimeChiSquared::DefaultSeed: Setting wavelength to: 400 nm.\n";
    fPhotonEnergy = h_Planck * c_light / ( 400.0 * nm );
  }
}

void
SOCPositionTimeChiSquared::SeedBySOC( DS::SOC& soc )
{
  FitVertex vertex;
  // Set the seed position and errors.
  vertex.SetPosition( soc.GetCalib().GetPos(), false, true );
  vertex.SetPositionErrors( ROOT::Math::XYZVectorF(6000.0, 6000.0, 6000.0), false, true );

  // Set the global-time offset from the SOC object
  // as the seed time.
  const double gOffset = soc.GetGlobalTimeOffset();

  // Ensure the global-offset fals within the conventional
  // event-time window of (0, 500)ns...
  if ( gOffset >= 0.0 && gOffset < 500.0 ){
    vertex.SetTime( gOffset, false, true );
  }
  // ..otherwise set to the default global-offset value of 230.0 ns.
  else{
    vertex.SetTime( 230.0, false, true );
  }
  vertex.SetTimeErrors( 100.0, false, true );

  // Set the seed.
  fSeedResult.SetVertex( 0, vertex );

  // Set the photon wavelength.
  if ( !fPhotonEnergyOverride ){
    fPhotonEnergy = h_Planck * c_light / ( (double)soc.GetCalib().GetMode() * nm );
    info << "SOCPositionTimeChiSquared::SeedBySOC: Setting wavelength to: " << (double)soc.GetCalib().GetMode() << " nm.\n";
  }
}

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

  SelectPMTData( fFitResult.GetVertex(0) );
  fFitResult.SetFOM( "SelectedNHit", static_cast<double>( fSelectedPMTData.size() ) );
  if( fSelectedPMTData.empty() )
    throw MethodFailError( "SOCPositionTimeChiSquared: No hits to fit" );

  fOptimiser->SetComponent( this );
  // Call the minimiser and save the fom
  const double chiSquared = fOptimiser->Minimise();
  fFitResult.SetFOM( "Chi2", chiSquared );
  // Now save the optimised values
  SetParams( fOptimiser->GetParams() );
  SetPositiveErrors( fOptimiser->GetPositiveErrors() );
  SetNegativeErrors( fOptimiser->GetNegativeErrors() );
  // Now check optimiser result, must assume optimiser is dumb
  FitVertex vertex = fFitResult.GetVertex( 0 );
  // Allow a 100mm error (push back to av) if greater then invalid fit
  if( vertex.GetPosition().Mag() > fAVRadius + 100.0 || std::isnan( vertex.GetPosition().Mag() ) )
    vertex.SetPositionValid( false );
  else if( vertex.GetPosition().Mag() > fAVRadius )
    vertex.SetPosition( vertex.GetPosition().Unit() * fAVRadius, true );

  fFitResult.SetVertex( 0, vertex );
  return fFitResult;
}

double
SOCPositionTimeChiSquared::operator()( const std::vector<double>& params )
{
  SetParams( params );
  FitVertex vertex = fFitResult.GetVertex( 0 );
  const double eventTime = vertex.GetTime();
  const TVector3 startPos = vertex.GetPosition();

  DU::GroupVelocity gVelo = DU::Utility::Get()->GetGroupVelocity();

  double chiSquared = 0.0;
  for( vector<FitterPMT>::const_iterator iPMT = fSelectedPMTData.begin(); iPMT != fSelectedPMTData.end(); ++iPMT )
    {

      fLightPath.CalcByPosition( startPos, DU::Utility::Get()->GetPMTInfo().GetPosition( iPMT->GetID() ) );
      double distInInnerAV = fLightPath.GetDistInInnerAV();
      double distInAV = fLightPath.GetDistInAV();
      double distInWater = fLightPath.GetDistInWater();

      double incidentAngle = acos( fLightPath.GetIncidentVecOnPMT().Dot( DU::Utility::Get()->GetPMTInfo().GetDirection( iPMT->GetID() ) ) );
      incidentAngle = incidentAngle*TMath::RadToDeg();

      /// Calculate the transit time with an extra 1.0ns for transit inside pmt
      const double transitTime = gVelo.PMTBucketTime( incidentAngle ) + gVelo.CalcByDistance( distInInnerAV,
                                                                                              distInAV,
                                                                                              distInWater,
                                                                                              fPhotonEnergy );
      double pmtTime= iPMT->GetTime();
      const double residualTime = pmtTime - eventTime - transitTime;
      chiSquared += pow( ( residualTime  /  iPMT->GetTimeError() ) , 2);
    }
  return chiSquared;
}

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

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

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

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

void
SOCPositionTimeChiSquared::SetPositiveErrors( const std::vector<double>& errors )
{
  FitVertex vertex = fFitResult.GetVertex(0);
  vertex.SetPositivePositionError( ROOT::Math::XYZVectorF(errors[0], errors[1], errors[2]), fOptimiser->GetValid() );
  vertex.SetPositiveTimeError( errors[3], fOptimiser->GetValid() );
  fFitResult.SetVertex( 0, vertex );
}

void
SOCPositionTimeChiSquared::SetNegativeErrors( const std::vector<double>& errors )
{
  FitVertex vertex = fFitResult.GetVertex(0);
  vertex.SetNegativePositionError( ROOT::Math::XYZVectorF(errors[0], errors[1], errors[2]), fOptimiser->GetValid() );
  vertex.SetNegativeTimeError( errors[3], fOptimiser->GetValid() );
  fFitResult.SetVertex( 0, vertex );
}