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

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

#include <complex>
#include <sstream>
#include <vector>
#include <string>

#include <iostream>
using namespace std;

#include <RAT/PhotonThinning.hh>
#include <RAT/Log.hh>
#include <RAT/DB.hh>
#include <RAT/ConcentratorOpticalModel.hh>
using namespace RAT;

ConcentratorOpticalModel::ConcentratorOpticalModel( const G4String name,
                                                    G4Region* const region,
                                                    const std::string params,
                                                    G4LogicalVolume* const volume )
  : G4VFastSimulationModel( name, region )
{
  fProcessName = "ConcentratorOpticalModel";

  // First load the overall model, Optics[something]...
  DBLinkPtr concModelTable = DB::Get()->GetLink( "CONC_OPTICAL_PARAMETERS", params );
  string reflParams;
  try
    {
      reflParams = concModelTable->GetS( "refl_parameters" );
    }
  catch( RAT::DBNotFoundError &e )
    {
      RAT::Log::Die( "ConcentratorOpticalModel::ConcentratorOpticalModel: Cannot Find CONC_OPTICAL_PARAMETERS Parameters" );
    }

  // Now load the correct reflectivity
  DBLinkPtr concReflTable = DB::Get()->GetLink( "CONC_REFL_PARAMETERS", reflParams );
  try
    {
      fReflectionSPol = concReflTable->GetDArray( "reflection_s_probability" );
      fReflectionPPol = concReflTable->GetDArray( "reflection_p_probability" );
      fDiffuseReflectionProb = concReflTable->GetD( "diffuse_reflection_probability" );
      Float_t reflectionSScaling = concReflTable->GetD( "reflection_s_probability_scaling" );
      Float_t reflectionPScaling = concReflTable->GetD( "reflection_p_probability_scaling" );
      Float_t diffuseScaling = concReflTable->GetD( "diffuse_reflection_probability_scaling" );

      // Only warn if scaling for (s,p)-reflectivity probabilities has changed from default values of 1.0
      if ( reflectionSScaling != 1.0 || reflectionPScaling != 1.0 ){
        warn << "ConcentratorOpticalModel::ConcentratorOpticalModel: Scaling concentrator (s,p)-reflectivity probabilities by: ( "
             << reflectionSScaling << ", " << reflectionPScaling << " )\n";
        for ( Int_t sVal = 0; sVal < (Int_t)fReflectionSPol.size(); sVal++ ){ fReflectionSPol[ sVal ] *= reflectionSScaling; }
        for ( Int_t pVal = 0; pVal < (Int_t)fReflectionPPol.size(); pVal++ ){ fReflectionPPol[ pVal ] *= reflectionPScaling; }
      }

      // Only warn if scaling for diffuse reflection probability has changed from default value of 1.0
      if ( diffuseScaling != 1.0 ){
        warn << "ConcentratorOpticalModel::ConcentratorOpticalModel: Scaling concentrator diffuse reflectivity probability by: "
             << diffuseScaling << "\n";
        fDiffuseReflectionProb *= diffuseScaling;
      }
    }
  catch( RAT::DBNotFoundError &e )
    {
      RAT::Log::Die( "ConcentratorOpticalModel::ConcentratorOpticalModel: Cannot Find CONC_REFL_PARAMETERS Parameters" );
    }

  fGVelocityWater = volume->GetMaterial()->GetMaterialPropertiesTable()->GetProperty("GROUPVEL");

  fConcEnvelopeVolume = volume->GetSolid();
  fPetalVolume = volume->GetDaughter(0)->GetLogicalVolume()->GetSolid(); // Assumes Petals are 0 volume
  fPlasticVolume = volume->GetDaughter(1)->GetLogicalVolume()->GetSolid(); // Assumes Plastic is 1 volume

  fConcVolNames.push_back( volume->GetName() );
  fConcVolNames.push_back( volume->GetDaughter(0)->GetName() );
  fConcVolNames.push_back( volume->GetDaughter(1)->GetName() );
  fConcVolNames.push_back( volume->GetName() + "_out" );
}

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

// IsApplicable() method overriding virtual function of G4VFastSimulationModel
// returns true if model is applicable to given particle.
// -- see also Geant4 docs
G4bool
ConcentratorOpticalModel::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
ConcentratorOpticalModel::ModelTrigger( const G4FastTrack& )
{
  return true; //Should track all
}


void
ConcentratorOpticalModel::DoIt( const G4FastTrack& fastTrack,
                                G4FastStep& fastStep )
{
  // Logic Summary
  // 1) Photons are propagated through

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

  const G4ThreeVector bucketNorm( 0.0, 0.0, 1.0 );

  int iIterations;
  for( iIterations = 0; iIterations < fMaxIterations; iIterations++ ) // Limit number of iterations
  {
    if( weight == 0.0 )
      break;

    if( fPetalVolume->Inside( position ) == kInside )
    {
      const G4double distBackOut = fPetalVolume->DistanceToOut( position, -direction );// + fSurfaceTolerance;
      position -= direction * distBackOut;
    }
    else if( fPlasticVolume->Inside( position ) == kInside )
    {
      const G4double distBackOut = fPlasticVolume->DistanceToOut( position, -direction );// + fSurfaceTolerance;
      position -= direction * distBackOut;
    }
    else if( fConcEnvelopeVolume->Inside( position ) != kOutside ) // In the water region Most likely
    {
      position += -direction * fConcEnvelopeVolume->DistanceToOut( position, -direction ); // Must start on surface
      TrackPhoton( fastTrack, position, direction, polarisation, time, weight );
      break; // Should be done by here
    }
    else
    {
      if( iIterations == 0 )
      {
        const G4int trackID = fastTrack.GetPrimaryTrack()->GetTrackID();
        G4cout << trackID << " " << "ConcTrack Error: Track is lost, Killing" << G4endl;
        weight = 0.0;
      }
      break;
    }
  }
 fastStep.SetPrimaryTrackFinalPosition( position );
 fastStep.SetPrimaryTrackFinalTime( time );
 fastStep.SetPrimaryTrackFinalMomentum( direction );
 fastStep.SetPrimaryTrackFinalPolarization( polarisation );
 if( weight <= 0.0 )
   fastStep.ProposeTrackStatus( fStopAndKill );
}

void
ConcentratorOpticalModel::TrackPhoton( const G4FastTrack& fastTrack,
                                       G4ThreeVector& position,
                                       G4ThreeVector& direction,
                                       G4ThreeVector& polarisation,
                                       G4double& time,
                                       G4double& weight )
{
  const G4double energy = fastTrack.GetPrimaryTrack()->GetKineticEnergy();
  EConcVolumes currentVolume = eEnvelope;
  EConcVolumes previousVolume = eOut;

  const G4double gvWater = fGVelocityWater->Value( energy );

  int iIterations;
  for( iIterations = 0; iIterations < fMaxIterations; iIterations++ ) // Limit number of iterations
    {
      if( weight == 0.0 )
        break;

      if( currentVolume == eEnvelope )
        {
          vector<G4double> distances(3);
          const EConcVolumes volume[3] = { eOut, ePetal, ePlastic };
          distances[0] = fConcEnvelopeVolume->DistanceToOut( position, direction ); // Track to surface
          distances[1] = fPetalVolume->DistanceToIn( position, direction ); // Track to surface
          distances[2] = fPlasticVolume->DistanceToIn( position, direction ); // Track to surface

          const int nextVolume = min_element( distances.begin(), distances.end() ) - distances.begin();
          position += direction * distances[nextVolume];
          time += distances[nextVolume] / gvWater;
          previousVolume = currentVolume;
          currentVolume = volume[nextVolume];
        }
      else if( currentVolume == ePetal )
        {
          if( Reflect( position, energy, direction, polarisation, weight ) )
            currentVolume = previousVolume;
          if( weight == 0.0 )
            break;
        }
      else if( currentVolume == ePlastic )
        {
          weight = 0.0;
          break;
        }
      else if( currentVolume == eOut && previousVolume == eEnvelope )
        {
          position += direction * fSurfaceTolerance; // Track out of surface
          time += fSurfaceTolerance / gvWater;
          previousVolume = currentVolume;
          currentVolume = eOut;
        }
      else
        break;

      //MCTrack Step here, not before
      RecordTrajectory( fastTrack, position, direction, 0.0, time, energy, fConcVolNames[previousVolume], fConcVolNames[currentVolume] );
    }

  if( iIterations == fMaxIterations )
    G4cout << "ConcentratorOpticalModel::TrackPhoton, photon is stuck" << G4endl;
}


// Reflect The photon, return true if reflects
bool
ConcentratorOpticalModel::Reflect( const G4ThreeVector& position,
                                   const G4double energy,
                                   G4ThreeVector& direction,
                                   G4ThreeVector& polarisation,
                                   G4double& weight )
{
  const G4ThreeVector norm = fPetalVolume->SurfaceNormal( position ).unit();

  polarisation = polarisation.unit();
  direction = direction.unit();
  // First check if reflects
  const G4double cosTheta = fabs(direction * norm);

  // Now get absorption & reflection probabilities
  const EPolarisation polarisationType = GetPolarisation( norm, direction, polarisation );

  G4double specularR, diffuseR;
  GetReflectionProb( position, acos( cosTheta ), energy, polarisationType, specularR, diffuseR );

  const G4double uniRand = G4UniformRand();
  if( uniRand > specularR )
    {
      weight = 0.0;
      return false;
    }
  else if( uniRand > specularR - diffuseR )
    {
      DiffuseReflect( norm, direction, polarisation );
      return true;
    }
  else
    {
      NormalReflect( norm, direction, polarisation );
      return true;
    }
}

void
ConcentratorOpticalModel::GetReflectionProb( const G4ThreeVector&,
                                             const G4double theta, // In Radians
                                             const G4double energy, // In MeV
                                             const EPolarisation polarisation, // As enum defined in OpticalModelBase.hh
                                             G4double& specularR,
                                             G4double& diffuseR )
{
  const G4double wavelength = twopi * hbarc / energy;
  const G4int iLocalTheta = static_cast<G4int>( theta / degree + 0.5 );
  G4int iWavelength = static_cast<G4int>( ( wavelength * mm ) / nm + 0.5 ); // Want in units of 1nm

  if( iWavelength < 305 )
    iWavelength = 305;
  else if( iWavelength > 800 )
    iWavelength = 800;

  iWavelength -= 305;
  if( polarisation == epPolarised )
    specularR = fReflectionPPol[ iLocalTheta + iWavelength * 91 ];
  else if( polarisation == esPolarised )
    specularR = fReflectionSPol[ iLocalTheta + iWavelength * 91 ];
  else
    Log::Die( "ConcentratorOpticalModel::GetReflectionProb: No polarisation" );

  diffuseR = fDiffuseReflectionProb;
}