// VertexGen_Axion.cc
// Contact person: Esther Turner <esther.turner@physics.ox.ac.uk>
// See VertexGen_Axion.hh for more details
//———————————————————————//

#include <Randomize.hh>
#include <globals.hh>
#include <cerrno>
using namespace CLHEP;

//GLG4Utils
#include <RAT/GLG4StringUtil.hh>

//CLHEP includes
#include <CLHEP/Units/PhysicalConstants.h>
#include <CLHEP/Units/SystemOfUnits.h>

//Geant 4 includes
#include <G4Event.hh>
#include <G4ParticleDefinition.hh>
#include <G4ParticleTable.hh>
#include <G4PrimaryParticle.hh>
#include <G4TransportationManager.hh>
#include <G4PrimaryVertex.hh>
#include <G4DynamicParticle.hh>

//RAT includes
#include <RAT/EventInfo.hh>
#include <RAT/Factory.hh>
#include <RAT/VertexGen_Axion.hh>
#include <RAT/Log.hh>
#include <RAT/StringUtil.hh>
#include <RAT/SunUtil.hh>
#include <RAT/ParentPrimaryVertexInformation.hh>

#include <math.h>

namespace RAT
{

  VertexGen_Axion::VertexGen_Axion() : GLG4VertexGen("vertex_axion"),fStateStr(""),fInteractionType("compton_conversion")
  {
    fDiffCrossSection = new TGraph();

    //set the start of the event
    fEvTime.SetStart();

    //gamma particle produced in the interaction
    fGamma = G4ParticleTable::GetParticleTable()->FindParticle("gamma");
    fElectron = G4ParticleTable::GetParticleTable()->FindParticle("e-");

    //kinetic energy of gamma produced by Axion (see Belliniet  al.,PRD 85, 092003  2012)
    //This is produced from the p + d -> 3He + A (5.5MeV)
    //Kinetic energy of Axions is always 5.5MeV from this interaction in the sun
    faxionEnergy = 5.5; //MeV

    G4int days = fEvTime.GetUTDays();
    G4int secs = fEvTime.GetUTSecs();
    G4int nsecs = fEvTime.GetUTNsecs();

    TVector3 solar_dir = RAT::SunDirection(days,secs,nsecs);

    //direction of the sun
    fSolarDirection = G4ThreeVector(solar_dir.x(),solar_dir.y(),solar_dir.z());

    //set the mass of the electron
    fMassElectron = electron_mass_c2;
  }

  VertexGen_Axion::~VertexGen_Axion()
  {
    delete fDiffCrossSection;
  }

  //cross section calculation taken from Avignone et al. (PRD 37, 3, p.618)
  double VertexGen_Axion::ComptonAxionDifferentialCrossSection(double theta) const {
    double particleThreeMomentum = TMath::Sqrt(faxionEnergy*faxionEnergy - faxionMass*faxionMass); //assuming very little rest mass energy

    double y = 2.0*fMassElectron*faxionEnergy + faxionMass*faxionMass;
    double z = 2.0*(fMassElectron + faxionEnergy - particleThreeMomentum*TMath::Cos(theta));
    double e_gamma = y/z;

    double front_term = e_gamma/(2.0*fMassElectron*fMassElectron*particleThreeMomentum);
    double first_term = 4.0*fMassElectron*fMassElectron*e_gamma*e_gamma/(y*y);
    double second_term = 4.0*fMassElectron*e_gamma/y;
    double third_term = 4.0*TMath::Sin(theta)*TMath::Sin(theta)*faxionMass*faxionMass*particleThreeMomentum*particleThreeMomentum*fMassElectron*e_gamma/(y*y*y);
    double differentialCrossSection = front_term*(1.0 + first_term - second_term - third_term);

    return differentialCrossSection;
  }

  //cross section calculation taken from Avignone et al. (PRD 37, 3, p.618) [equation 16]
  double VertexGen_Axion::InversePrimakoffConversionDifferentialCrossSection(double theta) const{
    //assumming the axion mass is much smaller than the electron mass
    if(theta < 0.0001){
      theta = 0.0001; //prevent any discontinuities
    }
    double diffCrossSect = ( 1.0 + TMath::Cos(theta) )/( 1.0 - TMath::Cos(theta) );
#ifdef RATDEBUG
    debug << "[VertexGen_Axion]::InversePrimakoffConversionDifferentialCrossSection Differential Cross Section= " << diffCrossSect << newline;
#endif
    return diffCrossSect;
  }

  double VertexGen_Axion::SampleDifferentialCrossSection() const{
    // Get the shape of the differential cross section.
    // Interpolate between the discrete points
    CLHEP::RandGeneral rndm(fDiffCrossSection->GetY(),fDiffCrossSection->GetN(),0);
    return rndm.shoot()*(this->fDiffCrossSection->GetX())[fDiffCrossSection->GetN()-1];
  }

  /* Draws the differential cross section*/
  void VertexGen_Axion::DrawDiffCrossSection() const {
    const int npoints = 10000;
    //theta values are given in radians
    const double thetaMin = 0.0, thetaMax = 1.0*pi;

    const double xwid = thetaMax - thetaMin;
    const double xstep = xwid/(double)npoints;
    double theta,dSigma;

    for (int ip = 0; ip < npoints; ip++)
      {
        theta = thetaMin + (double)ip*xstep;
        if(fInteractionType == "compton_conversion"){
          dSigma = ComptonAxionDifferentialCrossSection(theta);
        }
        else if(fInteractionType == "inverse_primakoff_conversion"){
          dSigma = InversePrimakoffConversionDifferentialCrossSection(theta);
        }
        else{
          Log::Die("[VertexGen_Axion]::DrawDiffCrossSection: VertexGen_Axion error: Unknown Interaction type");
        }
          fDiffCrossSection->SetPoint(ip,theta,dSigma);
      }
  }

  void VertexGen_Axion::GeneratePrimaryVertex(G4Event *event, G4ThreeVector &dx, G4double dt)
  {
    // Define vertex;
    G4PrimaryVertex* vertex= new G4PrimaryVertex(dx, dt);
    G4LorentzVector solar_axion;

    fAxionDirection = -fSolarDirection;
    double axionMomentum = sqrt(faxionEnergy*faxionEnergy - faxionMass*faxionMass);
    G4ThreeVector axionMomtumThreeVector(fAxionDirection*axionMomentum);
    solar_axion.setVectM(axionMomtumThreeVector,faxionMass);

    //Creating an axion primary parent particle. Please note that "0" is the code used for exotic particles.
    G4PrimaryParticle *solar_axion_parent = new G4PrimaryParticle(0, axionMomtumThreeVector.x(), axionMomtumThreeVector.y(), axionMomtumThreeVector.z(),faxionEnergy);

    //Set parent particle information
    ParentPrimaryVertexInformation *vertinfo = new ParentPrimaryVertexInformation();
    vertinfo->SetVertexParentParticle(solar_axion_parent);
    vertex->SetUserInformation(vertinfo);

    if(fInteractionType == "compton_conversion"){

      //components specific to this type of interaction
      G4LorentzVector gamma;
      G4LorentzVector electron;

      //Generate a random value in Phi uniform between -pi and pi
      double gammaPhiDir = G4UniformRand()*2.0* pi - 1.0*pi;

      //Generate a random value in theta between 0 and pi that is weighted by differential cross section
      double gammaThetaDir = SampleDifferentialCrossSection();

      // Energy of the Gamma particle taken from Avignone et al. (PRD 37, 3, p.618)
      double gammaEnergy = ( fMassElectron*faxionEnergy )/ (fMassElectron + faxionEnergy - faxionEnergy*TMath::Cos(gammaThetaDir));

      G4ThreeVector gammaMomentum(gammaEnergy*fAxionDirection);
#ifdef RATDEBUG
        debug << "[VertexGen_Axion]::GeneratePrimaryVertex: gamma vector: x" << gammaMomentum.x() << " y: " << gammaMomentum.y() << " z: " << gammaMomentum.z()<< newline;
#endif
      G4ThreeVector rotation_axis = fAxionDirection.orthogonal();
      rotation_axis.rotate(gammaPhiDir, fAxionDirection);
      gammaMomentum.rotate(gammaThetaDir, rotation_axis);

      // Compute 4-momentum of gamma
      gamma.setVect(gammaMomentum); // MeV (divide by c if need real units)
      gamma.setE(gammaEnergy);

      // Compute electron 4-momentum
      electron.setVectM(solar_axion.vect() - gamma.vect(),fMassElectron);
#ifdef RATDEBUG
      double energyCheck = (solar_axion.e() - electron.e() - gamma.e());
      double momentumCheckX = solar_axion.px() - electron.px() - gamma.px();
      double momentumCheckY = solar_axion.py() - electron.py() - gamma.py();
      double momentumCheckZ = solar_axion.pz() - electron.pz() - gamma.pz();
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: electron momentum: x" << electron.px() << " y: " << electron.py() << " z: " << electron.pz()<< newline;
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: gamma momentum: x" << gamma.px() << " y: " << gamma.py() << " z: " << gamma.pz()<< newline;
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: Momentum Check: " <<  momentumCheckX << " " <<  momentumCheckY << " " <<  momentumCheckZ << newline;
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: Energy Check: " <<  energyCheck << newline;
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: Interaction Angle (before generation): " << gammaThetaDir << newline;
      double interactionAngleCheck = TMath::ACos(solar_axion.vect().dot(gamma.vect())/(solar_axion.vect().mag()*gamma.vect().mag()));
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: Interaction Angle (after generation): " << interactionAngleCheck << newline;
#endif
      G4PrimaryParticle* pGammaParticle = new G4PrimaryParticle(fGamma,       // particle code
                                                                gamma.px(),     // x component of momentum
                                                                gamma.py(),     // y component of momentum
                                                                gamma.pz());    // z component of momentum
      pGammaParticle->SetMass(fGamma->GetPDGMass());
      vertex->SetPrimary( pGammaParticle );
      //event->AddPrimaryVertex(vertex);

      G4PrimaryParticle* e_minus_particle =
      new G4PrimaryParticle(fElectron,              // particle code
                              electron.px(),     // x component of momentum
                              electron.py(),     // y component of momentum
                              electron.pz());    // z component of momentum
      e_minus_particle->SetMass(fElectron->GetPDGMass());
      vertex->SetPrimary( e_minus_particle );
      event->AddPrimaryVertex(vertex);

    }        //end of compton conversion vertex

    else if(fInteractionType == "inverse_primakoff_conversion"){
       // Inverse primakoff conversion taken from Avignone et al. (PRD 37, 3, p.618)

      G4LorentzVector gamma;
      G4ThreeVector gammaMomentum(faxionEnergy*fAxionDirection);

      double gammaPhiDir = G4UniformRand()*2.0* pi - 1.0*pi;
      double gammaThetaDir = SampleDifferentialCrossSection();

      G4ThreeVector rotation_axis = fAxionDirection.orthogonal();
      rotation_axis.rotate(gammaPhiDir, fAxionDirection);
      gammaMomentum.rotate(gammaThetaDir, rotation_axis);

      gamma.setVect(gammaMomentum);
      gamma.setE(faxionEnergy);  //not sure if this should be sqrt(electron.vect().mag2() + fElectronMass*fElectronMass)
#ifdef RATDEBUG
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: gamma momentum: x" << gamma.px() << " y: " << gamma.py() << " z: " << gamma.pz() << newline;
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: Interaction Angle (before generation): " << gammaThetaDir << newline;

      double interactionAngleCheck = TMath::ACos(solar_axion.vect().dot(gamma.vect())/(solar_axion.vect().mag()*gamma.vect().mag()));
      debug << "[VertexGen_Axion]::GeneratePrimaryVertex: Interaction Angle (after generation): " << interactionAngleCheck << newline;
#endif
      G4PrimaryParticle* pGammaParticle = new G4PrimaryParticle(fGamma,       // particle code
                                                                gamma.px(),     // x component of momentum
                                                                gamma.py(),     // y component of momentum
                                                                gamma.pz());    // z component of momentum
      pGammaParticle->SetMass(fGamma->GetPDGMass());
      vertex->SetPrimary( pGammaParticle );
      event->AddPrimaryVertex(vertex);
    }        //end of inverse primakoff conversion vertex

    // Taken from Muratova! et al.! (hep-ex:! 1501.02943v1)
    else if(fInteractionType == "axioelectric_effect"){
      G4LorentzVector electron;
      G4double total_energy = faxionEnergy + fMassElectron;
      G4double p_e = sqrt(total_energy*total_energy - fMassElectron*fMassElectron);
      electron.setVect(p_e*fAxionDirection);
      electron.setE(total_energy);
      G4PrimaryParticle* e_minus_particle = new G4PrimaryParticle(fElectron,       // particle code
                                                                electron.px(),     // x component of momentum
                                                                electron.py(),     // y component of momentum
                                                                electron.pz());    // z component of momentum
      e_minus_particle->SetMass(fElectron->GetPDGMass());
      vertex->SetPrimary( e_minus_particle );
      event->AddPrimaryVertex(vertex);

    }        //end of axioelectric effect vertex
    else if(fInteractionType == "axion_decay"){
      G4LorentzVector gamma_one, gamma_two;
      G4ThreeVector gamma_one_momentum((faxionEnergy/2.0)*fAxionDirection);

      //Find an isotropic theta, phi for both gammas in axion frame
      double gamma_one_phi = G4UniformRand()*2.0* pi - 1.0*pi;
      double gamma_one_theta = G4UniformRand()*1.0* pi;

      G4ThreeVector rotation_axis = fAxionDirection.orthogonal();
      rotation_axis.rotate(gamma_one_phi, fAxionDirection);
      gamma_one_momentum.rotate(gamma_one_theta, rotation_axis);

      gamma_one.setVect(gamma_one_momentum);
      gamma_one.setE(faxionEnergy/2.0);

      gamma_two.setVect(-1.0*gamma_one_momentum);
      gamma_two.setE(faxionEnergy/2.0);

      //Generate the two gamma particles from the axion decay
      G4PrimaryParticle* gamma_one_particle = new G4PrimaryParticle(fGamma,       // particle code
                                                                gamma_one.px(),     // x component of momentum
                                                                gamma_one.py(),     // y component of momentum
                                                                gamma_one.pz());    // z component of momentum
      gamma_one_particle->SetMass(fGamma->GetPDGMass());
      vertex->SetPrimary( gamma_one_particle );

      G4PrimaryParticle* gamma_two_particle = new G4PrimaryParticle(fGamma,       // particle code
                                                                gamma_two.px(),     // x component of momentum
                                                                gamma_two.py(),     // y component of momentum
                                                                gamma_two.pz());    // z component of momentum
      gamma_two_particle->SetMass(fGamma->GetPDGMass());
      vertex->SetPrimary( gamma_two_particle );
      event->AddPrimaryVertex(vertex);
    }        //end of axion_decay vertex
    else{
      Log::Die("[VertexGen_Axion]::VertexGen_Axion error: Unknown Interaction type :" + fInteractionType) ;
    }

  }

  void VertexGen_Axion::SetState(G4String state)
  {
    state = util_strip_default(state); // from GLG4StringUtil

    if (state.length() == 0) {
      Log::Die("[VertexGen_Axion]::SetState: Error invalid commands in the macro. Please provide commands in the format /generator/vtx/set <interaction_type> <mass_of_axion_in_MeV>");
    }
    std::istringstream is(state.c_str());
    std::string rest;
    is >> rest;

    // Get separate the entries by the ':'
    std::string parseparams = ":";
    rest = strip(rest,parseparams);
    std::vector<std::string> params = util_split(rest,parseparams);
    char * errorX,*errorY,*errorZ,*errorMass;
    double xDirc,yDirc,zDirc;
    switch (params.size()) {
    case 1:
      //define the interaction state
      fInteractionType = params[0];
      if( (fInteractionType != "compton_conversion") &&
          (fInteractionType != "axioelectric_effect") &&
          (fInteractionType != "axion_decay") &&
          (fInteractionType != "inverse_primakoff_conversion") ){
        Log::Die("[VertexGen_Axion]::SetState: Error - unknown interaction type :" + fInteractionType) ;
      }
      faxionMass = fMassElectron/100.0;
      info << "[VertexGen_Axion]::SetState: No Axion mass specified. Defaulting to 0.01*electron_mass" << newline;
      break;

    case 2:
      fInteractionType = params[0];
      if( (fInteractionType != "compton_conversion") &&
          (fInteractionType != "axioelectric_effect") &&
          (fInteractionType != "axion_decay") &&
          (fInteractionType != "inverse_primakoff_conversion") ){
        Log::Die("[VertexGen_Axion]::SetState: Error - unknown interaction type :" + fInteractionType) ;
      }

      faxionMass = std::strtod(params[1].c_str(), &errorMass);
      if (*errorMass != '\0'||
          errno != 0){   // error, overflow or underflow
          Log::Die("[VertexGen_Axion]::SetState: Invalid Axion Mass Argument");
      }

      if(faxionMass >= 5.5){
        Log::Die("[VertexGen_Axion]::SetState: Error mass of Axion is in units of MeV. The mass of axion cannot be more than 5.5MeV");
      }
      info << "[VertexGen_Axion]::SetState:  Setting axion mass to " << faxionMass << " MeV" << newline;
      break;

    case 4:
      fInteractionType = params[0];
      if( (fInteractionType != "compton_conversion") &&
          (fInteractionType != "axioelectric_effect") &&
          (fInteractionType != "axion_decay") &&
          (fInteractionType != "inverse_primakoff_conversion") ){
        Log::Die("[VertexGen_Axion]::SetState: Error - unknown interaction type :" + fInteractionType);
      }

      xDirc = std::strtod(params[1].c_str(),&errorX);
      yDirc = std::strtod(params[2].c_str(),&errorY);
      zDirc = std::strtod(params[3].c_str(),&errorZ);
      if (*errorX != '\0' ||
          *errorY != '\0' ||
          *errorZ != '\0' ||
          errno != 0){
        Log::Die("[VertexGen_Axion]::SetState: Invalid Solar Direction Argument");
      }
      else{
        fSolarDirection = -1.0*G4ThreeVector(xDirc,yDirc,zDirc).unit();
      }
      faxionMass = fMassElectron/100.0;
      info << "[VertexGen_Axion]::SetState:  Setting axion mass to " << faxionMass << " MeV" << newline;
      break;

    case 5:
      fInteractionType = params[0];
      if( (fInteractionType != "compton_conversion") &&
          (fInteractionType != "axioelectric_effect") &&
          (fInteractionType != "axion_decay") &&
          (fInteractionType != "inverse_primakoff_conversion") ){
        Log::Die("[VertexGen_Axion]::SetState: Error - unknown interaction type :" + fInteractionType) ;
      }

      faxionMass = std::strtod(params[1].c_str(),&errorMass);
      if(faxionMass >= 5.5){
        Log::Die("[VertexGen_Axion]::SetState: Error mass of Axion is in units of MeV. The mass of axion cannot be more than 5.5MeV");
      }

      xDirc = std::strtod(params[2].c_str(),&errorX);
      yDirc = std::strtod(params[3].c_str(),&errorY);
      zDirc = std::strtod(params[4].c_str(),&errorZ);
      if (*errorMass != '\0' ||
          *errorX != '\0' ||
          *errorY != '\0' ||
          *errorZ != '\0' ||
          errno != 0){
        Log::Die("[VertexGen_Axion]::SetState: Invalid Solar Direction Argument");
      }
      else{
        fSolarDirection = -1.0*G4ThreeVector(xDirc,yDirc,zDirc).unit();
      }
      info << "[VertexGen_Axion]::SetState:  Setting axion mass to " << faxionMass << " MeV" << newline;
      break;

    default:
      Log::Die("[VertexGen_Axion]::SetState: Error invalid commands in the macro. Please provide commands in the format /generator/vtx/set <interaction_type>:<mass_of_axion_in_MeV>:<axion_direction>");
      break;

    }

    //Only call the ability to draw the cross section for these two instances.
    if( fInteractionType == "compton_conversion" ||
        fInteractionType == "inverse_primakoff_conversion" ){
        //fill the TGraph with the differential cross section
        DrawDiffCrossSection();
    }

  }
  //TODO:FIXME
  G4String VertexGen_Axion::GetState()
  {
    return fStateStr;
  }

}// namespace RAT