// VertexGen_IBD.cc
// Contact person: Nuno Barros <nfbarros@hep.upenn.edu>
// See VertexGen_IBD.hh for more details
//———————————————————————//
#include <sstream>

#include <globals.hh>

#include <G4ParticleDefinition.hh>
#include <G4ParticleTable.hh>
#include <G4PrimaryVertex.hh>
#include <G4PrimaryParticle.hh>
#include <G4Event.hh>
#include <G4PhysicalConstants.hh>


#include <Randomize.hh>
#include <RAT/GLG4PosGen.hh>
#include <RAT/VertexGen_IBD.hh>
#include <RAT/GLG4StringUtil.hh>
#include <RAT/ParentPrimaryVertexInformation.hh>
#include <RAT/IBDCrossSec.hh>
#include <RAT/DB.hh>
#include <RAT/Log.hh>
#include <RAT/NeutrinoVertexGen.hh>

#include <TGraph.h>
#include <TH1D.h>


using CLHEP::twopi;
using CLHEP::electron_mass_c2;
using CLHEP::proton_mass_c2;
using CLHEP::neutron_mass_c2;
using CLHEP::keV;
using CLHEP::MeV;

namespace RAT {


const double VertexGen_IBD::sDeltaM = neutron_mass_c2 - proton_mass_c2;

VertexGen_IBD::VertexGen_IBD(const char *arg_dbname)
  : NeutrinoVertexGen(arg_dbname), fNuDir(0.,0.,0.), fCrossSecMax(0.0),
    fFluxMax(0.0), fFlux(0), fSpectrumIndex("Reactornus"),
    fSpectrumLoaded(false), fStandaloneMode(true)
#ifdef RATDEBUG
, fProfiler("VertexIBD","Vertex")
#endif
{
  fCrossSec = new IBDCrossSec();
  fANu = G4ParticleTable::GetParticleTable()->FindParticle("anti_nu_e");
  fNeutron = G4ParticleTable::GetParticleTable()->FindParticle("neutron");
  fEplus = G4ParticleTable::GetParticleTable()->FindParticle("e+");
}

VertexGen_IBD::~VertexGen_IBD()
{
  delete fCrossSec;
}

void VertexGen_IBD::
GeneratePrimaryVertex(G4Event *argEvent,
          G4ThreeVector &dx,
          G4double dt,
          G4ThreeVector &nu_dir)
{
#ifdef RATDEBUG
  fProfiler.Start();
#endif
  fNuDir = nu_dir;
  GeneratePrimaryVertex(argEvent,dx,dt);
#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::GeneratePrimaryVertex : Returning an event " << argEvent << newline;
  fProfiler.Stop();
  #endif
}

void VertexGen_IBD::
GeneratePrimaryVertex(G4Event *argEvent,
		      G4ThreeVector &dx,
		      G4double dt)
{
#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::GeneratePrimaryVertex : Generating the IBD vertex." << newline;
#endif
  G4PrimaryVertex* vertex= new G4PrimaryVertex(dx, dt);
  G4ThreeVector ev_nu_dir(fNuDir); // By default use specified direction

#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::GeneratePrimaryVertex : Vertex generated at "
      << dx.x() << " , "<< dx.y() << " , "<< dx.z() << "  and t0 : " << dt << newline;
  detail << "[VertexGen_IBD]::GeneratePrimaryVertex : Direction "
      << ev_nu_dir.x() << " , "<< ev_nu_dir.y() << " , "<< ev_nu_dir.z() << newline;
#endif
  if (ev_nu_dir.mag2() == 0.0) {
    // Pick isotropic direction
    double theta = acos(2.0 * G4UniformRand() - 1.0);
    double phi = G4UniformRand() * twopi;
    ev_nu_dir.setRThetaPhi(1.0, theta, phi);
  }

  // -- Generate inverse beta decay interaction
  G4LorentzVector mom_nu, mom_eplus, mom_n;

  GenEvent(ev_nu_dir, mom_nu, mom_eplus, mom_n);

  // -- Create particles
  G4PrimaryParticle* antinuParent =
    new G4PrimaryParticle(fANu,
			  mom_nu.px(),     // x component of momentum
			  mom_nu.py(),     // y component of momentum
			  mom_nu.pz());    // z component of momentum
  // positron
#ifdef RATVERBOSE
  detail << "[VertexGen_IBD]::GeneratePrimaryVertex : nuebar details." << newline;
  detail << " Px : " << mom_nu.px() << " Py : " << mom_nu.py() << " Pz : " << mom_nu.pz() << newline;
#endif

  G4PrimaryParticle* eplus_particle =
    new G4PrimaryParticle(fEplus,              // particle code
			  mom_eplus.px(),     // x component of momentum
			  mom_eplus.py(),     // y component of momentum
			  mom_eplus.pz());    // z component of momentum
  eplus_particle->SetMass(fEplus->GetPDGMass()); // Geant4 is silly.

#ifdef RATVERBOSE
  detail << "[VertexGen_IBD]::GeneratePrimaryVertex : e+ details." << newline;
  detail << " Px : " << mom_eplus.px() << " Py : " << mom_eplus.py() << " Pz : " << mom_eplus.pz() << newline;
#endif

  vertex->SetPrimary( eplus_particle );

  // neutron
  G4PrimaryParticle* n_particle =
    new G4PrimaryParticle(fNeutron,                  // particle code
			  mom_n.px(),         // x component of momentum
			  mom_n.py(),         // y component of momentum
			  mom_n.pz());        // z component of momentum
  n_particle->SetMass(fNeutron->GetPDGMass()); // Geant4 is silly.
#ifdef RATVERBOSE
  detail << "[VertexGen_IBD]::GeneratePrimaryVertex : n details." << newline;
  detail << " Px : " << mom_n.px() << " Py : " << mom_n.py() << " Pz : " << mom_n.pz() << newline;
#endif

  vertex->SetPrimary( n_particle );


  // We Don't add this one to the vertex as we don't want to propagate it but alongside
  // so that the information is present for extraction in Gsim
  ParentPrimaryVertexInformation *vertinfo = new ParentPrimaryVertexInformation();
  vertinfo->SetVertexParentParticle(antinuParent);
  // The name of the reactor that gave origin to this vertex is added somewhere else

  vertex->SetUserInformation(vertinfo);
  argEvent->AddPrimaryVertex(vertex);
}


void VertexGen_IBD::SetState(G4String newValues)
{
  newValues = util_strip_default(newValues); // from GLG4StringUtil
  if (newValues.length() == 0) {
    // print help and current state
    warn << "Current state of this VertexGen_IBD:\n"
        << " \"" << GetState() << "\"\n" << newline;
    warn << "Format of argument to VertexGen_IBD::SetState: \n"
      " \"fNuDir_x fNuDir_y fNuDir_z fSpectrumIndex\"\n"
      " where fNuDir is the initial direction of the antineutrino and\n"
      " fSpectrumIndex is the index of the spectrum in the IBD type table.\n"
      " The direction does not have to be normalized. Set to \"0. 0. 0.\" for\n"
      " isotropic neutrino direction."
        << newline;
    return;
  }

  std::istringstream is(newValues.c_str());
  double x, y, z;
  is >> x >> y >> z >> fSpectrumIndex;
  if (is.fail())
    return;

  if ( x == 0. && y == 0. && z == 0. )
    fNuDir.set(0., 0., 0.);
  else
    fNuDir = G4ThreeVector(x, y, z).unit();

  SetSpectrum();
}


G4String VertexGen_IBD::
GetState()
{
  std::ostringstream os;

  os << fNuDir.x() << "\t" << fNuDir.y() << "\t" << fNuDir.z() << std::ends;

  G4String rv(os.str());
  return rv;
}


void VertexGen_IBD::GenEvent(const G4ThreeVector &nu_dir,
          G4LorentzVector &neutrino,
          G4LorentzVector &positron,
          G4LorentzVector &neutron)
{
  double Enu, CosThetaLab;

#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::GenEvent : Generating interaction." << newline;
#endif

  // Pick energy of neutrino and relative direction of positron
  GenInteraction(Enu, CosThetaLab);

#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::GenEvent : Interaction details: Enu : " << Enu/MeV << " CosTheta " << CosThetaLab << newline;
#endif

  // Zero'th order approximation of positron quantities (infinite nucleon mass)
  double E0 = Enu - sDeltaM;
  double p0 = sqrt(E0*E0-electron_mass_c2*electron_mass_c2);
  double v0 = p0/E0;
  // First order correction for finite nucleon mass
  double Ysquared = (sDeltaM*sDeltaM-electron_mass_c2*electron_mass_c2)/2;
  double E1 = E0*(1-Enu/proton_mass_c2*(1-v0*CosThetaLab))
               - Ysquared/proton_mass_c2;
  double p1 = sqrt(E1*E1-electron_mass_c2*electron_mass_c2);

  // Compute nu 4-momentum
  neutrino.setVect(nu_dir * Enu); // MeV (divide by c if need real units)
  neutrino.setE(Enu);

  // Compute positron 4-momentum
  G4ThreeVector pos_momentum(p1*nu_dir);

  // Rotation from nu direction to pos direction.
  double theta = acos(CosThetaLab);
  double phi = twopi*G4UniformRand();  // Random phi
  G4ThreeVector rotation_axis = nu_dir.orthogonal();
  rotation_axis.rotate(phi, nu_dir);
  pos_momentum.rotate(theta, rotation_axis);

  positron.setVect(pos_momentum);
  positron.setE(E1);

  // Compute neutron 4-momentum
  neutron.setVect(neutrino.vect() - positron.vect());
  neutron.setE(sqrt(neutron.vect().mag2() + neutron_mass_c2*neutron_mass_c2));

#ifdef RATDEBUG
  info << "VertexGen_IBD::GenEvent: Enu = " << Enu/keV << " keV Epos=" << E1/keV << " keV Diff=" << (Enu-E1)/keV << newline;
#endif
}

void VertexGen_IBD::GenInteraction(double &E, double &CosThetaLab)
{

  bool passed=false;

#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::GenInteraction : Generating IBD interaction with the following params: " << newline;
  detail << "Emin         : " << fEmin << newline;
  detail << "Emax         : " << fEmax << newline;
  detail << "CrossSecMax  : " << fCrossSecMax << newline;
  detail << "FluxMax      : " << fFluxMax << newline;

#endif
  // The issue here is that there are different spectra that have to be sampled from based on
  // which
  while(!passed){
    // Pick E and cos(theta) uniformly
    E = fEmin+(fEmax-fEmin)*G4UniformRand();
    CosThetaLab = -1.0+2.0*G4UniformRand();

    // Decided whether to draw again based on relative cross-section.
    double XCtest = fCrossSecMax * fFluxMax * G4UniformRand();
    double XCWeight = fCrossSec->dSigmadTheta(E,CosThetaLab);
    double FluxWeight = fFlux->Interpolate(E);
    passed = (XCWeight * FluxWeight) > XCtest;
  }
#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::GenInteraction : Returning E= " << E/MeV << " CosTheta " << CosThetaLab << newline;
#endif
}

void VertexGen_IBD::BeginOfRun() {
  if (fStandaloneMode) {
    warn << "*********************************************************" << endl;
    warn << "* !!     Running VertexGen_IBD in standalone mode     !!*" << endl;
    warn << "*                                                       *" << endl;
    warn << "* There is nothign wrong with this, provided that you   *" << endl;
    warn << "* know what you're doing. Keep in mind that in          *" << endl;
    warn << "* standalone mode, the generator grabs the spectrum from*" << endl;
    warn << "* the IBD RATDB table. Either modify that ratdb file for*" << endl;
    warn << "* your needs or change its contents in the macro file.  *" << endl;
    warn << "*********************************************************" << endl;

    // Running in standalone mode.
    // Load the spectrum information directly from the IBD database
    // entry
    if (!fSpectrumLoaded) {
      SetSpectrum();
    }
  } else if (!fSpectrumLoaded) {
    Log::Die("[VertexGen_IBD]::BeginOfRun: Antinu spectrum hasn't been loaded yet.");
  } else {
    // there's nothing to be done if the spectrum is being used with the ReactorGen
    // Just check that the necessary inputs are in place
    // This can actually cause a race condition, since it is not totally clear the order
    // the factories will call for BeginOfRun
    if (!fFlux) {
      Log::Die("[VertexGen_IBD]::BeginOfRun: Neutrino spectrum has not been set yet.",5000);
    }
  }

  info << "[VertexGen_IBD]::BeginOfRun: IBD spectrum index: " << fSpectrumIndex << newline;
}

void VertexGen_IBD::EndOfRun() {
  if (fStandaloneMode) {
    if (fFlux) {
      delete fFlux;
      fFlux = NULL;
    }
  }

  #ifdef RATDEBUG
  info << "------------------------------------------------------------" << endl;
  info << "VertexGen_IBD timing profile :" << endl;
  info << "------------------------------------------------------------" << endl;
  fProfiler.OutputTiming();
  info << "------------------------------------------------------------" << endl;
#endif
}

void VertexGen_IBD::SetSpectrum() {
  try {
    DBLinkPtr libd = DB::Get()->GetLink("IBD", fSpectrumIndex);

    fEmin = libd->GetD("emin");
    fEmax = libd->GetD("emax");

    // Flux function
    fSpectrumX = libd->GetDArray("spec_e");
    fSpectrumY = libd->GetDArray("spec_flux");
    Log::Assert(fSpectrumX.size() == fSpectrumY.size(),"VertexGen_IBD::BeginOfRun : Axes of input spectrum have mismatching dimensions.");

    double dstep = fSpectrumX.at(1)- fSpectrumX.at(0);
    fFlux = new TH1D("spectrum","",fSpectrumX.size(),fSpectrumX.at(0),fSpectrumX.at(fSpectrumX.size()-1)+dstep);
    for (size_t ibin = 1; ibin < fSpectrumX.size()+1; ++ibin) {
      fFlux->SetBinContent(ibin,fSpectrumY.at(ibin-1));
    }
    // this should be converted to a histogram...

    // Other useful numbers
    fCrossSecMax = fCrossSec->dSigmadTheta(fEmax,-1.0);

    // This assumption is dangerous.
    // We could be using some weird antineutrino spectrum that does not peak at the
    // low boundary
    // I guess it would be more appropriate to actually look specifically for the maximum

    fFluxMax = fFlux->GetMaximum();

    fSpectrumLoaded = true;
  }
  catch (RAT::DBNotFoundError &e) {
    RAT::Log::Die(dformat("VertexGen_IBD::SetSpectrum : Error loading contents from the database : %s",e.what()));
  }
}

void VertexGen_IBD::SetSpectrum(TH1D *flux) {

#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::SetSpectrum : Setting the neutrino flux : " << newline;
  detail << "[VertexGen_IBD]::SetSpectrum : " << flux->GetName() << newline;
#endif

  // if it exists, get rid of it
  if (fFlux) {
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::SetSpectrum : Clearing out old spectrum:  " << newline;
    detail << "[VertexGen_IBD]::SetSpectrum : " << fFlux->GetName() << newline;
#endif
    delete fFlux;
  }
  // If we want to be able to clear out, then we have to keep ownership.
  // In that case it is best to clone
  fFlux = dynamic_cast<TH1D*>(flux->Clone(Form("%s_Vtx_IBD",flux->GetName())));

  // Check that the integral is 1. If it isn't make it so
  if (fFlux->Integral("width") != 1.0) {
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::SetSpectrum : Renormalizing the spectrum." << newline;
#endif

    TH1D* tmph = fFlux;
    fFlux = (TH1D*)tmph->Clone(Form("%s_norm",tmph->GetName()));
    fFlux->Scale(1.0/fFlux->Integral("width"));
    delete tmph; // Delete the original
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::SetSpectrum : Spectrum renormalized." << newline;
    detail << "[VertexGen_IBD]::SetSpectrum : name " << fFlux->GetName() << " norm " << fFlux->Integral("width") << newline;
#endif
  }
  fEmin = flux->GetXaxis()->GetBinLowEdge(1);
  fEmax = flux->GetXaxis()->GetBinLowEdge(flux->GetXaxis()->GetNbins());
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::SetSpectrum : Emin " << fEmin/CLHEP::MeV << " MeV, Emax " << fEmax/CLHEP::MeV << " MeV" << newline;
    detail << "[VertexGen_IBD]::SetSpectrum : Estimating maximum of cross section." << newline;
    detail << "[VertexGen_IBD]::SetSpectrum : Pointer: " << fCrossSec << newline;
#endif
  fCrossSecMax = fCrossSec->dSigmadTheta(fEmax,-1.0);
#ifdef RATDEBUG
  detail << "[VertexGen_IBD]::SetSpectrum : value " << fCrossSecMax << newline;
#endif
  fFluxMax = fFlux->GetMaximum();
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::SetSpectrum : Maximum Flux set to " << fFluxMax << newline;
#endif

  fSpectrumLoaded = true;
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::SetSpectrum : Flux set and normalized." << newline;
#endif

}

G4double VertexGen_IBD::GetRatePerTarget() {
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::GetRatePerTarget : Convolving cross section with neutrino spectrum." << newline;
#endif

  if (!fSpectrumLoaded) {
    Log::Die("[VertexGen_IBD]::BeginOfRun: Antinu spectrum hasn't been loaded yet.",1);
  }
  double rate = 0.0;

  for (int ibin = 1; ibin < fFlux->GetXaxis()->GetNbins(); ++ibin) {
    rate += fFlux->GetXaxis()->GetBinWidth(ibin)*fFlux->GetBinContent(ibin)*fCrossSec->Sigma(fFlux->GetXaxis()->GetBinCenter(ibin));
  }
#ifdef RATDEBUG
    detail << "[VertexGen_IBD]::GetRatePerTarget : Returning a rate of " <<  rate << newline;
#endif

  return rate;
}
} // namespace RAT