#include "ISimHitMergerMain.hxx"
#include "IVInputFile.hxx"
#include "ICOMETOutput.hxx"
#include "IG4VHit.hxx"
#include "IG4HitSegment.hxx"
#include "IG4HitCalo.hxx"
#include "IG4HitGas.hxx"
#include "IG4Trajectory.hxx"
#include "IMergedEventList.hxx"
#include "IIntegerDatum.hxx"
#include "IProfileData.hxx"
#include "ITimeShiftedTrajectory.hxx"
#include "ITimeShiftedSegHit.hxx"
#include "ITimeShiftedCaloHit.hxx"
#include "ITimeShiftedGasHit.hxx"
#include "ITimeShiftedMonHit.hxx"
#include "ITimeShiftedCherenHit.hxx"
#include "IVTimeShiftGenerator.hxx"
#include "ITimeShiftedGasRegionDR.hxx"

#include "IOADatabase.hxx"
#include "IPAIDeltaRay.hxx"

#include <utility>
#include <limits>
#include <sstream>

#include <TRandom.h>

using namespace COMET;

ISimHitMergerMain::ISimHitMergerMain():
                        fCurrentInput(NULL),
                        fTrajectories(NULL),
                        fTimeShiftGenerator(NULL),
                        fMergedEvents(0),
			fMergedPOTs(0),
			fDesiredEventsPerBunch(0),
                        fDesiredEvents(0),
                        fNoBunches(1),
                        fUseRandomBunch(true),
                        fMinimumTime(-std::numeric_limits<double>::max()),
                        fMaximumTime(std::numeric_limits<double>::max()),
                        fGlobalShift(0),
                        fBunchSeparation(1.17e3),
                        fOutputEventNo(0),
                        fSaveTrajectories(true),
                        fProfileEvents(false),
                        fWasEventsSet(false),
                        fUseAllEvents(false),
                        fWasRunNoSet(false),
                        fWasTotalEventsSet(false)
                        {
                        }

void ISimHitMergerMain::PrepareNewOutputEvent(){
    // Reset the number of events
    fMergedEvents=0;
    // Reset the number of protons
    fMergedPOTs=0;
    // Clear the contents of the map
    fHitContainers.clear();
    // Delete the trajectories if they exist
    fTrajectories=NULL;
    // Reset the event contributors
    fEventContributors = new IMergedEventList("event_contributors");
}

ISimHitMergerMain::~ISimHitMergerMain(){
}

void ISimHitMergerMain::Usage()
{
  std::cout << 
    " -------------------------------------------------------------------------------------\n"
    "  The full details of SimHitMerger can be found in the GitLab wiki page below:\n"
    "       https://gitlab.in2p3.fr/comet/ICEDUST_packages/-/wikis/SimHitMerger\n"
    " -------------------------------------------------------------------------------------\n"
    "\n"
    "    -O no-trajectories                   Don't save trajectories.\n" 
    "\n" 
    "    -O events=<NoProtons>                Set the number of input events per bunch.\n" 
    "    -O protons=<NoProtons>               [Obsolete] Set the number of events in total for all the bunches.\n" 
    "\n" 
    "    -O all-events                        Use all availible events in one event.\n" 
    "                                         NOTE: This option cannot be used with the\n" 
    "                                               previous option or with the geometry \n" 
    "                                               options (-G or -g)\n" 
    "    -O all-protons                       [Obsolete] Equivalent to all-events.\n"
    "\n" 
    "    -O bunches=<NoBunches>               Set the number of bunches to include in an event.\n" 
    "\n" 
    "    -O bunch-separation=<Separation>     Set the separation of bunches (in nanoseconds)\n"
    "\n" 
    "    -O global-shift=<shift>     Set the global time shift (in nanoseconds)\n"
    "\n" 
    "    -O minimum-time=<time>      Set the minimum allowed hit time (in nanoseconds).\n"
    "                                This hit time is calculated after the time is shifted.\n"
    "\n" 
    "    -O maximum-time=<time>      Set the maximum allowed hit time (in nanoseconds).\n"
    "                                This hit time is calculated after the time is shifted.\n"
    "\n" 
    "    -O no-random-bunch                   Place the nth proton in bunch number:\n"
    "                                           (n % NoBunches)\n"
    "                                         instead of a random bunch. This is useful for bunch\n"
    "                                         train construction from bunch-events.\n"
    "\n" 
    "    -O time=<GeneratorName [arguments]>  Set the type of time shift generator to use.\n"
    "                                         Necessary arguments depend on the generator but \n"
    "                                         should be white space separated.\n" 
    "\n" 
    "    -O run=<RunNo>                       Set the run number\n" 
    "\n" 
    "    -O detector=<SD-name>                Add the name of a Sensitive Detector such that\n"
    "                                         only events containing hits in SDs named in this\n"
    "                                         list are merged. "
    "\n" 
    "    -O only-detectors                    Equivalent to multiple '-O detector=...' enabling\n"
    "                                         only real detectors: ECAL, CDC, CTH, StrawTrk\n"
    "\n" 
    "    -O profile-events                    Turn on event profiling, which is persisted to a\n"
    "                                         DataVector."
    "\n"
    "    -O random-seed=<seed>                Set the seed of the random generator (gRandom).\n"
    "\n"
	    << std::endl;
  
  // Cherry picking
  fCherryPickMan.Usage();
}

bool ISimHitMergerMain::SetOption(std::string option, std::string value)
{
  std::stringstream val(value);
  if (option == "events") {
    long int pot;
    val>>pot;
    SetDesiredEventsPerBunch(pot);
    fWasEventsSet=true;
  }else if (option == "protons") {
    long int pot;
    val>>pot;
    SetDesiredEvents(pot);
    fWasEventsSet=true;
    fWasTotalEventsSet=true;
  }else if (option == "all-events" ||
	    option == "all-protons" ) {
    // Use all the input protons
    SetDesiredEvents(std::numeric_limits<long int>::max());
    fUseAllEvents=true;
  }else if (option == "time") {
    std::string gen_type, args;
    val>>gen_type>>args;
    IVTimeShiftGenerator* generator=IVTimeShiftGenerator::Make(gen_type,args);
    if(!generator) return false;
    SetTimeShiftGenerator(generator);
  }else if (option == "run") {
    int run;
    val>>run;
    SetRunNumber(run);
    fWasRunNoSet=true;
  }else if (option == "bunches") {
    int bunches;
    val>>bunches;
    SetNumberBunches(bunches);
  }else if (option == "bunch-separation") {
    double separation;
    val>>separation;
    SetBunchSeparation(separation);
  }else if (option == "global-shift") {
    double shift;
    val>>shift;
    SetGlobalShift(shift);
  }else if (option == "minimum-time") {
    double min_time;
    val>>min_time;
    SetMinimumTime(min_time);
  }else if (option == "maximum-time") {
    double max_time;
    val>>max_time;
    SetMaximumTime(max_time);
  }else if (option == "no-random-bunch"){
    // Default behaviour is true, set to false if flag option used
    SetUseRandomBunches(false);
  }else if (option == "no-trajectories") {
    SetSaveTrajectories(false);
  }else if (option == "detector") {
    AddValidDetector(value);
  }else if (option == "only-detectors") {
    AddValidDetector("CDC");
    AddValidDetector("ECAL");
    AddValidDetector("CTH");
    AddValidDetector("StrawTrk");
  }else if (option == "profile-events"){
    SetProfileEvents(true);
  }else if (option == "random-seed"){
    int seed;
    val >> seed;
    gRandom->SetSeed(seed);
  }
  // Otherwise, for the cherry picking ...
  else return fCherryPickMan.SetOption(option, value);
	
  return true;
}

bool ISimHitMergerMain::CheckOptions()
{
  // Check if we set number of Events to value or to all availible
  if(!fWasEventsSet && !fUseAllEvents){
    COMETError("You need to set the number of events in each bunch : \n"
	       "  -O events=<NoProtons>\n"
	       "OR \n"
	       "  -O all-events");
    return false;
  }
  // Ensure that we only set use all Events or specified number, but not both
  if (fUseAllEvents && fWasEventsSet){
    COMETError("Use only ONE of the following options, not BOTH: \n"
	       "  -O events=<NoEvents>\n"
	       "OR \n"
	       "  -O all-events");
    return false;
  }
  if(fWasTotalEventsSet){
    COMETLog("CAUTION: You set the number of input events for total bunches, NOT per bunch.\n"
	     "         Use -O events=[NoEvents/Bunch] if you need that per bunch.");
    return true;
  }
  else if(not fUseAllEvents){
    // The number per bunch -> total number
    fDesiredEvents = fDesiredEventsPerBunch * fNoBunches;
  }
  if(!fWasRunNoSet){
    COMETLog("No run number set (-O run=<RunNo>) so defaulting to 1");
  }
  return true;
}

void ISimHitMergerMain::Initialize()
{
  fCherryPickMan.Initialize();
  PrepareNewOutputEvent();
}

void ISimHitMergerMain::NewEventSource(const COMET::IVInputFile* input){
  fCurrentInput=input;
}

bool ISimHitMergerMain::ProcessEvent(ICOMETEvent& in_event){

   // Extract hits and (optionally the trajectories)from the current event and
   // time shift
   double proton_time_shift=GetProtonTimeShift(in_event);
   bool event_contributed_hits = ExtractHits(in_event,proton_time_shift);

   // Increment the counters
   int n_proton = GetNumberOfProtons(in_event);
   fMergedPOTs  += n_proton;
   fMergedEvents++;

   // Log this input event as having contributed to the merged event
   if (event_contributed_hits) {
       fEventContributors->Add(fCurrentInput,
                               in_event,
                               n_proton,
                               proton_time_shift,
                               &fNewIds);
       COMETNamedTrace("HitMerger","Taken hits for event "
               <<in_event.GetEventId() <<" in input file '"
               <<fCurrentInput->GetFilename());
   }

   if (fMergedEvents >= fDesiredEvents){
       COMETNamedTrace("HitMerger","Writing out event");
       // Swap in the merged hit data
       ReplaceEventContents(in_event);
       // Get ready for the next event
       ++fOutputEventNo;
       PrepareNewOutputEvent();
       return true;
   }
   // We don't want to write a new output event
   // (which would be a copy of the current input event)
   return false;
}

bool ISimHitMergerMain::ExtractHits(const COMET::ICOMETEvent& event,
                                    double proton_time_shift){
    bool had_hits=false;
    if(fSaveTrajectories){
        fNewIds.clear();
        fParentIdMap.clear();
    }
    
    COMET::IHandle<COMET::IDataVector> hits = event.Get<COMET::IDataVector>("truth/g4Hits");
    //COMET::IHandle<COMET::IDataVector> deltaRays = event.Get<COMET::IPAIDeltaRayContainer>("truth/PAIDeltaRays/GasRegion");
    COMET::IHandle<COMET::IDataVector> paiVect = event.Get<COMET::IDataVector>("truth/PAIPhysics/");

    if (hits){
        COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories
            = event.Get<COMET::IG4TrajectoryContainer>("truth/G4Trajectories");
        for (COMET::IDataVector::iterator h = hits->begin();
                h != hits->end(); ++h) {
            COMET::IHandle<COMET::IG4HitContainer> hitContainer =
                (*h)->Get<COMET::IG4HitContainer>(".");
            // Check that a hit container if found
            if (hitContainer) {
                // If so, check it is in the allowed hit container names
                if(fDetectorList.empty() || 
                   fDetectorList.count(hitContainer->GetName()) > 0){
                    // If it is, check it has hits within the allowed time 
                    // window
                    bool container_had_hits=ContainerHasHits(*hitContainer,
                                                             proton_time_shift);
                    // If it does have hits, copy them over and mark the event 
                    // as having hits
                    if(container_had_hits){
                        CopyOverHits(*hitContainer,
                                     proton_time_shift,
                                     trajectories);
                        had_hits = true;
                    }
                }
            }
        }
    }

    if (paiVect){      
      for (COMET::IDataVector::iterator d = paiVect->begin(); d != paiVect->end(); ++d) {   
	COMET::IHandle<COMET::IPAIDeltaRayContainer> DeltaRayContainer = 
	  (*d)->Get<COMET::IPAIDeltaRayContainer>(".");
	// Check that a delta-ray container if found
	if (DeltaRayContainer) {	  
	  // If so, check it is in the allowed hit container names
	  if(fDetectorList.empty() || fDetectorList.count(DeltaRayContainer->GetName()) > 0){

	    // If it is, check it has hits within the allowed time 
	    // window
	    bool container_had_hits = ContainerHasHits(*DeltaRayContainer,
						       proton_time_shift);
	    // If it does have hits, copy them over and mark the event 
	    // as having hits
	    if(container_had_hits){
	      CopyOverDeltaRays(*DeltaRayContainer,
				proton_time_shift);
				//trajectories);	      
	      had_hits = true;	     
	    }
	  }	  	  
	}
      }
    }    
    return had_hits;    
}

bool ISimHitMergerMain::ContainerHasHits(
        const COMET::IG4HitContainer& hitContainer,
        double proton_time_shift){
    // Iterate through all the hits and check if they fall within the allowed 
    // time window
    for(IG4HitContainer::const_iterator i_hit=hitContainer.begin();
                                        i_hit!=hitContainer.end(); ++i_hit){
        // If the current hit is within the time window, then return true
        double new_time = (*i_hit)->GetPosT() + proton_time_shift;
        if ((new_time <= fMaximumTime) && (new_time >= fMinimumTime)) 
            return true;
    }
    // If no hit like this is found, return false
    return false;
}

bool ISimHitMergerMain::ContainerHasHits(
        const COMET::IPAIDeltaRayContainer& deltaRayContainer,
        double proton_time_shift){
    // Iterate through all the hits and check if they fall within the allowed 
    // time window
    for(IPAIDeltaRayContainer::const_iterator i_dr=deltaRayContainer.begin(); i_dr!=deltaRayContainer.end(); ++i_dr){
        // If the current hit is within the time window, then return true
        double new_time = (*i_dr)->GetPosT() + proton_time_shift;
        if ((new_time <= fMaximumTime) && (new_time >= fMinimumTime)) 
            return true;
    }
    // If no hit like this is found, return false
    return false;
}

void ISimHitMergerMain::ExtractTrajectoriesFor(
        std::set<int>& track_ids,
        IHandle<IG4TrajectoryContainer> trajectories,
        double proton_time_shift){
    // Initialize the cached trajectories if they do not already exist
    if(!fTrajectories) fTrajectories=new IG4TrajectoryContainer;

    // Now get all parent/ancestor track IDs
    std::set<int> all_requests;
    for(std::set<int>::const_iterator i_id=track_ids.begin();
                                      i_id!=track_ids.end();
                                    ++i_id){
        all_requests.insert(*i_id);
        MarkParentTrajectories(*trajectories,*i_id,all_requests);
    }
    COMETNamedTrace("HitMergerTrajectories","Total marked trajectories: "
                    << all_requests.size());

    // Remove trajectories that have already been added for previous hits (but
    // in the same event)
    track_ids.clear();
    for(std::set<int>::const_iterator i_id = all_requests.begin();
                                      i_id!= all_requests.end();
                                    ++i_id){
        if(fNewIds.count(*i_id)==0){
            track_ids.insert(*i_id);
            COMETNamedTrace("HitMergerTrajectories","Requested Id: "
                             << *i_id <<" will be added" );
        } else COMETNamedTrace("HitMergerTrajectories","Requested Id: "
                               << *i_id << " already added" );
    }
    // Copy the trajectories into the new container and reset their track ID
    CopyTrajectories(track_ids,*trajectories,proton_time_shift);
}

void ISimHitMergerMain::CopyTrajectories(
        const std::set<int>& requests,
        const IG4TrajectoryContainer& originalTrajectories,
        double proton_time_shift){
    COMETNamedTrace("HitMergerTrajectories","Begin copying");
    // Get the first ID of the new trajectories as the number of trajectories
    // already cached
    int firstID=fTrajectories->size() + 1;
    // Iterate through the requested track IDs from the original trajectories
    for(std::set<int>::const_iterator rqst = requests.begin();
                                      rqst!= requests.end();
                                    ++rqst){
        COMETNamedTrace("HitMergerTrajectories","CopyTrajectory: "<<*rqst);
        // Get the original trajectory
        IHandle<IG4Trajectory> old_traj =
            originalTrajectories.GetTrajectory(*rqst);
        if(!old_traj){
            // This can happen - not saving all trajectories for secondaries
            COMETVerbose("Trajectory not saved: Assigning ID -99 for track ID "
                         << *rqst);
            fNewIds[*rqst]=-99;
        } else {
            // Time shift the trajectory
            ITimeShiftedTrajectory new_traj(old_traj,proton_time_shift);
            int newID = fTrajectories->size() + 1;
            // Set the new track ID
            new_traj.SetTrackId(newID);
            (*fTrajectories)[newID]=new_traj;
            fNewIds[old_traj->GetTrackId()]=newID;
            // Increment the new track ID counter
            ++newID;
        }
    }

    // Reset the parent IDs
    for(IG4TrajectoryContainer::iterator i_traj=fTrajectories->find(firstID);
                                         i_traj!=fTrajectories->end();
                                       ++i_traj){
        int old_parent=i_traj->second.GetParentId();
        int new_parent=fNewIds[old_parent];
        i_traj->second.SetParentId(new_parent);
    }
}

void ISimHitMergerMain::MarkParentTrajectories(
        const IG4TrajectoryContainer& trajectories, int search,
        std::set<int>& all_requests){
    COMETNamedTrace("HitMergerTrajectories","MarkParentTrajectories: "<<search);
    IHandle<IG4Trajectory> traj = trajectories.GetTrajectory(search);
    // If there is not trajectory for this track can't find parents ...
    if (!traj) return ;
    int parent=traj->GetParentId();

    // If it has no parent or parent is already marked return
    if (parent <1 ||
        all_requests.count(parent)!=0 ||
        fParentIdMap.count(parent)!=0) return;
    // Add to the parent - child map and the "already requested" list
    fParentIdMap[search]=parent;
    all_requests.insert(parent);

    // Recurse:
    MarkParentTrajectories(trajectories,parent,all_requests);
}

void ISimHitMergerMain::CopyOverHits(const IG4HitContainer& hitContainer,
        double proton_time_shift, IHandle<IG4TrajectoryContainer> trajectories){
    // Set the output container to null
    IG4HitContainer* out_container=NULL;
    // Try to find the container by name
    HitContainerMap::iterator i_out_hc 
        = fHitContainers.find(hitContainer.GetName());
    // If its not found, make a new one and insert it into the map
    if(i_out_hc == fHitContainers.end()){
        out_container= new IG4HitContainer(hitContainer.GetName());
        fHitContainers.insert(std::make_pair( hitContainer.GetName(),out_container));
    }else {
        out_container=i_out_hc->second;
    }
    // Iterate through the input hit container and copy over the hits
    for(IG4HitContainer::const_iterator i_hit=hitContainer.begin();
                                        i_hit!=hitContainer.end();
                                        ++i_hit){
        // Check if we are saving trajectories
        bool saveTrajs = fSaveTrajectories && trajectories;
        // Check if we are within the timie limit
        double hTime = (*i_hit)->GetPosT() + proton_time_shift;
        // Skip this hit if it is not in the time limit
        if ((hTime > fMaximumTime) || (hTime < fMinimumTime)) continue;
        // Try casting to known g4hit types and adding the time shifted version
        // if it is found
        // TODO clean this up with polymophism/templating
        if (const IG4HitMonitor* g4Mon=dynamic_cast<IG4HitMonitor*>(*i_hit)){
            // Create the time shifted hit segment
            ITimeShiftedMonHit new_hit = ITimeShiftedMonHit(*g4Mon, proton_time_shift);
            // Save the relevant trajectories from this hit if possible/wanted
            if(saveTrajs) ManageHitTracks(new_hit, proton_time_shift, trajectories);
            // Add the new hit
            out_container->push_back(new IG4HitMonitor(new_hit));
        } else if (const IG4HitCalo* g4Cal=dynamic_cast<IG4HitCalo*>(*i_hit)){
            // Create the time shifted calo hit
            ITimeShiftedCaloHit new_hit = ITimeShiftedCaloHit(*g4Cal, proton_time_shift);
            // Save the relevant trajectories from this hit if possible/wanted
            if(saveTrajs) ManageHitTracks(new_hit, proton_time_shift, trajectories);
            // Add the new hit
            out_container->push_back(new IG4HitCalo(new_hit));
        } else if (const IG4HitGas* g4Gas=dynamic_cast<IG4HitGas*>(*i_hit)){
            // Create the time shifted gas hit
            ITimeShiftedGasHit new_hit = ITimeShiftedGasHit(*g4Gas, proton_time_shift);
            // Save the relevant trajectories from this hit if possible/wanted
            if(saveTrajs) ManageHitTracks(new_hit, proton_time_shift, trajectories);
            // Add the new hit
            out_container->push_back(new IG4HitGas(new_hit));
        } else if (const IG4HitSegment* g4Seg=dynamic_cast<IG4HitSegment*>(*i_hit)){
            // Create the time shifted hit segment
            ITimeShiftedSegHit new_hit = ITimeShiftedSegHit(*g4Seg, proton_time_shift);
            // Save the relevant trajectories from this hit if possible/wanted
            if(saveTrajs) ManageHitTracks(new_hit, proton_time_shift, trajectories);
            // Add the new hit
            out_container->push_back(new IG4HitSegment(new_hit));
        } else if (const IG4HitCheren* g4Ch=dynamic_cast<IG4HitCheren*>(*i_hit)){
            // Create the time shifted hit segment
            ITimeShiftedCherenHit new_hit = ITimeShiftedCherenHit(*g4Ch, proton_time_shift);
            // Save the relevant trajectories from this hit if possible/wanted
            if(saveTrajs) ManageHitTracks(new_hit, proton_time_shift, trajectories);
            // Add the new hit
            out_container->push_back(new IG4HitCheren(new_hit));
        } else COMETSevere("Hit type unknown!!");
    }
}

void ISimHitMergerMain::CopyOverDeltaRays(const COMET::IPAIDeltaRayContainer& deltaRayContainer,
					  double proton_time_shift){
					  //COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories){
    // Set the output container to null
    IPAIDeltaRayContainer* out_container=NULL;
    // Try to find the container by name
    
    DeltaRayContainerMap::iterator i_out_dc 
      = fDeltaRayContainers.find(deltaRayContainer.GetName());
    
    // If its not found, make a new one and insert it into the map
    
    if(i_out_dc == fDeltaRayContainers.end()){
        out_container= new IPAIDeltaRayContainer(deltaRayContainer.GetName());
        fDeltaRayContainers.insert(std::make_pair( deltaRayContainer.GetName(),out_container));
    }else {
        out_container=i_out_dc->second;
    }
    
    // Iterate through the input hit container and copy over the hits
    for(IPAIDeltaRayContainer::const_iterator i_dr=deltaRayContainer.begin();
                                         i_dr!=deltaRayContainer.end();
                                         ++i_dr){
        // Check if we are saving trajectories
        //bool saveTrajs = fSaveTrajectories && trajectories;
        // Check if we are within the timie limit
        double hTime = (*i_dr)->GetPosT() + proton_time_shift;
        // Skip this hit if it is not in the time limit
        if ((hTime > fMaximumTime) || (hTime < fMinimumTime)) continue;
        // Try casting to known g4hit types and adding the time shifted version
        // if it is found
        // TODO clean this up with polymophism/templating
	//const IG4HitMonitor* g4Mon=dynamic_cast<IG4HitMonitor*>(*i_hit);
	ITimeShiftedGasRegionDR new_dr = ITimeShiftedGasRegionDR(*(*i_dr), proton_time_shift);

	//if(saveTrajs) ManageDeltaRayTracks(new_dr, proton_time_shift, trajectories);
	out_container->push_back(new IPAIDeltaRay(new_dr));
    }   
}

void ISimHitMergerMain::ManageHitTracks(COMET::IVTimeShiftedHit& new_hit,
                                        double proton_time_shift,
                                        COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories){
    // Get the track IDs to save
    std::set<int> track_ids;
    new_hit.GetTrackIDs(track_ids);
    // Save the tracks
    ExtractTrajectoriesFor(track_ids,trajectories,proton_time_shift);
    // Set the new track IDs
    new_hit.SetTrackIDs(fNewIds);
}

void ISimHitMergerMain::ManageDeltaRayTracks(COMET::IVTimeShiftedDeltaRay& new_hit,
					     double proton_time_shift,
					     COMET::IHandle<COMET::IPAIDeltaRayContainer> trajectories){
    // Get the track IDs to save
    std::set<int> track_ids;
    new_hit.GetTrackIDs(track_ids);
    // Save the tracks
    ExtractTrajectoriesFor(track_ids,trajectories,proton_time_shift);
    // Set the new track IDs
    new_hit.SetTrackIDs(fNewIds);
}

void ISimHitMergerMain::AfterLastInputEvent(ICOMETOutput& outFile){
    // Write the last event if there are any proton-on-target events not written
    if (fMergedEvents > 0) {
        COMET::ICOMETEvent event;
        ReplaceEventContents(event);
        outFile.WriteEvent(event);
    }
    // Ensure geometry is written always
    outFile.WriteGeometry(IOADatabase::Get().UpdateGeometry());
    // Ensure the EM field is always written
    outFile.WriteEMField(IOADatabase::Get().EMFieldDescription());
    PrepareNewOutputEvent();
}

double ISimHitMergerMain::GetProtonTimeShift(const COMET::ICOMETEvent& event){
    double time_shift=0;
    if (fTimeShiftGenerator)
      time_shift= fTimeShiftGenerator->GetTimeShift(event);
    if(fNoBunches>1) {
        // Choose which bunch to put this proton into
        int bunch;
        // Random bunch
        if (fUseRandomBunch) bunch = gRandom->Integer(fNoBunches);
        // Fill the bunches in an orderly fashion
        else bunch = fMergedEvents % fNoBunches;
        time_shift += bunch*fBunchSeparation;
    }
    if(fGlobalShift !=0) time_shift+=fGlobalShift;
    return time_shift;
}

int ISimHitMergerMain::GetNumberOfProtons(const COMET::ICOMETEvent& event)const{
    // default number
    int n_proton = 1;

    // If this event contains a protons-on-target Datum use that:
    COMET::IHandle<COMET::IIntegerDatum> nProtonsDatum
        = event.Get<COMET::IIntegerDatum>("n_protons");
    if (nProtonsDatum)
      n_proton = nProtonsDatum->GetValue();

    return n_proton;
}

void ISimHitMergerMain::ReplaceEventContents(COMET::ICOMETEvent& eventInFile){
   
    // Specially merged events
    COMET::IDataVector* cherryPickedTracks = ProcessCherryPicks();
    
    // Clear existing data
    eventInFile.Clear();

    // Add the memory profiling
    IHandle<IDataVector> profile = eventInFile.Get<IDataVector>("SHM_profile");
    if (fProfileEvents){
        // If we don't have one add it
        if (!profile) {
            profile = IHandle<IDataVector>(new IDataVector("SHM_profile"));
            eventInFile.AddDatum(profile);
        }
    }

    // Prepare a new Context
    COMET::ICOMETContext context;
    context.SetRun(fRunNumber);
    context.SetEvent(fOutputEventNo);
    context.SetPartition(COMET::ICOMETContext::kMCData);
    eventInFile.SetContext(context);

    // Add all truth and hits etc
    IHandle<IDataVector> truth = eventInFile.Get<IDataVector>("truth");
    if (!truth) {
        truth=IHandle<IDataVector>(new IDataVector("truth"));
        eventInFile.AddDatum(truth);
    }

    // copy trajectories
    if(fTrajectories) truth->AddDatum(fTrajectories);

    // Get or make the g4Hits vector of hit containers
    IHandle<IDataVector> g4Hits = truth->Get<IDataVector>("g4Hits");
    if (!g4Hits) {
        g4Hits=IHandle<IDataVector>(new IDataVector("g4Hits"));
        truth->AddDatum(g4Hits);
    }
    // Copy over each hit container from this event
    for(HitContainerMap::const_iterator i_hc=fHitContainers.begin();
            i_hc!=fHitContainers.end(); ++i_hc){
        g4Hits->push_back(i_hc->second);
    }

    // Add all deltaRay container
    IHandle<IDataVector> paiPhys = truth->Get<IDataVector>("PAIPhysics");
    if (!paiPhys) {
        paiPhys=IHandle<IDataVector>(new IDataVector("PAIPhysics"));
        truth->AddDatum(paiPhys);
    }
    // Copy over each deltaRay container from this event
    for(DeltaRayContainerMap::const_iterator i_dc=fDeltaRayContainers.begin();
            i_dc!=fDeltaRayContainers.end(); ++i_dc){
        paiPhys->push_back(i_dc->second);
    }

    // Include the list of contributors that made this event
    eventInFile.push_back(fEventContributors);

    // Store the number of protons in this event
    eventInFile.push_back(new IIntegerDatum("n_events",  fMergedEvents));
    eventInFile.push_back(new IIntegerDatum("n_protons", fMergedPOTs));

    // Store the track Ids of specially merged events
    if(cherryPickedTracks)
      eventInFile.push_back( cherryPickedTracks );
    
    // Memory and time at the end of the persistence loop
    if (fProfileEvents) 
        profile->push_back(new IProfileData("end_event_persist"));
}

COMET::IDataVector* ISimHitMergerMain::ProcessCherryPicks()
{
  if(not fCherryPickMan.IsInitialized()) return NULL;

  COMETNamedVerbose("HitMerger", "Start to process specially merged events");

  TDirectory* pwd = TDirectory::CurrentDirectory(); // Push
  
  fCherryPickMan.PickupEvents();
  COMET::IDataVector* cherryPickedTracks =
    new COMET::IDataVector("CherryPicking",
			   "List of the cardinal track IDs of the cherry-picked events.");
  // Loop for all the special events
  COMET::ICOMETEvent* in_event;
  while(in_event = fCherryPickMan.GetNextEvent()){
    // Extract hits and (optionally the trajectories)from the current event and
    // time shift
    double time_shift = fCherryPickMan.IsTimeShiftEnabled()?
      fCherryPickMan.GetTimeShift(in_event) : GetProtonTimeShift(*in_event);
    
    bool event_contributed_hits = ExtractHits(*in_event, time_shift);
    
    // Not included in the P.O.T counter (should be included?)
    int n_proton = 1;
    
    // Log this input event as having contributed to the merged event
    if (event_contributed_hits) {
      fEventContributors->Add(fCherryPickMan.GetCurrentInput(),
			      *in_event,
			      n_proton,
			      time_shift,
			      &fNewIds);
      COMETNamedTrace("HitMerger","Taken hits for an special event "
		      << in_event->GetEventId() <<" in input file '"
		      << fCherryPickMan.GetCurrentInput()->GetFilename());

      // Get the new trackID of the cardinal tracks of this cherry-picked event
      if(0 == fCherryPickMan.FindCardinalTrackIds(*in_event)){
        COMETError("The cherry-picked event does not have any cardinal track!");
      }
      const std::set<Int_t>& cardinalTrackIds = fCherryPickMan.GetCardinalTrackIds();
      for(std::set<Int_t>::const_iterator 
	    it = cardinalTrackIds.begin(); it != cardinalTrackIds.end(); it++){
	COMET::IIntegerDatum* newTrackIdDatum =
	  new COMET::IIntegerDatum("CherryPickedTrackId", fNewIds[*it]);
	cherryPickedTracks->push_back( newTrackIdDatum );
      }
    }
  }
  pwd->cd(); // Pop
  
  return cherryPickedTracks;
}