// 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(""), fQpositionXYZ(false),fQpositionThetaPhi(false),fQDirectionThetaPhi(false),
                        fTheta(-999),fPhi(-999),fPosx(-999),fPosy(-999),fPosz(-999),fTheta_dir(-999), fPhi_dir(-999) {

    // 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 inner radius
    fR1 = 8500.; // 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 !");
    if (fQpositionXYZ && fQpositionThetaPhi) throw FluxException("Gen_Flux error: /generator/flux/setposition called with both Cartesian and spherical coordinates may lead on to position ambiguity, remove one.");
    G4ThreeVector dir;
    double dE0, CosTheta0;

    // if direction is given in macro use it
    if(fQDirectionThetaPhi){
      // You are writing the direction custom.
      SampleE(dE0,cos(fTheta_dir));
      dir = G4ThreeVector(sin(fTheta_dir)*cos(fPhi_dir), sin(fTheta_dir)*sin(fPhi_dir), cos(fTheta_dir));
    // otherwise sample from input spectrum
    }else{
      // Generate initial particle energy and direction
      GenInit(dE0,CosTheta0);
      double SinTheta0 = sqrt(1-CosTheta0*CosTheta0);
      double dPhi0 = twopi*G4UniformRand();
      dir = G4ThreeVector(SinTheta0*cos(dPhi0), SinTheta0*sin(dPhi0), CosTheta0);
    }

    // if position is given in macro use that
    G4ThreeVector position;
    if(fQpositionXYZ)
      position = G4ThreeVector(fPosx,fPosy,fPosz);
    else if (fQpositionThetaPhi){
      position= G4ThreeVector(sin(fTheta)*cos(fPhi), sin(fTheta)*sin(fPhi), cos(fTheta));
      position *= GetRadius();
    // otherwise sample it from the spectrum
    }else{
      // get vertex position
      double xpos, ypos, zpos;
      GenVertex(fR1, dir, xpos, ypos, zpos);
      position= G4ThreeVector(xpos, ypos, zpos);
    }

    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::SampleE(double &E, const double CosTheta) {

    bool passed=false;

    while(!passed){
      E = pow(10,fEmin+(fEmax-fEmin)*G4UniformRand());
      // decided whether to draw again based on flux
      double test = fMmax * G4UniformRand();
      double FluxWeight = fMagnitude(E, CosTheta);
      passed = (FluxWeight >= test);
    }
  }
///-------------------------------------------------------------------------
  void Gen_Flux::GenInit(double &E, double &CosTheta) {

    bool passed=false;

    while(!passed){
      // pick E and CosTheta uniformly
      E = pow(10,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 &x_vertex, double &y_vertex, double &z_vertex) {

    // Given direction vector we can now set up the perpendicular plane to said direction.
    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_radius1 = 0;
    double tang_radius2 = 0;
    bool passed=false;
    while(!passed){
      // Now randomly sample a point on the perpendicular plane and test
      // whether point is inside the circler cross section of radius
      // long_radius.
      // 1-2*G4UniformRand() is required to sample the whole plane.
      tang_radius1 = long_radius*(1-2*G4UniformRand());
      tang_radius2 = long_radius*(1-2*G4UniformRand());
      passed = tang_radius1*tang_radius1 + tang_radius2*tang_radius2 < long_radius*long_radius;
    }

    // Scale basis vectors to point to new point coordinate and create shift vector.
    tang_vec1 *= tang_radius1;
    tang_vec2 *= tang_radius2;
    const G4ThreeVector shift = tang_vec1+tang_vec2;

    // Now that we have the shift length we can work out the radius at which
    // the new point has to be pushed out by to sit on the inner surface of the
    // PSUP.
    const double adjusted_radius = sqrt(long_radius*long_radius - shift.mag2());
    const double x0 = adjusted_radius*direction.x();
    const double y0 = adjusted_radius*direction.y();
    const double z0 = adjusted_radius*direction.z();

    // Return generated position.
    x_vertex = x0 + shift.x() ;
    y_vertex = y0 + shift.y() ;
    z_vertex = z0 + shift.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::SetGeneratorPosition(const double x0_,const double y0_,const double z0_) {
    fPosx = x0_;
    fPosy = y0_;
    fPosz = z0_;
    fQpositionXYZ =true;
  }
///-------------------------------------------------------------------------
  void Gen_Flux::SetGeneratorPosition(const double theta_,const double phi_) {

    fTheta = theta_;
    fPhi = phi_;
    fQpositionThetaPhi =true;
  }
  void Gen_Flux::SetGeneratorDirection(const double theta_,const double phi_) {

    fTheta_dir = theta_;
    fPhi_dir = phi_;
    fQDirectionThetaPhi =true;
  }
///-------------------------------------------------------------------------

  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 = log10(fLFlux->GetD("emin"));
    fEmax = log10(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<-6. || fEmax<-6.) {
      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