// VertexGen_DinucleonDecay.cc
// Contact person:Ian Coulter <icoulter@hep.upenn.edu>
// See VertexGen_NucleonDecay.hh for more details
//———————————————————————//
//
// RAT::DinucleonDecayGen
//
//
// Code which takes user's choice of pp, pn or nn
// and generates the emitted decay products from invisible
// dinucleon 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_DinucleonDecay.hh>
#include <RAT/Log.hh>
#include <CLHEP/Units/SystemOfUnits.h>
#include <RAT/ParentPrimaryVertexInformation.hh>

namespace RAT
{

  VertexGen_DinucleonDecay::VertexGen_DinucleonDecay() : fStateStr("")
  {
    fNucleonType = "pp";                         // Generate for p-p decay by default

    info << "DinucleonDecayGen: Loading default 'pp' as decay type \n";

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

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

  void VertexGen_DinucleonDecay::BeginOfRun() {
    fDNDecay = DB::Get()->GetLink("DINUCLEON_DECAY");   // load database

    // Get the possible energies and branching fractions of the emitted gammas
    fEnergyGammaNN = fDNDecay->GetDArray( "nn_gamma" );
    fEnergyProtonNN = fDNDecay->GetDArray( "nn_proton" );
    fEnergyNeutronNN = fDNDecay->GetDArray( "nn_neutron" );
    fFractionNN = fDNDecay->GetDArray( "nn_prob" );
    fEnergyGammaPN = fDNDecay->GetDArray( "pn_gamma" );
    fEnergyProtonPN = fDNDecay->GetDArray( "pn_proton" );
    fEnergyNeutronPN = fDNDecay->GetDArray( "pn_neutron" );
    fFractionPN = fDNDecay->GetDArray( "pn_prob" );
    fEnergyGammaPP = fDNDecay->GetDArray( "pp_gamma" );
    fEnergyProtonPP = fDNDecay->GetDArray( "pp_proton" );
    fEnergyNeutronPP = fDNDecay->GetDArray( "pp_neutron" );
    fFractionPP = fDNDecay->GetDArray( "pp_prob" );

    // Check EGamma & Fraction are same length
    if((fEnergyGammaNN.size() != fEnergyProtonNN.size() || fEnergyGammaNN.size() != fEnergyNeutronNN.size() || fEnergyGammaNN.size() != fFractionNN.size()) && fEnergyGammaNN.size() > 0)
      Log::Die("VertexGen_DinucleonDecay error: Energy spectrums and branching fractions have different sizes");
    else if((fEnergyGammaPN.size() != fEnergyProtonPN.size() || fEnergyGammaPN.size() != fEnergyNeutronPN.size() || fEnergyGammaPN.size() != fFractionPN.size()) && fEnergyGammaPN.size() > 0)
      Log::Die("VertexGen_DinucleonDecay error: Energy spectrums and branching fractions have different sizes");
    else if((fEnergyGammaPP.size() != fEnergyProtonPP.size() || fEnergyGammaPP.size() != fEnergyNeutronPP.size() || fEnergyGammaPP.size() != fFractionPP.size()) && fEnergyGammaPP.size() > 0)
      Log::Die("VertexGen_DinucleonDecay error: Energy spectrums and branching fractions have different sizes");

  }

  void VertexGen_DinucleonDecay::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_DinucleonDecay::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 == "nn"){
      eGamma = fEnergyGammaNN[i];
      eProton = fEnergyProtonNN[i];
      eNeutron = fEnergyNeutronNN[i];
    }
    else if(fNucleonType == "pn"){
      eGamma = fEnergyGammaPN[i];
      eProton = fEnergyProtonPN[i];
      eNeutron = fEnergyNeutronPN[i];
    }
    else if(fNucleonType == "pp"){
      eGamma = fEnergyGammaPP[i];
      eProton = fEnergyProtonPP[i];
      eNeutron = fEnergyNeutronPP[i];
    }
    else Log::Die("Unknown decay type requested by dinucleon decay generator. Try pp, pn or nn 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_DinucleonDecay::SelectDecayMode(G4String nucleonType="")
  {
    double rnd = HepUniformRand();
    double prob = 0;
    double fracTotal = 0;

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

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

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

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

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

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

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

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

    if(fNucleonType != "nn" && fNucleonType != "pn" && fNucleonType != "pp"){
      Log::Die("VertexGen_DinucleonDecay error: Unknown decay type :" + state) ;
    }
  }

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

}// namespace RAT