//  InROOT_Gen.cc
// Contact person: R. Bonventre <rbonv@hep.upenn.edu>
// See  InROOT_Gen.hh for more details
//———————————————————————//

#include <RAT/InROOT_Gen.hh>
#include <RAT/Log.hh>
#include <RAT/ParentPrimaryVertexInformation.hh>
#include <RAT/Factory.hh>
#include <RAT/GLG4StringUtil.hh>
#include <RAT/GLG4TimeGen.hh>
#include <RAT/GLG4PosGen.hh>
#include <RAT/DS/Entry.hh>
#include <RAT/DS/MC.hh>
#include <RAT/DS/MCParticle.hh>

#include <G4Event.hh>
#include <G4PrimaryVertex.hh>
#include <G4PrimaryParticle.hh>
#include <G4ParticleTable.hh>
#include <G4ThreeVector.hh>
#include <G4RunManager.hh>

#include <TTree.h>
#include <TFile.h>
#include <TROOT.h>

#include <stdexcept>
#include <vector>
#include <string>
using namespace std;

namespace RAT {

void InROOT_Gen::GenerateEvent(G4Event *event)
{
  fTTree->GetEntry(fCurrentEvent);
  DS::MC& mc = fDS->GetMC();

  int vertexCount = 0;
  for (size_t i=0;i<mc.GetMCParticleCount();i++){
    int vid = mc.GetMCParticle(i).GetVertexID();
    if (vid >= vertexCount)
      vertexCount = vid+1;
  }
  std::vector<G4PrimaryVertex*> vertices(vertexCount);

  G4ThreeVector position;
  for (size_t i=0;i<mc.GetMCParentCount();i++){
    const DS::MCParticle& mcp = mc.GetMCParent(i);
    G4PrimaryParticle* particle = new G4PrimaryParticle(mcp.GetPDGCode(),
        mcp.GetMomentum().X(),mcp.GetMomentum().Y(),mcp.GetMomentum().Z());
    particle->SetMass(G4ParticleTable::GetParticleTable()->FindParticle(mcp.GetPDGCode())->GetPDGMass());

    ParentPrimaryVertexInformation *vertinfo = new ParentPrimaryVertexInformation();
    vertinfo->SetVertexParentParticle(particle);

    if (!vertices[mcp.GetVertexID()]){
      if (fPosGen)
        fPosGen->GeneratePosition(position);
      else
        position = G4ThreeVector(mcp.GetPosition().X(),mcp.GetPosition().Y(),mcp.GetPosition().Z());
      vertices[mcp.GetVertexID()] = new G4PrimaryVertex(position,NextTime()+mcp.GetTime());
    }
    vertices[mcp.GetVertexID()]->SetUserInformation(vertinfo);
  }

  for (size_t i=0;i<mc.GetMCParticleCount();i++){
    const DS::MCParticle& mcp = mc.GetMCParticle(i);
    G4PrimaryParticle* particle = new G4PrimaryParticle(mcp.GetPDGCode(),
        mcp.GetMomentum().X(),mcp.GetMomentum().Y(),mcp.GetMomentum().Z());
    particle->SetMass(G4ParticleTable::GetParticleTable()->FindParticle(mcp.GetPDGCode())->GetPDGMass());

    if (!vertices[mcp.GetVertexID()]){
      if (fPosGen)
        fPosGen->GeneratePosition(position);
      else
        position = G4ThreeVector(mcp.GetPosition().X(),mcp.GetPosition().Y(),mcp.GetPosition().Z());
      vertices[mcp.GetVertexID()] = new G4PrimaryVertex(position,NextTime()+mcp.GetTime());
    }
    vertices[mcp.GetVertexID()]->SetPrimary(particle);
  }

  for (size_t i=0;i<vertices.size();i++)
    event->AddPrimaryVertex(vertices[i]);

  fLastEventTime = mc.GetMCTime();
  fCurrentEvent++;

  if (fCurrentEvent >= fNumEvents || (fMaxEvent > 0 && fCurrentEvent >= fMaxEvent)){
    // we are out of events, stop the simulation after this one
    G4RunManager::GetRunManager()->AbortRun(/*softabort*/ true);
  }
}

void InROOT_Gen::ResetTime(double offset)
{
  if (fCurrentEvent < fNumEvents){
    if (fTimeGen){
      double eventTime = fTimeGen->GenerateEventTime();
      fNextTime = eventTime + offset;
    }else{
      fTTree->GetEntry(fCurrentEvent);
      DS::MC& mc = fDS->GetMC();
      fNextTime = mc.GetMCTime()-fLastEventTime + offset;
    }
  }else{
    fNextTime = 1e9;
  }
}

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

  std::string filename;
  if (nArgs >= 5){
    int num_skip;
    istringstream (parts[4]) >> num_skip;
    fCurrentEvent = num_skip;
  }
  if (nArgs >= 4){
    istringstream (parts[3]) >> fMaxEvent;
    if (fMaxEvent > 0)
      fMaxEvent += fCurrentEvent;
  }
  if (nArgs >= 3){
    if (parts[2] == "default"){
      fTimeGen = 0;
    }else{
      fTimeGen = GlobalFactory<GLG4TimeGen>::New(parts[2]);
    }
  }
  if (nArgs >= 2){
    if (parts[1] == "default"){
      fPosGen = 0;
    }else{
      fPosGen = GlobalFactory<GLG4PosGen>::New(parts[1]);
    }
  }
  if (nArgs >= 1){
    filename = parts[0];
  }else{
	  G4Exception(__FILE__, "Invalid Parameter", FatalException, ("inroot generator syntax error: '"+
          state+"' does not have a filename").c_str());
	}

  fStateStr = state;
  TFile *file = new TFile(filename.c_str());
  gROOT->cd(0);
  fTTree = (TTree*) ((TTree *) file->Get("T"))->CloneTree();
  fTTree->SetDirectory(0);

  file->Close();

  fNumEvents = fTTree->GetEntries();
  if (!fNumEvents)
	  G4Exception(__FILE__, "Invalid Parameter", FatalException, ("File '"+
          filename+"' is empty").c_str());

  info << "InROOT_Gen: Reading from " << filename << newline;

  fDS = new DS::Entry();
  fTTree->SetBranchAddress("ds", &fDS);

  fNumEvents = (int) fTTree->GetEntries();
}

G4String InROOT_Gen::GetState() const
{
	return fStateStr;
}

void InROOT_Gen::SetTimeState(G4String state)
{
  if (fTimeGen)
      fTimeGen->SetState(state);
  else
      warn << "InROOT_Gen error: Cannot set time state, no time generator selected\n";
}

G4String InROOT_Gen::GetTimeState() const
{
  if (fTimeGen)
     return fTimeGen->GetState();
  else
     return G4String("InROOT_Gen error: no time generator selected");
}

void InROOT_Gen::SetPosState(G4String state)
{
  if (fPosGen)
     fPosGen->SetState(state);
  else
     warn << "InROOT_Gen error: Cannot set position state, no position generator selected\n";
}

G4String InROOT_Gen::GetPosState() const
{
  if (fPosGen)
     return fPosGen->GetState();
  else
     return G4String("InROOT_Gen error: no pos generator selected");
}



} // namespace RAT