#include <G4OpticalPhoton.hh>
#include <G4TransportationManager.hh>
#include <Randomize.hh>
#include <G4Region.hh>
#include <G4UnitsTable.hh>
#include <G4GeometryTolerance.hh>
#include <local_g4compat.hh>

#include <complex>

#include <RAT/DiscOpticalModel.hh>
#include <RAT/Log.hh>
#include <RAT/DB.hh>
#include <RAT/PhotonThinning.hh>
#include <RAT/HitPMTCollection.hh>
#include <RAT/PMTVariation.hh>
#include <RAT/DS/MCPE.hh>
using namespace RAT;

using CLHEP::twopi;
using CLHEP::hbarc;
using CLHEP::degree;
using CLHEP::mm;
using CLHEP::nm;
DiscOpticalModel::DiscOpticalModel( const G4String name,
                                    G4Region* const region,
                                    const std::string params,
                                    G4LogicalVolume* const volume,
                                    const G4double discZ )
: G4VFastSimulationModel( name, region )
{
  fProcessName = "DiscOpticalModel";
  fDiscZ = discZ;

  RAT::DBLinkPtr greyDiscTable = RAT::DB::Get()->GetLink( "GREY_DISC_PARAMETERS", params );

  try
    {
      fTravelTime = greyDiscTable->GetD( "travel_time" );
      fDecayConstant = greyDiscTable->GetD( "decay_constant" );
      fBounceSpread = greyDiscTable->GetD( "bounce_spread" );
      fDiscRadius = greyDiscTable->GetD( "disc_radius" );

      fAbsorption = greyDiscTable->GetDArray( "absorption_probability" );
      fReflection = greyDiscTable->GetDArray( "reflection_probability" );
      fCollectionEfficiency = greyDiscTable->GetD( "collection_efficiency" );

      fPMTEllipsoidA = greyDiscTable->GetD("pmt_ellipsoid_a");
      fPMTEllipsoidB = greyDiscTable->GetD("pmt_ellipsoid_b");
      fPMTZMinContact = greyDiscTable->GetD("pmt_ellipsoid_minz_contact") ;

      fTDelaySpread = greyDiscTable->GetD("reflection_tdelay_spread") ;

      // Hit PMT first, single reflections
      fPMTOneReflRotationTheta  = greyDiscTable->GetD("reflection_model_spmt_trot") ;
      fPMTOneReflMean = greyDiscTable->GetD("reflection_model_spmt_tmean") ;
      fPMTOneReflSpread = greyDiscTable->GetD("reflection_model_spmt_tspread") ;
      fPMTOneReflPhiFlatProb  = greyDiscTable->GetD("reflection_model_spmt_pflat") ;
      fPMTOneReflPhiSpread = greyDiscTable->GetD("reflection_model_spmt_pspread") ;
      fPMTOneDelayFactor = greyDiscTable->GetD("reflection_tdelay_spmt") ;

      // Hit PMT first, multiple reflections
      fPMTMultReflRotationTheta  = greyDiscTable->GetD("reflection_model_mpmt_trot") ;
      fPMTMultReflMean = greyDiscTable->GetD("reflection_model_mpmt_tmean") ;
      fPMTMultReflSpread = greyDiscTable->GetD("reflection_model_mpmt_tspread") ;
      fPMTMultReflPhiSpread  = greyDiscTable->GetDArray("reflection_model_mpmt_pspread") ;
      fPMTMultReflPhiMean  = greyDiscTable->GetDArray("reflection_model_mpmt_pmean") ;
      fPMTMultDelayFactor  = greyDiscTable->GetD("reflection_tdelay_mpmpt");


      // Hitting PETALS first - only one theta model
      fCONCThetaSpread  = greyDiscTable->GetD("reflection_model_conc_tspread") ;
      // Single reflections
      fCONCPhiCoeffs  = greyDiscTable->GetDArray("reflection_model_conc_coeffs") ;
      fCONCOneDelay = greyDiscTable->GetDArray("reflection_tdelay_sconc") ;
      fCONCOneDelayMax = greyDiscTable->GetDArray("reflection_tdelay_sconc_max");
      // Multiple reflections
      fCONCMultPhiSpread  = greyDiscTable->GetD("reflection_model_mconc_pspread") ;
      fCONCMultPhiDiffuse = greyDiscTable->GetD("reflection_model_mconc_pdiffuse") ;
      fCONCMultDelay = greyDiscTable->GetDArray("reflection_tdelay_mconc") ;


      // Parameters that determine the type of reflection
      fCONCReflCoeffs = greyDiscTable->GetDArray("reflection_model_conc_nreflections_coeffs");

      fPMTReflRotation  = greyDiscTable->GetD("reflection_model_pmt_nreflections_rotation") ;
      fPMTReflCoeffs  = greyDiscTable->GetDArray("reflection_model_pmt_nreflections_coeffs");
    }
  catch ( RAT::DBNotFoundError &e)
    {
        RAT::Log::Die( "DiscOpticalModel::DiscOpticalModel: Cannot Find Grey Disc Parameters" );
    }
  fEnvelopeVolume = volume->GetSolid();
}

DiscOpticalModel::~DiscOpticalModel()
{
  // nothing to delete
}

G4bool
DiscOpticalModel::IsApplicable( const G4ParticleDefinition& type )
{
  return ( &type == G4OpticalPhoton::OpticalPhotonDefinition() );
}


// ModelTrigger() method overriding virtual function of G4VFastSimulationModel
// returns true if model should take over this specific track.
// -- see also Geant4 docs
G4bool
DiscOpticalModel::ModelTrigger( const G4FastTrack& )
{
  // Manage all photons in the region thus return true
  return true;
}

void
DiscOpticalModel::DoIt( const G4FastTrack& fastTrack,
                        G4FastStep& fastStep )
{
  // Logic summary
  // 1) Track photon to disc plane.
  //    i)  Check photon hits the disc, if not then kill photon
  // 2) If photon is alive, check probability of signal or reflection
  //    i)  If reflected, then diffuse reflect the photon
  //    ii) If signal then create some PE on the tube
  // 3) If photon is alive, track photon out of envelope.

  G4ThreeVector position =  fastTrack.GetPrimaryTrackLocalPosition();
  G4ThreeVector direction =  fastTrack.GetPrimaryTrackLocalDirection();
  G4ThreeVector polarisation =  fastTrack.GetPrimaryTrackLocalPolarization();
  G4double time = fastTrack.GetPrimaryTrack()->GetGlobalTime();
  G4double weight =  fastTrack.GetPrimaryTrack()->GetWeight();

  G4int ipmt = fastTrack.GetEnvelopePhysicalVolume()->GetCopyNo();

  // Disc intersection test, must not be parallel or travelling up
  if( direction.z() != 0.0 && direction.z() < 0.0 )
    {
      G4double ldt = ( fDiscZ - position.z() ) / direction.z();
      // Track back to the disc
      position += ldt * direction;

      if( sqrt( position.x() * position.x() + position.y() * position.y() ) < fDiscRadius )
        TrackPhoton( fastTrack, ipmt, position, direction, polarisation, time, weight );
      else
        weight = 0.0; // Missed the disc and hit the bucket
    }
    else
      weight = 0.0; // Hit side of envelope or base thus kill photon

  if( weight != 0.0 ) // Reflected, track back out of envelope
    {
      const G4double distOut = fEnvelopeVolume->DistanceToOut( position, direction ) + fSurfaceTolerance;
      position += direction * distOut;

      RecordTrajectory( fastTrack, position, direction, 0.0, time, fastTrack.GetPrimaryTrack()->GetKineticEnergy(), fastTrack.GetEnvelopePhysicalVolume()->GetName(), "GDOut" );
    }

  // Set track properties
  fastStep.SetPrimaryTrackFinalPosition( position );
  fastStep.SetPrimaryTrackFinalTime( time );
  fastStep.SetPrimaryTrackFinalMomentum( direction );
  fastStep.SetPrimaryTrackFinalPolarization( polarisation );
  if( weight <= 0.0 )
    fastStep.ProposeTrackStatus( fStopAndKill );
}

// adapted from the vxpmt_hit.for code.
void
DiscOpticalModel::TrackPhoton( const G4FastTrack& fastTrack,
                               const G4int ipmt,
                               G4ThreeVector& position,
                               G4ThreeVector& direction,
                               G4ThreeVector& polarisation,
                               G4double& time,
                               G4double& weight )
{
  const G4int trackID =  fastTrack.GetPrimaryTrack()->GetTrackID();
  const G4double energy = fastTrack.GetPrimaryTrack()->GetKineticEnergy();

  G4double wavelength = twopi * hbarc / energy;

  // Local co-ord normal to a disc approximation
  G4ThreeVector normal = G4ThreeVector( 0.0, 0.0, 1.0 );
  G4double cosTheta = direction * normal;

  // Now get absorption & reflection probabilities
  G4double absorptionProb = GetAbsorptionProb( acos( -cosTheta ), wavelength );
  G4double reflectionProb = GetReflectionProb( acos( -cosTheta ), wavelength );

  absorptionProb *= RAT::PhotonThinning::GetFactor() * weight * fCollectionEfficiency;
  absorptionProb *= PMTVariation::Get()->GetVariation( ipmt );

  const double uniRand = G4UniformRand();

  //Check if Photon is absorbed
  if( uniRand < absorptionProb )
    {
      // Calculate decay time approximation
      time += fTravelTime - fDecayConstant * log( G4UniformRand() );
      DS::MCPE photoelectron;
      photoelectron.SetPosition( TVector3( position.x(), position.y(), position.z() ) );
      photoelectron.SetCreationTime( time );
      photoelectron.SetCharge( 1 );
      photoelectron.SetPhotonTrackID( trackID );
      ExtractPhotonHistory(fastTrack,photoelectron);
      HitPMTCollection::Get()->DetectPhotoelectron( ipmt, photoelectron );
      RecordPhotoelectron( ipmt );
      weight = 0.0; //As absorbed
    }
  // Now check if reflected
  else if( uniRand < ( reflectionProb + absorptionProb ) )
    {
      DiscReflect( normal, position, direction, polarisation, time );
    }
  else
    {
      weight = 0.0; // No transmission in this model
    }
}

void
DiscOpticalModel::DiscReflect( const G4ThreeVector& normal,
                               G4ThreeVector& position,
                               G4ThreeVector& direction,
                               G4ThreeVector& polarisation,
                               G4double& time)
{
  G4double deltaTheta = 0.;
  G4double deltaPhi   = 0;
  G4double thetaBounceAngle = 0;
  G4double phiBounceAngle   = 0;
  G4double cosInTheta = -direction*normal;
  G4double timeDelayMean = 0;

  // Calculating contact time with the PMT (zero if the photon touches concentrator instead)
  G4double contactTime = GetPMTContactTime(position, direction);
  G4double radius = sqrt(position.x()*position.x() + position.y()*position.y());

  // There are four reflection models:
  // 1. Photon hits PMT first, escapes after a single reflection
  // 2. Photon hits PMT first, escapes after multiple reflections
  // 3. Photon hits CONCENTRATOR first, escapes after a single reflection
  // 4. Photon hits CONCENTRATOR first, escapes after multiple reflections
  if( contactTime > 0.)
    { // PMT contact

      // bool singleReflection = IsPMTReflectionSingle( radius, direction);
      if( IsPMTReflectionSingle( radius, direction) )
        {
          // Angles
          G4double rotationAngle = fPMTOneReflRotationTheta ;
          G4double rotatedSpread = G4RandGauss::shoot()*fPMTOneReflSpread+fPMTOneReflMean;

          deltaTheta = (rotatedSpread - acos(-direction.z())*sin(rotationAngle))/cos(rotationAngle);

          if( G4UniformRand() < fPMTOneReflPhiFlatProb)
            { deltaPhi = G4UniformRand()*twopi ;}
          else
            { deltaPhi = G4RandGauss::shoot()*fPMTOneReflPhiSpread ;}

          // Time delay
          timeDelayMean = contactTime/fPMTOneDelayFactor ;
        }
      else
        {
          // Angles
          G4double rotationAngle = fPMTMultReflRotationTheta ;
          G4double rotatedSpread = G4RandGauss::shoot()*fPMTMultReflSpread + fPMTMultReflMean ;
          deltaTheta = (rotatedSpread - acos(-direction.z())*sin(rotationAngle))/cos(rotationAngle);

          G4double phiSign = 1.;
          if (G4UniformRand() > 0.5)
            {phiSign = -1.;}

          deltaPhi = (G4RandGauss::shoot()*(fPMTMultReflPhiSpread[0]*radius + fPMTMultReflPhiSpread[1]) +
                      twopi/2.+phiSign*(fPMTMultReflPhiMean[0]*radius+ fPMTMultReflPhiMean[1]));

          // Time delay
          timeDelayMean = contactTime/fPMTMultDelayFactor;
        }
    }
  else
    {  // CONCENTRATOR contact

      // DeltaTheta for concentrator hits is independent of number of reflections
      deltaTheta =  G4RandGauss::shoot()*fCONCThetaSpread ;

      if( IsPetalReflectionSingle(acos(-direction.z())))
        {
          // Angle
          G4double phiSign = 1.;
          if (G4UniformRand() > 0.5)
            {phiSign = -1.;}
          deltaPhi = (G4RandGauss::shoot()*(fCONCPhiCoeffs[0]*radius + fCONCPhiCoeffs[1]) +
                      phiSign*(fCONCPhiCoeffs[2]*radius + fCONCPhiCoeffs[3])) ;

          // Time delay
          if (cosInTheta > fCONCOneDelayMax[0])
            {timeDelayMean = fCONCOneDelayMax[1];}
          else
            {
              timeDelayMean = fCONCOneDelay[0]*pow(cosInTheta,4) + fCONCOneDelay[1]*pow(cosInTheta,3) + \
                fCONCOneDelay[2]*pow(cosInTheta,2) + fCONCOneDelay[3]*cosInTheta + fCONCOneDelay[4];
            }
        }
      else
        {

          if (G4UniformRand() > fCONCMultPhiDiffuse)
            {deltaPhi = G4RandGauss::shoot()*fCONCMultPhiSpread + round(G4UniformRand()*9.)*twopi/9.;}
          else
            {deltaPhi = G4UniformRand()*twopi; }

          // This polynomial is very sensitive to the values used as input (10th degree!). Modify with care.
          timeDelayMean=fCONCMultDelay[0]*pow(cosInTheta,10)+fCONCMultDelay[1]*pow(cosInTheta,9) + \
            fCONCMultDelay[2]*pow(cosInTheta,8) + fCONCMultDelay[3]*pow(cosInTheta,7) + \
            fCONCMultDelay[4]*pow(cosInTheta,6) + fCONCMultDelay[5]*pow(cosInTheta,5) + \
            fCONCMultDelay[6]*pow(cosInTheta,4) + fCONCMultDelay[7]*pow(cosInTheta,3) + \
            fCONCMultDelay[8]*pow(cosInTheta,2) + fCONCMultDelay[9]*cosInTheta + fCONCMultDelay[10];
        }

    }

  // Cerifying that the angles obtained make sense
  thetaBounceAngle = acos(-direction.z()) + deltaTheta;
  phiBounceAngle = atan2(direction.y(), direction.x()) + deltaPhi ;

  if (thetaBounceAngle < 0. || thetaBounceAngle > twopi/4.)
    {
      thetaBounceAngle = G4UniformRand()*twopi/4.;
    }

  if (phiBounceAngle > 0)
    {phiBounceAngle = phiBounceAngle - twopi*floor(phiBounceAngle/twopi);}
  else
    {phiBounceAngle = twopi - (abs(phiBounceAngle) - twopi*floor(abs(phiBounceAngle)/twopi));}

  // Applying the reflection angle
  RayScatter( normal, thetaBounceAngle, phiBounceAngle, direction );

  // Applying the time delay
  G4double timeDelay = timeDelayMean +  G4RandGauss::shoot()*fTDelaySpread;
  if (timeDelay < 0)
    { timeDelay = 1. ;}
  time += timeDelay ;


  // Change the polarisation to follow OpticalModelBase::DiffuseReflect
  RayScatter( direction, twopi / 4.0, twopi * G4UniformRand(), polarisation );
  polarisation = polarisation.unit();
}

G4double
DiscOpticalModel::GetReflectionProb( const G4double theta, // In Radians
                                     const G4double wavelength ) const // In mm
{
  G4int iTheta = static_cast<G4int>( ( theta / degree ) + 0.5 );
  G4int iWavelength = static_cast<G4int>( ( wavelength * mm ) / ( 10.0 * nm ) + 0.5 ); // Want in units of 10nm

  if( iWavelength < 22 || iWavelength > 71 )
    iWavelength = 22;
  iWavelength -= 22; // Data File starts at 220nm thus 22 = index 0, and runs to 710nm
  return fReflection[ iTheta + iWavelength * 90 ];
}

G4double
DiscOpticalModel::GetAbsorptionProb( const G4double theta, // In Radians
                                     const G4double wavelength ) const // In mm
{
  G4int iTheta = static_cast<G4int>( ( theta / degree ) + 0.5 );
  G4int iWavelength = static_cast<G4int>( ( wavelength * mm ) / ( 10.0 * nm ) + 0.5 ); // Want in units of 10nm

  if( iWavelength < 22 || iWavelength > 71 )
    iWavelength = 22;
  iWavelength -= 22; // Data File starts at 220nm thus 22 = index 0, and runs to 710nm
  return fAbsorption[ iTheta + iWavelength * 90 ];
}


G4bool
DiscOpticalModel::IsPMTReflectionSingle(const G4double radius,
                                        G4ThreeVector& direction) const
{

  G4double rotationAngle =  fPMTReflRotation ;
  G4double rotatedY = -direction.z()*sin(rotationAngle) + radius*cos(rotationAngle);

  // Evaluate the rotated value in the single reflection PDF
  G4double singleReflectionPDF = (fPMTReflCoeffs[0]*pow(rotatedY,2) + fPMTReflCoeffs[1]*rotatedY  +
                                  fPMTReflCoeffs[2]) ;

  if(G4UniformRand() < singleReflectionPDF )
    { return true;}
  else
    { return false;}

}

G4bool
DiscOpticalModel::IsPetalReflectionSingle(const G4double theta) const
{
  G4double multRefProb = (fCONCReflCoeffs[0]*pow(theta,5) + fCONCReflCoeffs[1]*pow(theta,4) +
                          fCONCReflCoeffs[2]*pow(theta,3) + fCONCReflCoeffs[3]*pow(theta,2) +
                          fCONCReflCoeffs[4] *theta + fCONCReflCoeffs[5]);

  if (G4UniformRand() > multRefProb)
    {return true;}
  else
    {return false;}
}

G4double
DiscOpticalModel::GetPMTContactTime( G4ThreeVector& position,
                                     G4ThreeVector& direction) const
{

  // The PMT surface is characterized by an ellipsoid with semiaxes of ae, ce size (mm)
  // Solving the quadratic equation to find the intersection of an ellipsoid and a ray (photon)
  G4double ae = fPMTEllipsoidA ;
  G4double ce = fPMTEllipsoidB ;

  G4double contact_time = 0.;

  G4double a = pow(direction.z()/ce,2) + ( pow(direction.x(),2) + pow(direction.y(),2))/pow(ae,2);
  G4double b_numA = direction.x()*position.x()+direction.y()*position.y() ;
  G4double b_numB = direction.z()*position.z() ;
  G4double b = 2.*(b_numA/pow(ae,2) + b_numB/pow(ce,2));
  G4double c = (pow(position.x(),2)+pow(position.y(),2))/pow(ae,2)+pow(position.z()/ce,2)-1.;

  G4double inner_sqrt = pow(b,2) - 4*a*c;

  if( (inner_sqrt < 0) && (inner_sqrt > -0.0001))
    {inner_sqrt = 0.;}

  if( inner_sqrt > 0)
    {
      G4double sol1 = (-b+sqrt(inner_sqrt))/(2*a);
      G4double sol2 = (-b-sqrt(inner_sqrt))/(2*a);

      if ((sol1 > 0) && (sol2 > 0))
        {
          if((sol1 > 0) && (sol2 < 0))
            {contact_time = sol1;}
          else if((sol2 > 0) && (sol1 < 0))
            {contact_time = sol2;}
          else if ((sol1 > 0) && (sol2 > 0))
            { if(sol1 < sol2)
                {contact_time = sol1;}
              else
                {contact_time = sol2;}
            }
          // In the geometry, the PMT face is only exposed for z positions > 26mm - below is hidden
          // Contact times below this value are with the bottom face of the PMT, but skipping the top
          // These are concentrator hits
          if ((contact_time*direction.z()+position.z() ) < fPMTZMinContact )
            {contact_time = 0;}
        }

    }
  return contact_time;
}