#include "PrimaryGeneratorAction.hh"
#include "BeamConfig.hh"
#include "G4ParticleGun.hh"
#include "G4ParticleTable.hh"
#include "G4IonTable.hh"
#include "G4Proton.hh"
#include "G4SystemOfUnits.hh"
#include "G4Event.hh"
#include "G4PhysicalConstants.hh"
#include "G4RandomTools.hh"
#include "G4UnitsTable.hh"

#include <cmath>

PrimaryGeneratorAction::PrimaryGeneratorAction()
: G4VUserPrimaryGeneratorAction(),
  fParticleGun(nullptr),
  fBeamRadius(BeamConfig::kRadius),
  fBeamStartZ(BeamConfig::kStartZ)
{
  G4int nParticle = 1;
  fParticleGun = new G4ParticleGun(nParticle);
  fParticleGun->SetParticleDefinition(G4Proton::ProtonDefinition());
  fParticleGun->SetParticleMomentumDirection(G4ThreeVector(0., 0., 1.));
  fParticleGun->SetParticleEnergy(BeamConfig::kTotalBeamEnergy);
}

PrimaryGeneratorAction::~PrimaryGeneratorAction()
{
  delete fParticleGun;
}

void PrimaryGeneratorAction::GeneratePrimaries(G4Event* event)
{
  auto* ion = G4IonTable::GetIonTable()->GetIon(BeamConfig::kAtomicNumber, BeamConfig::kAtomicMass, 0.0);
  if (ion) {
    fParticleGun->SetParticleDefinition(ion);
  }

  G4double r = fBeamRadius * std::sqrt(G4UniformRand());
  G4double phi = twopi * G4UniformRand();
  G4double x = r * std::cos(phi);
  G4double y = r * std::sin(phi);
  G4double z = fBeamStartZ;

  fParticleGun->SetParticlePosition(G4ThreeVector(x, y, z));
  fParticleGun->SetParticleMomentumDirection(G4ThreeVector(0., 0., 1.));
  fParticleGun->SetParticleEnergy(BeamConfig::kTotalBeamEnergy);
  fParticleGun->GeneratePrimaryVertex(event);
}
