#include <Randomize.hh>
#include <G4ThreeVector.hh>
#include <G4FastTrack.hh>
#include <G4VTouchable.hh>
#include <G4RunManager.hh>
#include <local_g4compat.hh>

#include <RAT/OpticalModelBase.hh>
#include <RAT/TrackingAction.hh>
#include <RAT/PhotonTrajectory.hh>
#include <RAT/DS/MCTrackStep.hh>
#include <RAT/RunManager.hh>
#include <RAT/Gsim.hh>
#include <RAT/Trajectory.hh>

#include <G4VProcess.hh>
#include <G4OpProcessSubType.hh>
#include <G4EmProcessSubType.hh>
#include <RAT/UserTrackInformation.hh>
#include <RAT/Log.hh>

using namespace RAT;

using CLHEP::mm;
using CLHEP::twopi;

G4int OpticalModelBase::fsTrackID = -1;
G4int OpticalModelBase::fsPMTID = -1;
G4double OpticalModelBase::fsBucketCos = 0.0;

OpticalModelBase::OpticalModelBase()
{
  fMaxIterations = 100;
  fSurfaceTolerance = 0.005 * mm;
}

double
OpticalModelBase::GetSPolarisationFraction( const G4ThreeVector& polarisation,
                                            const G4ThreeVector& normal,
                                            const G4ThreeVector& direction )
{
  G4ThreeVector planeNormal = direction.cross( normal );
  planeNormal = planeNormal.unit();
  const double sFraction = polarisation.dot( planeNormal );
  return sFraction * sFraction;
}

OpticalModelBase::EPolarisation
OpticalModelBase::GetPolarisation( const G4ThreeVector& normal,
                                   const G4ThreeVector& direction,
                                   const G4ThreeVector& polarisation )
{
  const double sFraction = GetSPolarisationFraction( normal, direction, polarisation );

  const double uniRand = G4UniformRand();
  if( uniRand  < sFraction )
    {
      return esPolarised;
    }
  else
    {
      return epPolarised;
    }
}

void
OpticalModelBase::NormalReflect( const G4ThreeVector& normal,
                                 G4ThreeVector& direction,
                                 G4ThreeVector& polarisation )
{
  // TESTED
  //assert( polarisation.mag() == 1.0 && normal.mag() == 1.0 && direction.mag() == 1.0 );
  direction = direction - 2.0 * direction.dot( normal ) * normal;
  polarisation = polarisation - 2.0 * polarisation.dot( normal ) * normal;
  direction = direction.unit();
  polarisation = polarisation.unit();
}

void
OpticalModelBase::DiffuseReflect( const G4ThreeVector& normal,
                                  G4ThreeVector& direction,
                                  G4ThreeVector& polarisation )
{
  // TESTED
  // First pick a random direction theta and phi
  double theta = 0.99 * G4UniformRand() * twopi / 8.0; //No Perpendicular to normal
  double phi = twopi * G4UniformRand();
  RayScatter( normal, theta, phi, direction );

  // Now do the same for the polarisation, Note must be normalal to the direction
  theta = twopi / 4.0;
  phi = twopi * G4UniformRand();
  RayScatter( direction, theta, phi, polarisation );
  direction = direction.unit();
  polarisation = polarisation.unit();
}

void
OpticalModelBase::Refract( const G4ThreeVector& normal,
                           const G4double n1,
                           const G4double n2,
                           G4ThreeVector& direction,
                           G4ThreeVector& polarisation )
{
  // TESTED
  const double nRatio = n1 / n2;
  double cosTheta1 = direction.dot( normal );
  const double sinTheta1_2 = 1.0 - cosTheta1 * cosTheta1;
  const double cosTheta2 = sqrt( 1 - nRatio * nRatio * sinTheta1_2 );
  direction = nRatio * direction - nRatio * cosTheta1 * normal - cosTheta2 * normal;
  polarisation = polarisation - polarisation.dot( direction ) * direction;

  direction = direction.unit();
  polarisation = polarisation.unit();
}

// As in pmt_rayscatter.for
void
OpticalModelBase::RayScatter( const G4ThreeVector& normal,
                              const G4double theta,
                              const G4double phi,
                              G4ThreeVector& resp )
{
  // TESTED
  const double cosTheta = cos( theta );
  const double sinTheta = sin( theta );
  const double cosPhi = cos( phi );
  const double sinPhi = sin( phi );

  const G4double u = normal.x();
  const G4double v = normal.y();
  const G4double w = normal.z();

  G4double sinAlpha = sqrt( 1 - w*w );

  G4double cosBeta, sinBeta;
  if( sinAlpha > 0.0 )
    {
      cosBeta = u / sinAlpha;
      sinBeta = v / sinAlpha;
    }
  else
    {
      cosBeta = 1.0;
      sinBeta = 0.0;
    }

  resp.setX( cosTheta * u + sinTheta * ( w * cosPhi * cosBeta - sinPhi * sinBeta ) );
  resp.setY( cosTheta * v + sinTheta * ( w * cosPhi * sinBeta + sinPhi * cosBeta ) );
  resp.setZ( cosTheta * w - sinTheta * cosPhi * sinAlpha );
  resp = resp.unit();
}

void
OpticalModelBase::RecordTrajectory( const G4FastTrack& fastTrack,
                                    const G4ThreeVector& endPos,
                                    const G4ThreeVector& direction,
                                    const G4double stepLength,
                                    const G4double time,
                                    const G4double energy,
                                    const std::string& startVolume,
                                    const std::string& endVolume )
{
  // I'm not sure why the run manager returns a const, it doesn't need to...
  TrackingAction* trackingAction = dynamic_cast<TrackingAction*>( const_cast<G4UserTrackingAction*>( G4RunManager::GetRunManager()->GetUserTrackingAction() ) );
  if( trackingAction == NULL )
    return; // User is using a special tracking action, thus leave
  PhotonTrajectory* trajectory = dynamic_cast<PhotonTrajectory*>( trackingAction->GetTrajectory() );
  if( trajectory == NULL )
    return; // Not storing trajectories thus leave this

  const G4ThreeVector globalEndPos = fastTrack.GetInverseAffineTransformation()->TransformPoint( endPos );
  const G4ThreeVector globalDir = fastTrack.GetInverseAffineTransformation()->TransformAxis( direction );
  const G4double localTime = time - fastTrack.GetPrimaryTrack()->GetGlobalTime() + fastTrack.GetPrimaryTrack()->GetLocalTime();

  TVector3 globalMom( globalDir.x(), globalDir.y(), globalDir.z() );
  globalMom *= energy; // As photon

  RAT::DS::MCTrackStep thisStep;
  thisStep.SetLength( stepLength );
  thisStep.SetPosition( TVector3( globalEndPos.x(), globalEndPos.y(), globalEndPos.z() ) );
  thisStep.SetGlobalTime( time );
  thisStep.SetLocalTime( localTime );
  thisStep.SetProperTime( 0.0 ); //?
  thisStep.SetMomentum( globalMom );
  thisStep.SetKineticEnergy( energy );
  thisStep.SetProcess( fProcessName );
  thisStep.SetEndVolume( endVolume );
  thisStep.SetStartVolume( startVolume );
  thisStep.SetStatus( RAT::DS::MCTrackStep::GeomBoundary );
  trajectory->AppendStep( thisStep );
}

void
OpticalModelBase::RecordPhotoelectron( const int pmtID )
{
  // I'm not sure why the run manager returns a const, it doesn't need to...
  TrackingAction* trackingAction = dynamic_cast<TrackingAction*>( const_cast<G4UserTrackingAction*>( G4RunManager::GetRunManager()->GetUserTrackingAction() ) );
  if( trackingAction == NULL )
    return; // User is using a special tracking action, thus leave
  PhotonTrajectory* trajectory = dynamic_cast<PhotonTrajectory*>( trackingAction->GetTrajectory() );
  if( trajectory == NULL )
    return; // Not storing trajectories thus leave this
  trajectory->AppendPhotoelectron( pmtID );
}

int
OpticalModelBase::GetPMTID( const G4FastTrack& fastTrack )
{
  int iDepth;
  for( iDepth = 0; iDepth < fastTrack.GetPrimaryTrack()->GetTouchableHandle()->GetHistoryDepth(); iDepth++ )
    {
      const string volName = fastTrack.GetPrimaryTrack()->GetTouchableHandle()->GetVolume( iDepth )->GetName();
      const size_t envelopeHistory = volName.find( "_pmtenv" );
      if( envelopeHistory != string::npos )
        return fastTrack.GetPrimaryTrack()->GetTouchableHandle()->GetCopyNumber( iDepth );
    }
  // Should never get here
  Log::Die( "OpticalModelBase::GetPMTID: Cannot decode PMT ID." );
  return -1;
}

void
OpticalModelBase::SetBucketCos( const G4int trackID,
                                const G4int pmtID,
                                const G4double bucketCos )
{
  if( trackID != fsTrackID || pmtID != fsPMTID )
    {
      fsTrackID = trackID;
      fsPMTID = pmtID;
      if( bucketCos < 0.0 )
	fsBucketCos = -bucketCos;
      else
	fsBucketCos = bucketCos;
    }
}

void OpticalModelBase::ExtractPhotonHistory(const G4FastTrack& fastTrack, DS::MCPE &photoelectron) {
  UserTrackInformation *history = dynamic_cast<UserTrackInformation *>(fastTrack.GetPrimaryTrack()->GetUserInformation());
  if (history) {
    // History already exists. Simply copy it
    photoelectron.SetHistory(history->GetHistory());
  } else {
#ifdef RATDEBUG
    warn << "OpticalModelBase::ExtractPhotonHistory : Didn't find track user information. Storing only track creating process." << newline;
#endif
    const G4VProcess *crProc = fastTrack.GetPrimaryTrack()->GetCreatorProcess();
    if (!crProc) {
#ifdef RATDEBUG
      warn << "OpticalModelBase::ExtractPhotonHistory : Creator Process not found." << newline;
#endif
      return;
    }
    if (crProc->GetProcessType() == fOptical) {
      if ( crProc->GetProcessSubType() == fOpRayleigh) {
          photoelectron.AddToHistory(DS::MCPE::hRayleighScatt);
        } else if (crProc->GetProcessSubType() == fOpMieHG) {
          photoelectron.AddToHistory(DS::MCPE::hMieScatt);
        } else {
          warn << "OpticalModelBase::ExtractPhotonHistory :  Unknown optical process subtype : " << crProc->GetProcessName() << newline;
          photoelectron.AddToHistory(DS::MCPE::hUnknown);
        }
    } else if (crProc->GetProcessType() == fUserDefined) {
        // This could means re-emission or scintillation
      if (crProc->GetProcessName()  == "Scintillation") {
        photoelectron.AddToHistory(DS::MCPE::hScintillation);
      } else if (crProc->GetProcessName().find("Reemission")  != std::string::npos) {
        // This envelopes several types of reemission (from the different componentds)
        photoelectron.AddToHistory(DS::MCPE::hWLS);
      } else {
        warn << "OpticalModelBase::ExtractPhotonHistory : Unknown user process : " << crProc->GetProcessName() << newline;
        photoelectron.AddToHistory(DS::MCPE::hUnknown);
      }
    } else if (crProc->GetProcessType() == fElectromagnetic) {
      if (crProc->GetProcessSubType() == fCerenkov) {
        photoelectron.AddToHistory(DS::MCPE::hCherenkov);
      } else {
        warn << "OpticalModelBase::ExtractPhotonHistory : Unknown EM process : " << crProc->GetProcessName() << newline;
        photoelectron.AddToHistory(DS::MCPE::hUnknown);
      }
    } else {
      warn << "OpticalModelBase::ExtractPhotonHistory : Unknown process : " << crProc->GetProcessName() << newline;
      photoelectron.AddToHistory(DS::MCPE::hUnknown);
    }
  }

}