#include "ReactionGenerator.hh"
#include "BeamConfig.hh"
#include "ClassTransfer.h"
#include "G4SystemOfUnits.hh"
#include "TRandom.h"
#include "TVector3.h"
#include "TLorentzVector.h"
#include <cmath>

struct ReactionGenerator::Impl {
  TransferReaction transferP;
  TransferReaction transferAlpha;
};

ReactionGenerator::ReactionGenerator()
: fImpl(nullptr),
  fEffectiveSigma(5e5 * barn) //originally 5e6
{
  // Ensure the isotope mass table is found when running from the AnasenG4 folder.
  massData = "/home/jamesszalkie/anasen/Armory/mass20.txt";
  fImpl = new Impl();

  // p-channel: 27Al + 4He -> p + 30Si
  fImpl->transferP.SetA(BeamConfig::kAtomicMass, BeamConfig::kAtomicNumber, 0);
  fImpl->transferP.Seta(4, 2);
  fImpl->transferP.Setb(1, 1);
  fImpl->transferP.SetB(30, 14, 0);

  // alpha-channel: 27Al + 4He -> alpha + 27Al
  fImpl->transferAlpha.SetA(BeamConfig::kAtomicMass, BeamConfig::kAtomicNumber, 0);
  fImpl->transferAlpha.Seta(4, 2);
  fImpl->transferAlpha.Setb(4, 2);
  fImpl->transferAlpha.SetB(BeamConfig::kAtomicMass, BeamConfig::kAtomicNumber, 0);
}

ReactionGenerator::~ReactionGenerator()
{
  delete fImpl;
}

ReactionOutput ReactionGenerator::SampleEvent(double beamKineticEnergy, const G4ThreeVector& beamDirection)
{
  // Forward-peaked sampling gives the expected inverse E(theta)-like trend in lab observables.
  double thetaCM = 0.0;
  while (true) {
    const double trial = gRandom->Rndm() * M_PI;
    const double weight = 1.0 / (1.0 + 20.0 * trial * trial);
    if (gRandom->Rndm() < weight) {
      thetaCM = trial;
      break;
    }
  }
  const double phiCM = gRandom->Rndm() * 2.0 * M_PI;

  const double beamEnergyPerU = beamKineticEnergy / (BeamConfig::kAtomicMass * MeV);
  const double beamThetaDeg = beamDirection.theta() / deg;
  const double beamPhiDeg = beamDirection.phi() / deg;

  // Use a mild energy-dependent branching to mix p and alpha events.
  const double protonFraction = (beamEnergyPerU > 8.0) ? 0.65 : 0.45;
  const bool isProtonChannel = (gRandom->Rndm() < protonFraction);

  TransferReaction* transfer = isProtonChannel ? &fImpl->transferP : &fImpl->transferAlpha;
  transfer->SetIncidentEnergyAngle(beamEnergyPerU, beamThetaDeg, beamPhiDeg);
  transfer->CalReactionConstant();

  TLorentzVector* reactionOutput = transfer->Event(thetaCM, phiCM);
  TLorentzVector light = reactionOutput[2];
  delete[] reactionOutput;

  ReactionOutput output;
  output.channel = isProtonChannel ? ReactionChannel::Proton : ReactionChannel::Alpha;
  output.kineticEnergy = light.E() - light.M();
  output.px = light.Px();
  output.py = light.Py();
  output.pz = light.Pz();
  return output;
}

double ReactionGenerator::GetEffectiveCrossSection() const
{
  return fEffectiveSigma;
}
