#include <G4ParticleTable.hh>
#include <G4PrimaryVertex.hh>
#include <G4PrimaryParticle.hh>
#include <Randomize.hh>
#include <G4Event.hh>

#include <RAT/HitPMTCollection.hh>

#include <vector>

#include "ExampleMiniSim.hh"

namespace RAT
{

ExampleMiniSim::ExampleMiniSim(const G4ThreeVector &center,
                               double betaEnergy)
  : fCenter(center), fBetaEnergy(betaEnergy)
{
  // maybe load some RATDB tables
}

ExampleMiniSim::~ExampleMiniSim() {
  delete fNhits;
}

void ExampleMiniSim::Run(const int nevents)
{
  const int nchannels = 19 * 16 * 32;
  fNhits = new TH1F("nhits", "nhits", nchannels, 0, nchannels);

  // MiniSim::BeamOn takes over control and runs RAT
  BeamOn(nevents);

  fMeanNhits = fNhits->GetMean();
}

void ExampleMiniSim::GeneratePrimaries(G4Event* event)
{
  double dt = 0.0;

  G4ParticleDefinition* particle = G4ParticleTable::GetParticleTable()->FindParticle("e-");
  G4ThreeVector dx = fCenter;

  // random polarization
  double px = (G4UniformRand() * 2.0 - 1.0);
  double py = (G4UniformRand() * 2.0 - 1.0);
  double pz = (G4UniformRand() * 2.0 - 1.0);
  G4ThreeVector mom(px, py, pz);
  mom *= sqrt(fBetaEnergy * (fBetaEnergy + 2.0 * particle->GetPDGMass()));

  G4PrimaryVertex* vertex = new G4PrimaryVertex(dx, dt);
  G4PrimaryParticle* beta = new G4PrimaryParticle(particle, mom.x(), mom.y(), mom.z());

  // random polarization
  double phi = (G4UniformRand() * 2.0 - 1.0) * pi;
  G4ThreeVector e1 = mom.orthogonal().unit();
  G4ThreeVector e2 = mom.unit().cross(e1);
  G4ThreeVector pol = e1*cos(phi)+e2*sin(phi);
  beta->SetPolarization(pol.x(), pol.y(), pol.z());

  vertex->SetPrimary(beta);
  event->AddPrimaryVertex(vertex);
}

void ExampleMiniSim::BeginOfEventAction(const G4Event* /*event*/)
{
  HitPMTCollection::Get()->BeginOfEvent();
}

void ExampleMiniSim::EndOfEventAction(const G4Event* /*event*/)
{
  const std::vector<DS::MCPMT*>& hitPMTs = HitPMTCollection::Get()->GetSortedPMTCollection();

  fNhits->Fill(hitPMTs.size());
}

} // namespace RAT