#include "BeamReactionSteppingAction.hh"

#include "ReactionGenerator.hh"
#include "WireTrackingManager.hh"

#include "G4Step.hh"
#include "G4Track.hh"
#include "G4DynamicParticle.hh"
#include "G4EventManager.hh"
#include "G4StackManager.hh"
#include "G4ParticleDefinition.hh"
#include "G4Proton.hh"
#include "G4Alpha.hh"
#include "G4SystemOfUnits.hh"
#include "G4RandomTools.hh"

#include <cmath>

BeamReactionSteppingAction::BeamReactionSteppingAction()
: G4UserSteppingAction(),
  fReaction(new ReactionGenerator()),
  fTargetVolumeName("Target"),
  fBeamZ(13),
  fBeamA(27)
{}

BeamReactionSteppingAction::~BeamReactionSteppingAction()
{
  delete fReaction;
}

void BeamReactionSteppingAction::UserSteppingAction(const G4Step* step)
{
  WireTrackingManager::Instance().UpdateForStep(step);

  const G4Track* track = step->GetTrack();
  if (track->GetTrackStatus() != fAlive) {
    return;
  }

  const G4ParticleDefinition* particle = track->GetDefinition();
  if (particle->GetParticleType() != "nucleus") {
    return;
  }

  if (particle->GetAtomicNumber() != fBeamZ || particle->GetAtomicMass() != fBeamA) {
    return;
  }

  const auto* prePoint = step->GetPreStepPoint();
  if (!prePoint || !prePoint->GetPhysicalVolume()) {
    return;
  }

  if (prePoint->GetPhysicalVolume()->GetName() != fTargetVolumeName) {
    return;
  }

  const G4double stepLength = step->GetStepLength();
  if (stepLength <= 0.0) {
    return;
  }

  const auto* material = prePoint->GetMaterial();
  if (!material) {
    return;
  }

  const G4double atomDensity = material->GetTotNbOfAtomsPerVolume();
  const G4double sigma = fReaction->GetEffectiveCrossSection();
  const G4double probability = 1.0 - std::exp(-atomDensity * sigma * stepLength);

  if (G4UniformRand() > probability) {
    return;
  }

  ReactionOutput output = fReaction->SampleEvent(track->GetKineticEnergy(), track->GetMomentumDirection());
  if (output.kineticEnergy <= 0.0) {
    return;
  }

  const G4ThreeVector momentum(output.px, output.py, output.pz);
  G4ThreeVector direction = momentum;
  if (direction.mag2() > 0.0) {
    direction = direction.unit();
  } else {
    direction = G4ThreeVector(0.0, 0.0, 1.0);
  }

  G4ParticleDefinition* product = nullptr;
  if (output.channel == ReactionChannel::Proton) {
    product = G4Proton::ProtonDefinition();
  } else {
    product = G4Alpha::AlphaDefinition();
  }

  auto* dynamicParticle = new G4DynamicParticle(product, direction, output.kineticEnergy);
  auto* secondaryTrack = new G4Track(dynamicParticle, track->GetGlobalTime(), prePoint->GetPosition());
  secondaryTrack->SetParentID(track->GetTrackID());
  secondaryTrack->SetTouchableHandle(track->GetTouchableHandle());

  G4EventManager::GetEventManager()->GetStackManager()->PushOneTrack(secondaryTrack, nullptr);

  G4Track* mutableTrack = const_cast<G4Track*>(track);
  mutableTrack->SetTrackStatus(fStopAndKill);
}