// Gen_Flux.cc
// Contact person:  Eric Vazquez-Jauregui <ericvj@fisica.unam.mx>
// See Gen_Flux.hh for more details
//———————————————————————//

#include <CLHEP/Units/PhysicalConstants.h>
#include <RAT/Gen_Flux.hh>
#include <RAT/Gen_FluxMessenger.hh>
#include <RAT/EventInfo.hh>
#include <RAT/Log.hh>
#include <RAT/GLG4TimeGen.hh>
#include <Randomize.hh>
#include <G4Event.hh>
#include <G4Track.hh>
#include <G4ParticleDefinition.hh>
#include <G4ParticleTable.hh>
#include <G4PrimaryVertex.hh>
#include <G4PrimaryParticle.hh>
#include <RAT/Factory.hh>
#include <CLHEP/Units/SystemOfUnits.h>
#include <time.h>

using namespace CLHEP;
using namespace std;

namespace RAT {

  Gen_Flux::Gen_Flux() : GLG4Gen("flux"),fStateStr("") {

    // create a messenger to allow the user to change some parameters
    fMessenger = new Gen_FluxMessenger(this);

    // default time generator
    fTimeGen = new GLG4TimeGen_Poisson();

    // default radius is the psup h2o outer radius
    fR1 = 8800.; // mm

    fEmin = 0.;
    fEmax = 0.;
    fCosTmin = 0.;
    fCosTmax = 0.;
    fMmax = 0.;

    info << "Gen_Flux: Loading default 'poisson' time generator \n";
    info << "Gen_Flux: Loading default 'psup h20 outer' radius: \"" << fR1 << "\" mm \n";
    info << "Gen_Flux: now set spectrum \n";
  }

///-------------------------------------------------------------------------
  Gen_Flux::~Gen_Flux() {

    if ( fMessenger != 0 ){
      delete fMessenger;
      fMessenger = 0;
    }
  }

///-------------------------------------------------------------------------

  void Gen_Flux::BeginOfRun() {
    GLG4Gen::BeginOfRun();
    // load database values
    fLFlux = DB::Get()->GetLink("FLUX");

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

  void Gen_Flux::GenerateEvent(G4Event *event) {

    if(GetSpectrum().length() == 0) Log::Die("Gen_Flux error: input spectrum not defined. Please set input spectrum via /generator/flux/setspectrum !");

    double E0, CosTheta0;
    // generate initial particle energy and direction
    GenInit(E0,CosTheta0);
    double dE0 = E0;
    double dTheta0 = acos(CosTheta0);
    double dPhi0 = twopi*G4UniformRand();

    G4ThreeVector dir = G4ThreeVector(sin(dTheta0)*cos(dPhi0), sin(dTheta0)*sin(dPhi0), cos(dTheta0));

    // get vertex position
    double xpos, ypos, zpos;
    GenVertex(fR1, dir, dTheta0, xpos, ypos, zpos);
    double dx = xpos;
    double dy = ypos;
    double dz = zpos;

    G4ThreeVector position(dx, dy, dz);

    G4double ntime = NextTime();

    // get particle momentum (-1 to propagate into detector)
    G4ThreeVector mom( -1. * dir * dE0 );

    G4PrimaryVertex *vertex = new G4PrimaryVertex(position, ntime);
    G4PrimaryParticle *particle = new G4PrimaryParticle(fPDef, mom.x(), mom.y(), mom.z());

    // mu+/mu- at SNO+ ~ 1.4;
    if(GetSpectrum()=="cosmic" && G4UniformRand()*12. < 7.) particle->SetPDGcode(-13);

    vertex->SetPrimary(particle);
    event->AddPrimaryVertex(vertex);
  }
///-------------------------------------------------------------------------

  void Gen_Flux::GenInit(double &E, double &CosTheta) {

    bool passed=false;

    while(!passed){
      // pick E and CosTheta uniformly
      E = fEmin+(fEmax-fEmin)*G4UniformRand();
      CosTheta = fCosTmin+(fCosTmax-fCosTmin)*G4UniformRand();

      // decided whether to draw again based on flux
      double test = fMmax * G4UniformRand();
      double FluxWeight = fMagnitude(E, CosTheta);
      passed = (FluxWeight >= test);
    }
  }
///-------------------------------------------------------------------------

  void Gen_Flux::GenVertex(double long_radius, G4ThreeVector direction, double zenith_angle, double &x_vertex, double &y_vertex, double &z_vertex) {

    // original vertex position
    double x0 = long_radius*direction.x();
    double y0 = long_radius*direction.y();
    double z0 = long_radius*direction.z();

    // to not only propagate particles through AV centre but homogeneously through entire AV relocate original position on tangential ellipse
    G4ThreeVector zaxis = G4ThreeVector(0., 0., 1.);
    G4ThreeVector tang_vec1 = zaxis.cross(direction); // direction of major axis a
    G4ThreeVector tang_vec2 = tang_vec1.cross(direction); // direction of minor axis b

    tang_vec1 = tang_vec1.unit();
    tang_vec2 = tang_vec2.unit();

    double tang_radius = -888., ratio = -888.;
    bool passed=false;

    while(!passed){
      // decided whether to draw again based on relative arc length  l/lmax = a/amax = b/bmax
      tang_radius = long_radius*cos(zenith_angle)*G4UniformRand();

      double test = G4UniformRand();
      ratio = tang_radius/(long_radius*cos(zenith_angle));

      passed = (ratio >= test);
    }

    double alpha = twopi*G4UniformRand();

    tang_vec1 *= sin(alpha);
    tang_vec2 *= cos(alpha);

    x_vertex = x0 + ratio*long_radius*tang_vec1.x() + tang_radius*tang_vec2.x();
    y_vertex = y0 + ratio*long_radius*tang_vec1.y() + tang_radius*tang_vec2.y();
    z_vertex = z0 + ratio*long_radius*tang_vec1.z() + tang_radius*tang_vec2.z();

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

  void Gen_Flux::ResetTime(double offset) {
    fNextTime = fTimeGen->GenerateEventTime() + offset;
  }
///-------------------------------------------------------------------------

  void Gen_Flux::SetState(G4String state) {
    // use src/cmd/Gen_FluxMessenger to set state
    __unused_parameter(state);

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

  G4String Gen_Flux::GetState() const {
    return fStateStr;
  }
///-------------------------------------------------------------------------

  // taken from Gen_LED:
  void Gen_Flux::SetTimeGen(G4String state) {
    try {
      delete fTimeGen;
      fTimeGen = 0; // in case of exception in next line
      fTimeGen = RAT::GlobalFactory<GLG4TimeGen>::New(state);
      info << "Gen_Flux: Setting Flux time generator to " << state << "\n";
    }
    catch (RAT::FactoryUnknownID &unknown) {
      G4cerr << "Gen_Flux error: Unknown time generator \"" << unknown.id << "\"" << G4endl;
    }
  }
///-------------------------------------------------------------------------

  void Gen_Flux::SetTimeState(G4String state) {
    if (fTimeGen){
      fTimeGen->SetState(state);
    }
    else G4cerr << "Gen_Flux error: Cannot set time state, no time generator selected" << G4endl;
  }
///-------------------------------------------------------------------------

  G4String Gen_Flux::GetTimeState() const {
    if (fTimeGen)
      return fTimeGen->GetState();
    else
      return G4String("Gen_Flux error: no time generator selected");
  }
///-------------------------------------------------------------------------

  void Gen_Flux::SetRadius(double radius) {

    if(radius > 0. && radius <= 20000.){
      info << "Gen_Flux: Setting radius:" << " \"" << radius << "\"\n";
      fR1 = radius;
    }
    else{
      G4cerr << "Gen_Flux error: Cannot set radius: " << " \"" << radius << " \"" << G4endl;
      Log::Die("Gen_Flux error: radius not in valid range ]0.,20000.] mm");
    }

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

  double Gen_Flux::GetRadius() const {
    return fR1;
  }
///-------------------------------------------------------------------------

  void Gen_Flux::SetSpectrum(G4String state) {

    DBLinkGroup fLFluxTable = DB::Get()->GetLinkGroup("FLUX");
    DBLinkGroup::iterator i_index;

    bool isindex = false;

    for (i_index = fLFluxTable.begin(); i_index != fLFluxTable.end(); ++i_index) {

      if(i_index->first == state){
        info << "Gen_Flux: Setting flux spectrum:" << " \"" << state << "\"\n";

        fLFlux = DB::Get()->GetLink("FLUX", state);
        spectrum = state;
        isindex = true;
        break;
      }
    }
    if(!isindex){
      info << "Format of argument to Gen_Flux::SetSpectrum: \n"
           << "          \"specname\"\n"
           << "          specname = spectrum name as given in ratdb 'index' \n";
      Log::Die("Gen_Flux error: Could not find input spectrum: " + state);
    }

    // initialise new spectrum
    this->InitialiseSpectrum();
  }
///-------------------------------------------------------------------------

  G4String Gen_Flux::GetSpectrum() {
    return spectrum;
  }

///-------------------------------------------------------------------------
  void Gen_Flux::InitialiseSpectrum() {

    fEmin = fLFlux->GetD("emin");
    fEmax = fLFlux->GetD("emax");
    fCosTmin = fLFlux->GetD("costmin");
    fCosTmax = fLFlux->GetD("costmax");
    fMmax = fLFlux->GetD("mmax");

    fPDef = G4ParticleTable::GetParticleTable()->FindParticle(fLFlux->GetI("pdg"));

    if(fEmax<fEmin || fEmin<0. || fEmax<0.) {
      Log::Die("Gen_Flux error: Nonsensical energy range for spectrum: " + spectrum);
    }
    if(fCosTmax<fCosTmin || fCosTmin<-1. || fCosTmax>1.) {
      Log::Die("Gen_Flux error: Nonsensical CosTheta range or range exceeds [-1,1] for spectrum: " + spectrum);
    }
    if(fPDef==NULL) {
      Log::Die("Gen_Flux error: Could not find particle type for spectrum: " + spectrum);
    }

    fMagnitude.Set(fLFlux->GetDArray("energy"), fLFlux->GetDArray("costheta"), fLFlux->GetDArray("magnitude"));
  }

} // namespace RAT