// VertexGen_Spectrum.cc
// Contact person:Jeanne Wilson <j.r.wilson@qmul.ac.uk>
// See VertexGen_Spectrum.hh for more details
//———————————————————————//
#include <RAT/VertexGen_Spectrum.hh>
#include <RAT/Log.hh>
#include <Randomize.hh>
#include <G4Event.hh>
#include <G4ParticleTable.hh>
#include <G4IonTable.hh>
#include <G4PrimaryParticle.hh>
#include <G4PrimaryVertex.hh>
#include <CLHEP/Units/PhysicalConstants.h>
#include <G4ThreeVector.hh>

using namespace CLHEP;
#define G4std  std
namespace RAT {
VertexGen_Spectrum::VertexGen_Spectrum(const char *arg_dbname) :
    GLG4VertexGen(arg_dbname) {
  // constructor
  fParticle = "e-";
  fSpectrum = "flat";
  fLSpec = DB::Get()->GetLink("SPECTRUM", fSpectrum); // default
  fEMin = 0;  // no valid range for spectrum until properly initialised
  fEMax = 0;
  fELimUniLow = 0;
  fELimUniHigh = 0;
  fELimTempLow = 0;
  fELimTempHigh = 0;
}

///-------------------------------------------------------------------------
VertexGen_Spectrum::~VertexGen_Spectrum() {
  // destructor
}

void VertexGen_Spectrum::BeginOfRun() {
  fLSpec = DB::Get()->GetLink("SPECTRUM", fSpectrum); // default
  InitialiseSpectrum();
}

///-------------------------------------------------------------------------
void VertexGen_Spectrum::GeneratePrimaryVertex(G4Event *event,
    G4ThreeVector &dx, G4double dt) {
  // Where the main work is done - the astute may notice a strong similarity to the Gun generator!

  G4PrimaryVertex* vertex = new G4PrimaryVertex(dx, dt);
  G4PrimaryParticle* particle;

  fPDef = G4ParticleTable::GetParticleTable()->FindParticle(fParticle);
  // particle mass
  G4double mass = fPDef->GetPDGMass();

  G4double _ke = this->SampleEnergy();
  // after generating the event we must reset any temporary limits on the energy range
  fELimTempLow = fELimUniLow;
  fELimTempHigh = fELimUniHigh;

  // isotropic direction
  G4double phi = twopi * G4UniformRand();
  G4double cosTheta = -1. + 2. * G4UniformRand();
  G4double sinTheta = sqrt(1. - cosTheta * cosTheta);
  G4double ux = sinTheta * cos(phi);
  G4double uy = sinTheta * sin(phi);
  G4double uz = cosTheta;

  G4ThreeVector dir = G4ThreeVector(ux, uy, uz);

  G4ThreeVector rmom(dir * sqrt(_ke * (_ke + 2. * mass)));
  particle = new G4PrimaryParticle(fPDef,     // particle code
      rmom.x(),                               // x component of momentum
      rmom.y(),                               // y component of momentum
      rmom.z());                              // z component of momentum
  particle->SetMass(mass); // Geant4 is silly.

  // random polarisation
  phi = (G4UniformRand()*2.0-1.0)*pi;
  G4ThreeVector e1 = dir.orthogonal().unit();
  G4ThreeVector e2 = dir.cross(e1);
  G4ThreeVector rpol = e1 * cos(phi) + e2 * sin(phi);
  particle->SetPolarization(rpol.x(), rpol.y(), rpol.z());

  vertex->SetPrimary(particle);

  event->AddPrimaryVertex(vertex);
}

///-------------------------------------------------------------------------
void VertexGen_Spectrum::SetState(G4String newValues) {
  if (newValues.length() == 0) {
    // print help and current state
    info << "Current state of this VertexGen_Spectrum:\n" << " \"" << GetState()
        << "\"\n";
    info << "Format of argument to VertexGen_Spectrum::SetState: \n"
        << " \"pname  specname  (Elo Ehi)\"\n" << " pname = particle name \n"
        << " specname = spectrum name as given in ratdb \n"
        << " Elo Ehi = optional limits on energy range of generated particles \n";
    return;
  }

  std::istringstream is(newValues.c_str());
  G4std::string pname, specname;
  // Read in the particle type
  is >> pname;
  if (is.fail() || pname.length() == 0) {
    Log::Die("VertexGen_Spectrum: Incorrect vertex setting " + newValues);
  }
  G4ParticleDefinition * newTestGunG4Code =
      G4ParticleTable::GetParticleTable()->FindParticle(pname);
  if (newTestGunG4Code == NULL) {
    // not a particle name
    // see if we can parse it as an ion, e.g., U238 or Bi214
    G4std::string elementName;
    int A, Z;
    if (pname[1] >= '0' && pname[1] <= '9') {
      A = atoi(pname.substr(1).c_str());
      elementName = pname.substr(0, 1);
    } else {
      A = atoi(pname.substr(2).c_str());
      elementName = pname.substr(0, 2);
    }
    if (A > 0) {
      for (Z = 1; Z <= GLG4VertexGen_Gun::fNumberOfElements; Z++) {
        if (elementName == GLG4VertexGen_Gun::fTheElementNames[Z - 1])
          break;
        if (Z <= GLG4VertexGen_Gun::fNumberOfElements) {
          newTestGunG4Code = G4IonTable::GetIonTable()->GetIon(Z, A, 0.0);
          info << " Spectrum Vertex: Setting ion with A = " << A << " Z = " << Z
              << "\n";
        }
      }
    }
    if (newTestGunG4Code == NULL) {
      warn << "Spectrum Vertex: Could not find particle type " << pname
          << " defaulting to electron \n";
      fParticle = "e-";
      return;
    }
  } else {
    info << "Spectrum Vertex: Setting particle = " << pname << "\n";
  }
  // so store the name and the particle definition
  fParticle = pname;
  fPDef = newTestGunG4Code;

  // Read in the spectrum name
  is >> specname;
  // check that spectrum is in database - not sure what happens if this fails - tidy!
  DBLinkPtr lspec = DB::Get()->GetLink("SPECTRUM", specname);
  if (lspec) {
    fLSpec = lspec;
    fSpectrum = specname;
    info << "Spectrum Vertex: Setting spectrum " << specname << "\n";
  } else {
    warn << "Could not find spectrum " << specname
        << " using default, flat spectrum \n";
  }

  // Has the user supplied energy limits too?
  float Elo, Ehi;
  is >> Elo >> Ehi;
  if (is.fail()) {
    // obviously not, so make the universal limits huge
    fELimUniLow = 0;
    fELimUniHigh = 9999999;
  } else {
    // user has selected universal limits
    fELimUniLow = Elo;
    fELimUniHigh = Ehi;
    info << "Limiting spectrum to range " << fELimUniLow << " - "
        << fELimUniHigh << "\n";
  }

  // finally ready to initialise the spectrum!
  this->InitialiseSpectrum();
  return;
}

///-------------------------------------------------------------------------
G4String VertexGen_Spectrum::GetState() {
  // State setting is the particle name and spectrum name - return it
  G4std::ostringstream os;
  os << fParticle << " " << fSpectrum << G4std::ends;
  G4String rv(os.str());
  return rv;
}

///-------------------------------------------------------------------------
void VertexGen_Spectrum::InitialiseSpectrum() {
  // Initialise the spectrum
  // clear any stored values
  fSpecE.clear();
  fSpecCumMag.clear();
  // firstly get the energies and store the endpoints
  fSpecE = fLSpec->GetFArrayFromD("spec_e");
  fEMin = fSpecE.front();
  fEMax = fSpecE.back();
  if (fEMax - fEMin <= 0 || fEMin < 0 || fEMax < 0) {
    // We have a problem with the energy range
    Log::Die(
        "VertexGen_Spectrum: Nonsensical energy range for spectrum "
            + fSpectrum);
  }

  // now get the magnitudes
  fSpecMag = fLSpec->GetFArrayFromD("spec_mag");
  // now create a cumulative magnitude vector
  float magsum = 0;
  fSpecCumMag.push_back(0.0);
  for (unsigned int istep = 1; istep < fSpecE.size(); ++istep) {
    // use trapezoid rule to produce linear interpolation
    // of spectrum
    magsum += (fSpecMag[istep] + fSpecMag[istep - 1]) / 2.0
        * (fSpecE[istep] - fSpecE[istep - 1]);
    fSpecCumMag.push_back(magsum);
  }

  // check that this works with the universal energy limits
  if ((fEMax < fELimUniLow) || (fEMin > fELimUniHigh)) {
    Log::Die(
        "VertexGen_Spectrum: selected spectrum does not lie within chosen energy range: "
            + fSpectrum);
  }

  // and finally make sure temporary limits are set to universal ones
  fELimTempLow = fELimUniLow;
  fELimTempHigh = fELimUniHigh;
  info << "Spectrum " << fSpectrum << " initialised \n";
}

///-------------------------------------------------------------------------
float VertexGen_Spectrum::SampleEnergy() {
  // Return a value for KE sampled from the spectrum between the limits Elo and Ehi
  float Elo, Ehi;
  // decide whether, Elo and Ehi should be universal or temporary limits (which is more stringent?)
  if (fELimTempLow > fELimUniLow) {
    Elo = fELimTempLow;
  } else {
    Elo = fELimUniLow;
  }
  if (fELimTempHigh < fELimUniHigh) {
    Ehi = fELimTempHigh;
  } else {
    Ehi = fELimUniHigh;
  }
  // (assume linear interpolation between specified points in spectrum)
  float start = fSpecCumMag.front();
  float stop = fSpecCumMag.back();
  if (Elo > fEMin) {
    // We have to apply a more stringent lower energy limit
    int istep = 0;
    while (fSpecE[istep] < Elo) {
      start = fSpecCumMag[istep];
      istep++;
    }
    // and interpolate last bit
    start += ((Elo - fSpecE[istep - 1]) / (fSpecE[istep] - fSpecE[istep - 1]))
        * (fSpecCumMag[istep] - fSpecCumMag[istep - 1]);
  }
  if (Ehi < fEMax) {
    int istep = fSpecE.size() - 1;
    while (fSpecE[istep] > Ehi) {
      stop = fSpecCumMag[istep];
      detail << "* " << istep << " " << stop << " " << fSpecE[istep] << " "
          << Ehi << "\n";
      istep--;
    }
    // and interpolate the last bit
    stop -= ((fSpecE[istep + 1] - Ehi) / (fSpecE[istep + 1] - fSpecE[istep]))
        * (fSpecCumMag[istep + 1] - fSpecCumMag[istep]);
  }

  // now throw a random number between these limits
  float random = start + G4UniformRand()*(stop-start);
  // now loop through cumulative distribution to choose energy
  int istep = 0;
  while (fSpecCumMag[istep + 1] < random) {
    istep++;
  }

  // The desired energy is bracketed by istep and istep+1 points.
  // Since the cumulative function is an integral of a linearly
  // interpolated spectrum, it is quadratic between data points.
  // Find energy assuming quadratic shape.

  float spec_slope = (fSpecMag[istep + 1] - fSpecMag[istep])
      / (fSpecE[istep + 1] - fSpecE[istep]);

  float energy;
  float delta_y = random - fSpecCumMag[istep];
  if (fabs(spec_slope) < 1e-6) {
    energy = fSpecE[istep] + delta_y / fSpecMag[istep];
  } else {
    float a = 0.5 * spec_slope;
    float b = fSpecMag[istep] - spec_slope * fSpecE[istep];
    float c = -a * fSpecE[istep] * fSpecE[istep] - b * fSpecE[istep] - delta_y;
    energy = (-b + sqrt(b * b - 4 * a * c)) / (2 * a);
  }

  return energy;
}

///-------------------------------------------------------------------------
void VertexGen_Spectrum::LimitEnergies(float Elo, float Ehi) {
  // Set the limits for the generated energy range
  // first check this makes sense
  if ((Elo > fEMax) || (Elo > fELimUniHigh) || (Ehi < fEMin)
      || (Ehi < fELimUniLow) || ((Ehi - Elo) <= 0)) {
    warn << "Spectrum Vertex: temporary energy limits " << Elo << " - " << Ehi
        << " don't make sense, not applied \n" << "spectrum in range " << fEMin
        << " - " << fEMax << " MeV with universal limits " << fELimUniLow
        << " - " << fELimUniHigh << "\n";
    return;
  }
  fELimTempLow = Elo;
  fELimTempHigh = Ehi;
  return;
}

///-------------------------------------------------------------------------
float VertexGen_Spectrum::EMaximum() {
  // return the maximum possible energy - accounting for spectrum and universal limits on it
  if (fEMax < fELimUniHigh) {
    return fEMax;
  } else {
    return fELimUniHigh;
  }
}

///-------------------------------------------------------------------------
float VertexGen_Spectrum::EMinimum() {
  // return the minimum possible energy - accounting for spectrum and universal limits on it
  if (fEMin > fELimUniLow) {
    return fEMin;
  } else {
    return fELimUniLow;
  }
}

///-------------------------------------------------------------------------
}