#include <CLHEP/Units/PhysicalConstants.h>
using namespace CLHEP;

#include <RAT/IsotropyRegions.hh>
#include <RAT/DB.hh>
#include <RAT/DS/FitResult.hh>
#include <RAT/DS/FitVertex.hh>
#include <RAT/DU/Utility.hh>
#include <RAT/DU/PMTInfo.hh>
#include <RAT/Log.hh>
using namespace RAT;
using namespace RAT::Classifiers;
using namespace RAT::DS;

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

void
IsotropyRegions::DefaultSeed()
{
  fEventPos = TVector3();
  fEventTime = 0.0;
}

void
IsotropyRegions::SetSeed( const DS::FitResult& seed )
{
  FitVertex seedVertex;
  try
    {
      seedVertex = seed.GetVertex(0);
    }
  catch( FitResult::NoVertexError& error )
    {
      // Nothing to seed from, strange request
      return;
    }
  try { fEventPos = seedVertex.GetPosition(); }
  catch( FitVertex::NoValueError& error ) { /* No data to seed from */ }
  try { fEventTime = seedVertex.GetTime(); }
  catch( FitVertex::NoValueError& error ) { /* No data to seed from */ }
}

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

  // Primary axis is the fitted event axis, regions are then split around this
  // equivalent to an orange aligned about this axis.
  const TVector3 primaryAxis = fEventPos.Unit();
  const TVector3 secondaryAxis = primaryAxis.Cross( TVector3( 0.0, 0.0, -1.0 ) ).Unit();
  const TVector3 tertiaryAxis = primaryAxis.Cross( secondaryAxis ).Unit();

  // Find the neck limiting longitude
  const TVector3 neckLimit = TVector3( 0.0, 0.0, 6060.6 ) + secondaryAxis * 1000.0; // 1m to the side of neck centre (neck is ~700mm)
  const double limitPhi = atan2( neckLimit.Dot( secondaryAxis ), neckLimit.Dot( tertiaryAxis ) );
  // Loop over the hit PMTs and fill the regions, split by latitude then longitude
  vector< vector< double > > fRegions( 4, vector<double>( 4, 0.0 ) );
  for( vector<FitterPMT>::const_iterator iPMT = fPMTData.begin(); iPMT != fPMTData.end(); ++iPMT )
    {
      const TVector3 pmtDir = ( DU::Utility::Get()->GetPMTInfo().GetPosition( iPMT->GetID() ) - fEventPos ).Unit();
      // Which longitude
      double phi = atan2( pmtDir.Dot( secondaryAxis ), pmtDir.Dot( tertiaryAxis ) ); // NEW switched
      if( phi < 0.0 ) // Must be [0, 2pi]
        phi += twopi;
      if( phi < limitPhi || phi > twopi - limitPhi ) // NEW
        continue;
      // Which latitude
      double cosTheta = pmtDir.Dot( primaryAxis ) + 1.0;
      int longitude = static_cast<int>( ( phi - limitPhi ) / ( twopi - 2.0 * limitPhi ) * 4 ); // NEW Truncates i.e. floors
      int latitude = static_cast<int>( cosTheta / 2.0 * 4 ); // Truncates i.e. floors
      (fRegions[latitude])[longitude] += 1.0;
    }
  // Now each latitude should have the same number of hits in each longitude, so test this
  double fullLatChi = 0.0;
  // NEW Correction for missing longitude segment
  const double expectedHitsPerRegion = static_cast<double>( fPMTData.size() ) / ( 4.0 * 4.0 ) * ( 1.0 - ( 2.0 * limitPhi ) / twopi );
  double fullChi = 0.0;
  for( int iLat = 0; iLat < 4; iLat++ )
    {
      const double latAverage = accumulate( fRegions[iLat].begin(), fRegions[iLat].end(), 0.0 ) / 4.0;
      double chi = 0.0;
      for( int iLong = 0; iLong < 4; iLong++ )
        {
          if( latAverage > 0.2 ) // 0.25 if there is one hit
            chi += pow( (fRegions[iLat])[iLong] - latAverage, 2 ) / latAverage;
          fullChi += pow( (fRegions[iLat])[iLong] - expectedHitsPerRegion, 2 ) / expectedHitsPerRegion;
        }
      fullLatChi += chi;
    }
  fullLatChi /= 4.0;

  fClassifierResult.SetClassification( "fullChi", fullChi );
  fClassifierResult.SetClassification( "fullLatChi", fullLatChi );
  fClassifierResult.SetValid( true );
  return fClassifierResult;
}