// %%%%%%%%%%
// G4 headers
// %%%%%%%%%%
#include "G4Event.hh"
#include "G4ParticleGun.hh"
#include "G4ParticleTable.hh"
#include "G4ParticleDefinition.hh"
#include "Randomize.hh"
#include "G4UnitsTable.hh"

// %%%%%%%%%%%%%
// gemc headers
// %%%%%%%%%%%%%
#include "MPrimaryGeneratorAction.h"
#include "detector.h"

// %%%%%%%%%%%
// C++ headers
// %%%%%%%%%%%
#include <iostream>
using namespace std;

MPrimaryGeneratorAction::MPrimaryGeneratorAction(gemc_opts opts)
{
 gemcOpt = opts;
 hd_msg        = gemcOpt.args["LOG_MSG"].args + " Beam Settings >> " ;
 input_gen     = gemcOpt.args["INPUT_GEN_FILE"].args;
 GEN_VERBOSITY = gemcOpt.args["GEN_VERBOSITY"].arg;

 particleTable = G4ParticleTable::GetParticleTable();

 setBeam(gemcOpt);

 particleGun = new G4ParticleGun(1);

 if(input_gen == "gemc_internal")
 {
    cout << endl << hd_msg << " Beam Type: "      << Particle->GetParticleName() << endl;
    cout << hd_msg << " Beam Momentum: "    << G4BestUnit(mom, "Energy") ;
    if(dmom > 0) cout << " +- " << G4BestUnit(dmom, "Energy") ;
    cout << endl;
    cout << hd_msg << " Beam Direction: (theta, phi) = (" << theta/deg << ", " << phi/deg << ")" ;
    if(dtheta > 0 || dphi > 0) cout << " +- (" << dtheta/deg << ", " << dphi/deg << ")" ;
    cout << " deg " << endl;
    cout << hd_msg << " Beam Vertex: (" << vx/cm << ", " << vy/cm << ", " << vz/cm << ")" ;
    if(dvx + dvy + dvz > 0) cout << " +- (" << dvx/cm << ", " << dvy/cm << ", " << dvz/cm << ")" ;
    cout << " cm " << endl;
 }


 if(NP>0)
 {
    cout << endl << hd_msg << " Luminosity Particle Type: "      << L_Particle->GetParticleName() << endl;
    cout << hd_msg << " Luminosity Particle Momentum: "    << G4BestUnit(L_Mom, "Energy") ;
    cout << endl;
    cout << hd_msg << " Luminosity Particle Direction: (theta, phi) = (" << L_Theta/deg << ", " << L_Phi/deg << ")" ;
    cout << " deg " << endl;
    cout << hd_msg << " Luminosity Particle Vertex: (" << L_vx/cm << ", " << L_vy/cm << ", " << L_vz/cm << ") cm" << endl;
    cout << hd_msg << " Number of Luminosity Particles: " << NP << endl;
    cout << hd_msg << " Luminosity Time Window: " << TWINDOW/ns << " nanoseconds." << endl ;
    cout << hd_msg << " Luminosity Time Between Bunches: " << TBUNCH/ns << " nanoseconds." << endl;
 }

}


void MPrimaryGeneratorAction::GeneratePrimaries(G4Event* anEvent)
{
 // internal generator. Particle defined by command line
 if(input_gen == "gemc_internal")
 {
    // Primary Particle
    particleGun->SetParticleDefinition(Particle);

    // 4-momenta
    Mom   = mom/MeV   + (2.0*G4UniformRand()-1.0)*dmom/MeV;
    Theta = theta/rad + (2.0*G4UniformRand()-1.0)*dtheta/rad;
    Phi   = phi/rad   + (2.0*G4UniformRand()-1.0)*dphi/rad;
    double mass = Particle->GetPDGMass();
    double akine = sqrt(Mom*Mom + mass*mass) - mass ;
    beam_dir = G4ThreeVector(cos(Phi/rad)*sin(Theta/rad), sin(Phi/rad)*sin(Theta/rad), cos(Theta/rad));
    particleGun->SetParticleEnergy(akine);
    particleGun->SetParticleMomentumDirection(beam_dir);

    // vertex
    Vx = vx/mm + (2.0*G4UniformRand()-1.0)*dvx/mm;
    Vy = vy/mm + (2.0*G4UniformRand()-1.0)*dvy/mm;
    Vz = vz/mm + (2.0*G4UniformRand()-1.0)*dvz/mm;
    beam_vrt = G4ThreeVector(Vx, Vy, Vz);
    particleGun->SetParticlePosition(beam_vrt);

    // Primary particle generated int the middle of Time window
    particleGun->SetParticleTime(TWINDOW/2);
    particleGun->GeneratePrimaryVertex(anEvent);
 }
 else
 // external generator: input file
 {
    // LUND format
    if(gformat == "LUND" && !gif.eof())
    {
       double tmp, px, py, pz;
       int NPART, pdef;
       gif >> NPART >> tmp >> tmp >> tmp >> tmp >> tmp >> tmp >> tmp >> tmp >> tmp;
       for(int p=0; p<NPART; p++)
       {
          gif >> tmp >> tmp >> tmp >> pdef >> tmp >> tmp >> px >> py >> pz >> tmp >> tmp >> Vx >> Vy >> Vz;

          // Primary Particle
          Particle = particleTable->FindParticle(pdef);
          particleGun->SetParticleDefinition(Particle);

          // 4-momenta
          G4ThreeVector pmom(px*GeV, py*GeV, pz*GeV);
          Mom = pmom.mag();
          Phi   = pmom.getPhi();
          Theta = pmom.getTheta();
          double mass = Particle->GetPDGMass();
          double akine = sqrt(Mom*Mom + mass*mass) - mass ;

          beam_dir = G4ThreeVector(cos(Phi/rad)*sin(Theta/rad), sin(Phi/rad)*sin(Theta/rad), cos(Theta/rad));
          particleGun->SetParticleEnergy(akine);
          particleGun->SetParticleMomentumDirection(beam_dir);

          // vertex
          beam_vrt = G4ThreeVector(Vx*cm, Vy*cm, Vz*cm);
          particleGun->SetParticlePosition(beam_vrt);

          // Primary particle generated int the middle of Time window
          particleGun->SetParticleTime(TWINDOW/2);
          particleGun->GeneratePrimaryVertex(anEvent);
          if(GEN_VERBOSITY > 3)
            cout << hd_msg << " Particle Number:  " << p << ", id=" << pdef
                 << "  Vertex=" << beam_vrt << "cm,  momentum=" << pmom/GeV << " GeV" << endl;
       }
    }
 }



 // Luminosity Particles
 int NBUNCHES   = (int) floor(TWINDOW/TBUNCH);
 int PBUNCH     = (int) floor((double)NP/NBUNCHES);

 particleGun->SetParticleDefinition(L_Particle);
 particleGun->SetParticlePosition(L_beam_vrt);
 double L_mass = L_Particle->GetPDGMass();
 double L_akine = sqrt(L_Mom*L_Mom + L_mass*L_mass) - L_mass ;

 particleGun->SetParticleEnergy(L_akine);
 particleGun->SetParticleMomentumDirection(L_beam_dir);
 for(int b=0; b<NBUNCHES; b++)
 {
    for(int p=0; p<PBUNCH; p++)
    {
       particleGun->  SetParticleTime(TBUNCH*b);
       particleGun->GeneratePrimaryVertex(anEvent);
    }
 }

}



void MPrimaryGeneratorAction::setBeam(gemc_opts)
{
 string hd_msg    = gemcOpt.args["LOG_MSG"].args + " Beam Settings >> " ;

 if(input_gen == "gemc_internal")
 {

    // %%%%%%%%%%%%
    // Primary Beam
    // %%%%%%%%%%%%
    string beam     = gemcOpt.args["BEAM_P"].args;
    string vert     = gemcOpt.args["BEAM_V"].args;
    string spread_b = gemcOpt.args["SPREAD_P"].args;
    string spread_v = gemcOpt.args["SPREAD_V"].args;

    string blockstr;
    int ppos = beam.find(",");
    blockstr.assign(beam, 0, ppos);

    Particle = particleTable->FindParticle(blockstr);

    if(!Particle)
    {
       if(blockstr == "show_all")
       {
         for(int i=0; i<particleTable->entries(); i++) 
            cout << hd_msg << " g4 particle: " << particleTable->GetParticleName(i) << endl;;
      }
      else
         cout << hd_msg << " Particle " << blockstr << " not found in G4 table. Exiting" << endl << endl;

      exit(0);
    }

    // %%%%%%%%
    // Momentum
    // %%%%%%%%
    // energy
    beam.assign(beam, ppos+1, beam.size());
    ppos = beam.find(",");
    blockstr.assign(beam, 0, ppos);
    mom = get_number(blockstr);
    // theta angle
    beam.assign(beam, ppos+1, beam.size());
    ppos = beam.find(",");
    blockstr.assign(beam, 0, ppos);
    theta =  get_number(blockstr);
    // phi angle
    beam.assign(beam, ppos+1, beam.size());
    ppos = beam.find("\"");
    blockstr.assign(beam, 0, ppos);
    phi =  get_number(blockstr);

    // %%%%%%
    // Spread
    // %%%%%%
    // delta energy
    ppos = spread_b.find(",");
    blockstr.assign(spread_b, 0, ppos);
    dmom = get_number(blockstr);
    // delta theta angle
    spread_b.assign(spread_b, ppos+1, spread_b.size());
    ppos = spread_b.find(",");
    blockstr.assign(spread_b, 0, ppos);
    dtheta =  get_number(blockstr);
    // delta phi angle
    spread_b.assign(spread_b, ppos+1, spread_b.size());
    ppos = spread_b.find(",");
    blockstr.assign(spread_b, 0, ppos);
    dphi =  get_number(blockstr);


    // %%%%%%
    // Vertex
    // %%%%%%
    // units
    ppos = vert.find(")");
    string units;
    units.assign(vert, ppos+1, vert.size() - ppos);
    // vx
    ppos = vert.find("(");
    vert.assign(vert, ppos+1, vert.size()); 
    ppos = vert.find(",");
    blockstr.assign(vert, 0, ppos);
    vx = get_number(blockstr + "*" + units); 
    // vy
    ppos = vert.find(",");
    vert.assign(vert, ppos+1, vert.size());
    ppos = vert.find(",");
    blockstr.assign(vert, 0, ppos);
    vy = get_number(blockstr + "*" + units);
    // vz
    ppos = vert.find(",");
    vert.assign(vert, ppos+1, vert.size());
    ppos = vert.find(",");
    blockstr.assign(vert, 0, ppos);
    vz = get_number(blockstr + "*" + units);


    // %%%%%%%%%%%%%%%%%%%%%%%%%
    // Vertex Spread (cartesian)
    // %%%%%%%%%%%%%%%%%%%%%%%%%
    // units
    ppos = spread_v.find(")");
    units.assign(spread_v, ppos+1, spread_v.size() - ppos);
    // dvx
    ppos = spread_v.find("(");
    spread_v.assign(spread_v, ppos+1, spread_v.size()); 
    ppos = spread_v.find(",");
    blockstr.assign(spread_v, 0, ppos);
    dvx = get_number(blockstr + "*" + units); 
    // dvy
    ppos = spread_v.find(",");
    spread_v.assign(spread_v, ppos+1, spread_v.size());
    ppos = spread_v.find(",");
    blockstr.assign(spread_v, 0, ppos);
    dvy = get_number(blockstr + "*" + units);
    // dvz
    ppos = spread_v.find(",");
    spread_v.assign(spread_v, ppos+1, spread_v.size());
    ppos = spread_v.find(",");
    blockstr.assign(spread_v, 0, ppos);
    dvz = get_number(blockstr + "*" + units);
 }
 else
 {
    gformat.assign(  input_gen, 0, input_gen.find(",")) ;
    gfilename.assign(input_gen,    input_gen.find(",") + 1, input_gen.size()) ;
    cout << hd_msg << " Opening  " << gformat << " file: " << TrimSpaces(gfilename).c_str() << endl;
    gif.open(TrimSpaces(gfilename).c_str());


    if(!gif)
    {
       cerr << hd_msg << " Can't open output file " << TrimSpaces(gfilename).c_str() << ". Exiting. " << endl;
       exit(1);
    }
 }


 // %%%%%%%%%%%%%%%
 // Luminosity Beam
 // %%%%%%%%%%%%%%%
 string L_beam     = gemcOpt.args["LUMI_P"].args;
 string L_vert     = gemcOpt.args["LUMI_V"].args;

 string L_blockstr;
 int L_ppos = L_beam.find(",");
 L_blockstr.assign(L_beam, 0, L_ppos);

 L_Particle = particleTable->FindParticle(L_blockstr);

 if(!L_Particle)
 {
    if(L_blockstr == "show_all")
    {
      for(int i=0; i<particleTable->entries(); i++) 
         cout << hd_msg << " g4 particle: " << particleTable->GetParticleName(i) << endl;
    }
    else
       cout << hd_msg << " Particle " << L_blockstr << " not found in G4 table. Exiting" << endl << endl;

    exit(0);
 }

 // %%%%%%%%%%%%%%%%%%%
 // Luminosity Momentum
 // %%%%%%%%%%%%%%%%%%%
 // energy
 L_beam.assign(L_beam, L_ppos+1, L_beam.size());
 L_ppos = L_beam.find(",");
 L_blockstr.assign(L_beam, 0, L_ppos);
 L_Mom = get_number(L_blockstr);
 // theta angle
 L_beam.assign(L_beam, L_ppos+1, L_beam.size());
 L_ppos = L_beam.find(",");
 L_blockstr.assign(L_beam, 0, L_ppos);
 L_Theta =  get_number(L_blockstr);
 // phi angle
 L_beam.assign(L_beam, L_ppos+1, L_beam.size());
 L_ppos = L_beam.find("\"");
 L_blockstr.assign(L_beam, 0, L_ppos);
 L_Phi =  get_number(L_blockstr);

 L_beam_dir = G4ThreeVector(cos(L_Phi/rad)*sin(L_Theta/rad), sin(L_Phi/rad)*sin(L_Theta/rad), cos(L_Theta/rad));

 // %%%%%%%%%%%%%%%%%
 // Luminosity Vertex
 // %%%%%%%%%%%%%%%%%
 // units
 L_ppos = L_vert.find(")");
 string L_units;
 L_units.assign(L_vert, L_ppos+1, L_vert.size() - L_ppos);
 // vx
 L_ppos = L_vert.find("(");
 L_vert.assign(L_vert, L_ppos+1, L_vert.size()); 
 L_ppos = L_vert.find(",");
 L_blockstr.assign(L_vert, 0, L_ppos);
 L_vx = get_number(L_blockstr + "*" + L_units); 
 // vy
 L_ppos = L_vert.find(",");
 L_vert.assign(L_vert, L_ppos+1, L_vert.size());
 L_ppos = L_vert.find(",");
 L_blockstr.assign(L_vert, 0, L_ppos);
 L_vy = get_number(L_blockstr + "*" + L_units);
 // vz
 L_ppos = L_vert.find(",");
 L_vert.assign(L_vert, L_ppos+1, L_vert.size());
 L_ppos = L_vert.find(",");
 L_blockstr.assign(L_vert, 0, L_ppos);
 L_vz = get_number(L_blockstr + "*" + L_units);

 L_beam_vrt = G4ThreeVector(L_vx, L_vy, L_vz);

 // %%%%%%%%%%%%%%%%%%%%%
 // Luminosity Parameters
 // %%%%%%%%%%%%%%%%%%%%%
 string L_pars = gemcOpt.args["LUMI_EVENT"].args;
 // NP
 L_ppos = L_pars.find(",");
 L_blockstr.assign(L_pars, 0, L_ppos);
 NP = atoi(L_blockstr.c_str());
 // TWINDOW
 L_ppos = L_pars.find(",");
 L_pars.assign(L_pars, L_ppos+1, L_pars.size());
 L_ppos = L_pars.find(",");
 L_blockstr.assign(L_pars, 0, L_ppos);
 TWINDOW = get_number(L_blockstr);
 // TBUNCH
 L_ppos = L_pars.find(",");
 L_pars.assign(L_pars, L_ppos+1, L_pars.size());
 L_ppos = L_pars.find(",");
 L_blockstr.assign(L_pars, 0, L_ppos-1);
 TBUNCH = get_number(L_blockstr);

}


MPrimaryGeneratorAction::~MPrimaryGeneratorAction()
{
 delete particleGun;
 gif.close();
}