// M Hentz, 2016

#include "PrimaryGeneratorAction.hh"
#include "PrimaryGeneratorMessenger.hh"
#include "DetectorConstruction.hh"
#include "FileReader.hh"

#include "G4Event.hh"
#include "G4ParticleGun.hh"
#include "G4ParticleTable.hh"
#include "G4ParticleDefinition.hh"
#include "G4GeneralParticleSource.hh"
#include "G4SystemOfUnits.hh"

#include "Randomize.hh"
#include "globals.hh"

FileReader* PrimaryGeneratorAction::fReader = 0;

PrimaryGeneratorAction::PrimaryGeneratorAction()
    :G4VUserPrimaryGeneratorAction(),
     fParticleGunGPS(0),
     fParticleGun(0),
     fEbeamCumul(0),
     fFirst(true)
{
  fPrimaryGenMessenger = new PrimaryGeneratorMessenger(this);

  // To be used if reading primaries from a phase space file
  fParticleGun = new G4ParticleGun(1);

  // To be used if generating particles according to settings in macro
  fParticleGunGPS = new G4GeneralParticleSource();
}

PrimaryGeneratorAction::~PrimaryGeneratorAction()
{
  delete fPrimaryGenMessenger;

  // fInputFileName only exists if reading from phase space file
  if( fInputFileName ){
    delete fParticleGun;
    delete fReader;
    fReader = 0;
  }
  
  if( !fInputFileName ){
    delete fParticleGunGPS;
  }
}

void PrimaryGeneratorAction::GeneratePrimaries( G4Event* event )
{
  // Before first event, instantiate file reader if fInputFileName is not empty
  if( fFirst && fInputFileName != "" ){
      if( !fReader ) fReader = new FileReader( fInputFileName );
      fFirst = false;
  }

  // Generate primaries from a phase space file read in by FileReader
  if( fInputFileName != "" && fReader ){
    FileReader::Primary* primary = fReader->GetPrimary();
    
    G4String particleName = primary->GetName();
    
    G4ThreeVector pos = primary->GetPosition();

    // Shift the z position of the particle so it is relative to the room
    G4double posZ = (fPrimaryPosition - 4200.)*CLHEP::mm;
    pos = G4ThreeVector( pos.x(), pos.y(), posZ );

    G4ThreeVector mom = primary->GetMomentum();
    G4double energy = primary->GetEnergy();
    
    G4ParticleTable* pTable = G4ParticleTable::GetParticleTable();
    G4ParticleDefinition* particle = pTable->FindParticle( particleName.data() );
    
    fParticleGun->SetParticleDefinition( particle );
    fParticleGun->SetParticlePosition( pos );
    fParticleGun->SetParticleMomentumDirection( mom );
    fParticleGun->SetParticleEnergy( energy*CLHEP::MeV );
    
    fParticleGun->GeneratePrimaryVertex( event );
    
    fPrimaryEnergy = fParticleGun->GetParticleEnergy();
  }

  // Generate primaries with GeneralParticleSource if not reading from phase space file
  if( fInputFileName == "" ){
    // Confine the particle source to the first aluminium tube to avoid "leakage"
    fParticleGunGPS->GetCurrentSource()->GetPosDist()->ConfineSourceToVolume( "Source" );
    fParticleGunGPS->GeneratePrimaryVertex( event );
    
    fPrimaryEnergy = fParticleGunGPS->GetParticleEnergy();
  }

  fEbeamCumul += fPrimaryEnergy;
}