// VertexGen_NucleonDecay.cc
// Contact person:Ian Coulter <icoulter@hep.upenn.edu>
// See VertexGen_NucleonDecay.hh for more details
//———————————————————————//
//
// RAT::NucleonDecayGen
//
//
// Code which takes user's choice of proton or neutron
// and generates the emitted decay products from invisible
// nucleon decay in O16
//
//
//
//
//
/////////////////////////////////////////////////////



#include <CLHEP/Units/PhysicalConstants.h>
#include <Randomize.hh>
using namespace CLHEP;

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

#include <RAT/EventInfo.hh>
#include <RAT/Factory.hh>
#include <RAT/VertexGen_NucleonDecay.hh>
#include <RAT/Log.hh>
#include <CLHEP/Units/SystemOfUnits.h>
#include <RAT/ParentPrimaryVertexInformation.hh>

namespace RAT
{

  VertexGen_NucleonDecay::VertexGen_NucleonDecay() : fStateStr("")
  {
    fNucleonType = "proton";                         // Generate for proton decay by default

    info << "NucleonDecayGen: Loading default 'proton' as particle type \n";

    fGamma = G4ParticleTable::GetParticleTable()->FindParticle("gamma");
    fProton = G4ParticleTable::GetParticleTable()->FindParticle("proton");
    fNeutron = G4ParticleTable::GetParticleTable()->FindParticle("neutron");
  }

  VertexGen_NucleonDecay::~VertexGen_NucleonDecay()
  {
    // Do Nothing
  }

  void VertexGen_NucleonDecay::BeginOfRun() {
    fNDecay = DB::Get()->GetLink("NUCLEON_DECAY");   // load database

    // Get the possible energies and branching fractions of the emitted gammas
    fEnergyGammaN = fNDecay->GetDArray( "neutron_gamma" );
    fEnergyProtonN = fNDecay->GetDArray( "neutron_proton" );
    fEnergyNeutronN = fNDecay->GetDArray( "neutron_neutron" );
    fFractionN = fNDecay->GetDArray( "neutron_prob" );
    fEnergyGammaP = fNDecay->GetDArray( "proton_gamma" );
    fEnergyProtonP = fNDecay->GetDArray( "proton_proton" );
    fEnergyNeutronP = fNDecay->GetDArray( "proton_neutron" );
    fFractionP = fNDecay->GetDArray( "proton_prob" );

    // Check EGamma & Fraction are same length
    if(fEnergyGammaN.size() != fEnergyProtonN.size() && fEnergyGammaN.size() != fEnergyNeutronN.size() && fEnergyGammaN.size() != fFractionN.size() && fEnergyGammaN.size() > 0)
      Log::Die("NucleonDecayGen error: Energy spectrums and branching fractions have different sizes");
    else if(fEnergyGammaP.size() != fEnergyProtonP.size() && fEnergyGammaP.size() != fEnergyNeutronP.size() && fEnergyGammaP.size() != fFractionP.size() && fEnergyGammaP.size() > 0)
      Log::Die("NucleonDecayGen error: Energy spectrums and branching fractions have different sizes");

  }

  void VertexGen_NucleonDecay::CreateVertex(G4String type, G4double energy, G4Event* event, G4ThreeVector& dx, G4double dt)
  {
    // Define vertex;
    G4PrimaryVertex* vertex= new G4PrimaryVertex(dx, dt);
    G4ThreeVector dir(0.0,0.0,0.0);
    G4LorentzVector mom;

    // Pick isotropic direction
    double theta = acos(2.0 * HepUniformRand() - 1.0);
    double phi = 2.0 * HepUniformRand() * pi;
    dir.setRThetaPhi(1.0, theta, phi);

    // Compute 4-momentum of particle
    mom.setVect(dir * energy); // MeV (divide by c if need real units)
    mom.setE(energy);

    // Set particle type
    G4ParticleDefinition* particle;
    if(type == "proton") particle= fProton;
    else if(type == "neutron") particle= fNeutron;
    else if(type == "gamma") particle= fGamma;
    else Log::Die("Unknown particle type requested by nucleon decay generator");

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

  void VertexGen_NucleonDecay::GeneratePrimaryVertex(G4Event *event, G4ThreeVector &dx, G4double dt)
  {
    // Determine energy of emitted particles
    int i = SelectDecayMode(fNucleonType);

    double eGamma=0, eProton=0, eNeutron=0;
    if(fNucleonType == "neutron"){
      eGamma = fEnergyGammaN[i];
      eProton = fEnergyProtonN[i];
      eNeutron = fEnergyNeutronN[i];
    }
    else if(fNucleonType == "proton"){
      eGamma = fEnergyGammaP[i];
      eProton = fEnergyProtonP[i];
      eNeutron = fEnergyNeutronP[i];
    }
    else Log::Die("Unknown nucleon type requested by nucleon decay generator. Try proton or neutron instead.");

    if(eGamma) CreateVertex("gamma",eGamma,event,dx,dt);
    if(eProton) CreateVertex("proton",eProton,event,dx,dt);
    if(eNeutron) CreateVertex("neutron",eNeutron,event,dx,dt);
  }

  int VertexGen_NucleonDecay::SelectDecayMode(G4String nucleonType="")
  {
    double rnd = HepUniformRand();
    double prob = 0;
    double fracTotal = 0;

    if(nucleonType == "neutron")
    {
      // Get total visible branching fractions for renormalisation
      for(size_t i=0; i<fFractionN.size(); i++)
	fracTotal += fFractionN[i];

      // Select mode of decay by drawing random number
      for(size_t i=0; i<fEnergyGammaN.size(); i++)
      {
	prob += fFractionN[i]/fracTotal;
	if(rnd <= prob) return i;
      }

      // Shouldn't reach here
      warn << "NucleonDecayGen: Sum of branching fractions not equal to one\n";
      return -1;
    }
    else if(nucleonType == "proton")
      {
	// Get total visible branching fractions for renormalisation
	for(size_t i=0; i<fFractionP.size(); i++)
	  fracTotal += fFractionP[i];

	// Select mode of decay by drawing random number
	for(size_t i=0; i<fEnergyGammaP.size(); i++)
	  {
	    prob += fFractionP[i]/fracTotal;
	    if(rnd <= prob) return i;
	  }

	// Shouldn't reach here
	warn << "NucleonDecayGen: Sum of branching fractions not equal to one\n";
	return -1;
      }
    else
    {
      warn << "NucleonDecayGen: Need to specify the type of decay, proton or neutron\n";
      return -1;
    }
  }

  void VertexGen_NucleonDecay::SetState(G4String state)
  {
    fNucleonType = state;
    info << "NucleonDecayGen: Setting nucleon type to " << state << "\n";

    if(fNucleonType != "neutron" && fNucleonType != "proton"){
      Log::Die("NucleonDecayGen error: Unknown nucleon type :" + state) ;
    }
  }

  G4String VertexGen_NucleonDecay::GetState()
  {
    return fStateStr;
  }

}// namespace RAT