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

#include <RAT/VertexGen_DecayChain.hh>
#include <RAT/DecayChain.hh>
#include <RAT/ParentPrimaryVertexInformation.hh>
#include <RAT/Log.hh>

#include <RAT/GLG4PosGen.hh>
#include <RAT/GLG4TimeGen.hh>
#include <RAT/Factory.hh>
#include <RAT/GLG4StringUtil.hh>

#include <G4Event.hh>
#include <G4PrimaryVertex.hh>
#include <G4PrimaryParticle.hh>
#include <G4IonTable.hh>
#include <G4Electron.hh>
#include <G4Positron.hh>
#include <G4Gamma.hh>
#include <G4Alpha.hh>
#include <G4ThreeVector.hh>
#include <G4UnitsTable.hh>

#include <cstring>

#undef DEBUG

namespace RAT{


  VertexGen_DecayChain::VertexGen_DecayChain(const char *arg_dbname)
    : GLG4VertexGen(arg_dbname), fStateStr(""), fDecayChain(0)
  {
    fInMiddle=false;
    fInAlphaDecay=false;
    fInGammaDecay=false;
    fScreenCorr=false;
  }


  VertexGen_DecayChain::~VertexGen_DecayChain(){}


  void VertexGen_DecayChain::GeneratePrimaryVertex(G4Event *event,
						   G4ThreeVector& dx,
						   G4double dt)
  {
    if (fDecayChain == 0) {
      warn << "RAT::VertexGen_DecayChain::GeneratePrimaryVertex: No decay chain found; "
	     << "please use /generator/add 'a top level generator' decaychain:ISOTOPE:POSITION[:TIME]\n";
      return;
    }

    // Generate the chain of decay products for this isotope.
    fDecayChain->GenerateFullChain();
    G4int nPrimaries = fDecayChain->GetNGenerated();

#ifdef DEBUG
    debug << "RAT::VertexGen_DecayChain::GeneratePrimaryVertex: "
          << nPrimaries << " primaries in decay chain "
	  << fDecayChain->GetChainName() << "\n";
#endif

    for ( G4int iPrimary = 0; iPrimary < nPrimaries; iPrimary++ ) {
	   // particle type
      G4int pid = fDecayChain->GetEventID(iPrimary); // not the PDG code!

      G4ParticleDefinition* particleDef = 0;
      if ( pid >= 100000000 ) {
	G4int A = (pid-100000000)/1000;
	G4int Z = (pid-100000000) - A*1000;
	G4double excitationEnergy = 0.0; // assume all in ground state
	G4IonTable* ionTable = G4IonTable::GetIonTable();
	particleDef = ionTable->GetIon(Z, A, excitationEnergy);
      }
      else {
	if ( pid == DecayBeta ) {
	  particleDef = G4Electron::Electron();
	}
        else if ( pid == DecayEC ) {
          particleDef = G4Positron::Positron();
        }
	else if ( pid == DecayGamma ) {
	  particleDef = G4Gamma::Gamma();
	}
	else if( pid == DecayAlpha ) {
	  particleDef = G4Alpha::Alpha();
	}
      }
      if ( particleDef == 0 ) {
	warn << "RAT::VertexGen_DecayChain::GenerateEvent: "
	     << "didn't know how to handle particle with ID " << pid << "\n";
	continue;
      }

      // Get the particle's information (momentum and time)
      DecayChain::ParticleInfo_t particleInfo = fDecayChain->GetParticleInfo(iPrimary);
      // generate a vertex with a primary particle
      G4PrimaryVertex* vertex = new G4PrimaryVertex(dx, dt);
      G4PrimaryParticle* particle = new G4PrimaryParticle(particleDef,
							  particleInfo.vector.px(),
							  particleInfo.vector.py(),
							  particleInfo.vector.pz());
      particle->SetMass(particleDef->GetPDGMass()); // Apparently this is useful in IBD code.
      vertex->SetPrimary(particle);
      event->AddPrimaryVertex(vertex);

// Set the parent information for this vertex (but don't generate parent!)
      G4IonTable *theIonTable = G4IonTable::GetIonTable();
      G4ParticleDefinition *parentDef = theIonTable->GetIon(fDecayChain->GetParentCharge(iPrimary),fDecayChain->GetParentMass(iPrimary),0.);
      G4PrimaryParticle *parent = new G4PrimaryParticle(parentDef,0,0,0);

      ParentPrimaryVertexInformation *vertinfo = new ParentPrimaryVertexInformation();
      vertinfo->SetVertexParentParticle(parent);
      vertex->SetUserInformation(vertinfo);

#ifdef DEBUG
      debug << "RAT::VertexGen_DecayChain::GeneratePrimaryVertex: "
	    << "Primary " << iPrimary+1 << " of " << nPrimaries
	    << ", pid=" << pid
	    << ", name=" << particleDef->GetParticleName() << "\n";
      debug << "    time=" << G4BestUnit(dt,"Time")
	    << ", position=" << G4BestUnit(dx,"Length")
	    << ", momentum=" << G4BestUnit(particleInfo.vector,"Energy")
	    << "\n";
#endif

    } // for each primary
  }


  void VertexGen_DecayChain::SetState(G4String state)
  {
#ifdef DEBUG
    debug << "RAT::VertexGen_DecayChain::SetState called with state='"
	  << state << "'" << G4endl;
#endif

    // Break the argument to the this generator into sub-strings
    // separated by ":".
    state = util_strip_default(state);
    vector<string> parts = util_split(state, ":");
    size_t nArgs = parts.size();

#ifdef DEBUG
    debug << "RAT::VertexGen_DecayChain::SetState: nArgs = " << nArgs << "\n";
#endif

    try {

      if( nArgs >= 3 ){

        string ScreenCorr = parts[2];
        if (ScreenCorr=="screen"){
          fScreenCorr=true;
          }
       if (ScreenCorr=="none"){
          fScreenCorr=false;
          }

      }

      if ( nArgs >= 2 ){

	string InMiddle = parts[1];
	if (InMiddle=="midchain"){
          info << "RAT::VertexGen_DecayChain: starting chain at Isotope\n";
          fInMiddle=true;
	}
	if (InMiddle=="alpha"){
          info << "RAT::VertexGen_DecayChain: alpha decay of Isotope\n";
          fInAlphaDecay=true;
	}
        else if (InMiddle=="gamma"){
          info << "RAT::VertexGen_DecayChain: gamma decay of Istope\n";
          fInGammaDecay=true;
        }
      }

      if ( nArgs >= 1 )	{
	// The first argument is the isotope that starts the decay chain.
	string isotope = parts[0];

	// Don't bother building a new decay chain if we already have
	// one for that isotope.
	if ( fDecayChain == 0  || isotope != fDecayChain->GetChainName() ) {
	  delete fDecayChain;
	  fDecayChain = new DecayChain(isotope);
#ifdef DEBUG
	  fDecayChain->SetVerbose(true);
#endif
	  if (fInMiddle) fDecayChain->SetInMiddleChain(true);
	  if (fInAlphaDecay) fDecayChain->SetAlphaDecayStart(true);
          if (fInGammaDecay) fDecayChain->SetGammaDecayStart(true);
          if (fScreenCorr) fDecayChain->SetScreenCorr(true);

	  bool found = fDecayChain->ReadInputFile(fDecayChain->GetChainName());
	  info << "found: "<<found<<", fDecayChain->GetChainName(): "<<fDecayChain->GetChainName() << "\n";
	  if (!found) {
	    warn << "RAT::VertexGen_DecayChain::SetState: couldn't find data for isotope "
		   << isotope << "\n";
	    delete fDecayChain; fDecayChain=0;
	  }
	  info << "RAT::VertexGen_DecayChain::SetState: successfully created decay chain for "
		 << fDecayChain->GetChainName() << "\n";
#ifdef DEBUG
	  fDecayChain->Show();
#endif
	}

      }
      else	{
	  G4Exception(__FILE__, "Invalid Parameter", FatalException, ("decaychain generator syntax error: '"+
								      state+
								      "' does not have a position generator").c_str());
      }

      fStateStr = state; // Save for later call to GetState()
    } catch (FactoryUnknownID &unknown) {
      warn << "Unknown generator \"" << unknown.id << "\"\n";
    }
  }


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

}