#include <IShowerTruthInfo.hxx>
#include <ICOMETLog.hxx>
#include <ICOMETEvent.hxx>
#include <IEventFolder.hxx>
#include <IG4Trajectory.hxx>
#include <IDataVector.hxx>
#include <IHit.hxx>
#include <IComboHit.hxx>
#include <IReconHit.hxx>
#include <IMultiHit.hxx>
#include <IDataHit.hxx>
#include <IMCHit.hxx>
#include <IG4VHit.hxx>
#include <IG4HitSegment.hxx>
#include <IReconCluster.hxx>
#include <IReconShower.hxx>
#include <IReconTrack.hxx>
#include <IReconPID.hxx>
#include <IReconVertex.hxx>
#include <IGeometryId.hxx>
#include <IGeomIdManager.hxx>
#include <IOADatabase.hxx>
#include <TrackTruthInfo.hxx>
#include <retrieveHitTruthInfo.hxx>
#include <algorithm>

COMET::IShowerTruthInfo::IShowerTruthInfo() {
}

COMET::IShowerTruthInfo::IShowerTruthInfo(COMET::IHandle<COMET::IReconBase> object,
                                       bool enableSingle,
                                       bool enablePrimary,
                                       bool enableRecursive) {
  FillFromObject(object, enableSingle, enablePrimary, enableRecursive);
}

COMET::IShowerTruthInfo::IShowerTruthInfo(COMET::IHandle<COMET::IHitSelection> hits,
                                       COMET::IShowerTruthInfo::Views views,
                                       bool enableSingle,
                                       bool enablePrimary,
                                       bool enableRecursive) {
  FillFromHits(hits, views, enableSingle, enablePrimary, enableRecursive);
}

COMET::IShowerTruthInfo::~IShowerTruthInfo() {
}

std::vector<COMET::IShowerTruthInfo::TruthInfo>
    COMET::IShowerTruthInfo::GetSingleTrajectories(
                                       COMET::IShowerTruthInfo::SortOrder order) {
  if (!fEnableSingle) {
    COMETWarn("Single truth results requested, but calculation not enabled.");
  }

  return SortResult(fSingleTrajectoryResults, order);
}

std::vector<COMET::IShowerTruthInfo::TruthInfo>
    COMET::IShowerTruthInfo::GetPrimaryTrajectories(
                                       COMET::IShowerTruthInfo::SortOrder order) {
  if (!fEnablePrimary) {
    COMETWarn("Primary truth results requested, but calculation not enabled");
  }

  return SortResult(fPrimaryTrajectoryResults, order);
}

std::vector<COMET::IShowerTruthInfo::TruthInfo>
    COMET::IShowerTruthInfo::GetRecursiveTrajectories(
                                       COMET::IShowerTruthInfo::SortOrder order) {
  if (!fEnableRecursive) {
    COMETWarn("Recursive truth results requested, but calculation not enabled.");
  }

  return SortResult(fRecursiveTrajectoryResults, order);
}

std::vector<COMET::IShowerTruthInfo::TruthInfo> COMET::IShowerTruthInfo::SortResult(
                           std::vector<COMET::IShowerTruthInfo::TruthInfo> result,
                           COMET::IShowerTruthInfo::SortOrder order) {

  if (order == kNumHits) {
    std::sort(result.begin(), result.end(), ISortByNumHits());
  } else if (order == kG4Energy) {
    std::sort(result.begin(), result.end(), ISortByG4Energy());
  } else if (order == kParticleID) {
    std::sort(result.begin(), result.end(), ISortByParticleID());
  }

  return result;
}

void COMET::IShowerTruthInfo::FillFromObject(COMET::IHandle<COMET::IReconBase> object,
                                          bool enableSingle,
                                          bool enablePrimary,
                                          bool enableRecursive) {
  // Clear the details of saved results
  Reset();

  if (!object) {
    COMETVerbose("Input object does not exist");
    return;
  }

  // Work out whether the object is 2D or 3D.
  Views views = GetViews(object);
  COMET::IHandle<COMET::IHitSelection> hits = object->GetHits();

  FillFromHits(hits, views, enableSingle, enablePrimary, enableRecursive);
}

void COMET::IShowerTruthInfo::FillFromHits(COMET::IHandle<COMET::IHitSelection> hits,
                                        COMET::IShowerTruthInfo::Views views,
                                        bool enableSingle,
                                        bool enablePrimary,
                                        bool enableRecursive) {
  // Clear the details of saved results
  Reset();

  if (!hits) {
    COMETVerbose("Input hit selection does not exist");
    return;
  }

  // Make sure we actually have some trajectories to search for.
  const COMET::ICOMETEvent& event = *(COMET::IEventFolder::GetCurrentEvent());
  COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories =
    event.Get<COMET::IG4TrajectoryContainer> ("truth/G4Trajectories");

  if (!trajectories) {
    COMETVerbose("No truth trajectories exist");
    return;
  }

  // Convert the user's input to fields so we don't have to pass them around.
  fHits = hits;
  fViews = views;
  fEnableSingle = enableSingle;
  fEnablePrimary = enablePrimary;
  fEnableRecursive = enableRecursive;

  // Work out the true information for the input hit selection. This works out
  // which hits in the hit selection each true trajectory caused.
  FillSelectionTruthDetails();

  // Work out the true information for all the hits in the event. We need to
  // do this to work out the completeness of the truth-matching.
  FillEventTruthDetails();

  if (fEnableRecursive) {
    // Work out the parent-child relationships for the trajectories that are
    // relevant to the recursive matching method.
    MarkRecursiveTrajectories();
  }

  // Calculate the completeness and cleanliness for the algorithms that were
  // enabled.
  CalculateStatistics();
}

void COMET::IShowerTruthInfo::Reset() {
  // Reset all the fields to default values.
  fEnableSingle = true;
  fEnablePrimary = true;
  fEnableRecursive = true;
  fDetectorsUsed = 0;
  fHits = COMET::IHandle<COMET::IHitSelection>(NULL);
  fViews = kViewNotSet;
  fSingleTrajectoryResults.clear();
  fPrimaryTrajectoryResults.clear();
  fRecursiveTrajectoryResults.clear();
  fSingleTrajList.clear();
  fPrimaryTrajList.clear();
  fRecursiveTrajList.clear();
  fPrimaryTrajHierarchy.clear();
  fRecursiveTrajHierarchy.clear();
  fEventDetails.clear();
  fSelecDetails.clear();

  // This needs to be tuned.
  fDarkNoiseEnergy = 0;
}

void COMET::IShowerTruthInfo::MarkSingleTrajectory(int trajID) {
  if (std::find(fSingleTrajList.begin(), fSingleTrajList.end(), trajID) ==
        fSingleTrajList.end()) {
    // Add the trajectory to the list of ones that directly contribute to the
    // hits in the input hit selection.
    fSingleTrajList.push_back(trajID);
  }
}

void COMET::IShowerTruthInfo::MarkPrimaryTrajectory(int trajID, int primaryID) {
  if (std::find(fPrimaryTrajList.begin(), fPrimaryTrajList.end(), primaryID) ==
        fPrimaryTrajList.end()) {
    // Add the primary to the list of primaries we should calculate results
    // for.
    fPrimaryTrajList.push_back(primaryID);
  }

  // Add the primary->daughter connection.
  if (fPrimaryTrajHierarchy.find(primaryID) == fPrimaryTrajHierarchy.end()) {
    fPrimaryTrajHierarchy[primaryID] = std::vector<int>();
  }

  if (std::find(fPrimaryTrajHierarchy[primaryID].begin(),
                fPrimaryTrajHierarchy[primaryID].end(),
                trajID) == fPrimaryTrajHierarchy[primaryID].end()) {
    fPrimaryTrajHierarchy[primaryID].push_back(trajID);
  }
}

void COMET::IShowerTruthInfo::MarkRecursiveTrajectories() {
  const COMET::ICOMETEvent& event = *(COMET::IEventFolder::GetCurrentEvent());
  COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories =
    event.Get<COMET::IG4TrajectoryContainer> ("truth/G4Trajectories");

  // Temporary map of whether a trajecotry in the event enters the subdetectors
  // used by the input object. The one exception is for neutrinos, which are
  // never marked as relevant.
  std::map<int, bool> singleTrajRelevant;

  // Temporary map of the parent->child relationships for all the trajectories
  // in the event.
  ParticleTree parentChildMap;

  // For each trajectory in the event, work out the parent->child relationship
  // and see whether it enters a subdetector we use.
  for (COMET::IG4TrajectoryContainer::iterator it = trajectories->begin();
       it != trajectories->end();
       it++) {
    // First fill the parent->child map
    int trajID = it->first;
    COMET::IG4Trajectory traj = it->second;
    int parentID = traj.GetParentId();

    if (parentChildMap.find(trajID) == parentChildMap.end()) {
      parentChildMap[trajID] = std::vector<int>();
    }

    if (parentChildMap.find(parentID) == parentChildMap.end()) {
      parentChildMap[parentID] = std::vector<int>();
    }

    parentChildMap[parentID].push_back(trajID);

    // Now see if this individual trajectory is relevant
    singleTrajRelevant[trajID] = false;

    // Trajectory 0 is never relevant
    if (trajID == 0) {
      continue;
    }

    // Never class a neutrino as relevant
    int pdg = traj.GetPDGEncoding();

    if (abs(pdg) == 12 || abs(pdg) == 14 || abs(pdg) == 16) {
      continue;
    }

    // Check if the particle went through one of the detectors our hit
    // selection uses.
    COMET::IG4Trajectory::Points points = traj.GetTrajectoryPoints();
    COMET::IGeomIdManager& gidManager = COMET::IOADatabase::Get().GeomId();

    for (COMET::IG4Trajectory::Points::iterator it = points.begin();
         it != points.end();
         it++) {
      COMET::IG4TrajectoryPoint point = (*it);
      TLorentzVector pos = point.GetPosition();
      COMET::IGeometryId geomID;

      if (gidManager.GetGeometryId(pos.X(), pos.Y(), pos.Z(), geomID)) {
        int thisDet = TrackTruthInfo::GeomIdToDetectorId(geomID);

        if (thisDet & fDetectorsUsed) {
          singleTrajRelevant[trajID] = true;
          break;
        }
      }
    }
  }

  // Now we know whether each individual trajectory enters our subdetector, we
  // can work out whether a trajectory is relevant for the recursive matching
  // or not. The trajectory is obviously relevant if it enters a subdetector
  // we use. Intermediate trajectories that have parents and children that
  // enter the subdetector also need to be marked as relevant.
  for (COMET::IG4TrajectoryContainer::iterator it = trajectories->begin();
       it != trajectories->end();
       it++) {

    int trajID = it->first;
    COMET::IG4Trajectory traj = it->second;
    int parentID = traj.GetParentId();

    bool thisTrajRel = singleTrajRelevant[trajID];
    bool parentsIrrel =
      AllParentsIrrelevant(trajectories->GetTrajectory(trajID),
                           singleTrajRelevant);
    bool childrenIrrel =
      AllChildrenIrrelevant(trajectories->GetTrajectory(trajID),
                            singleTrajRelevant,
                            parentChildMap);

    if (thisTrajRel || (!parentsIrrel && !childrenIrrel)) {
      // This trajectory is relevant, so add it to the parent-child map of
      // relevant trajectories.
      if (fRecursiveTrajHierarchy.find(trajID) ==
            fRecursiveTrajHierarchy.end()) {
        fRecursiveTrajHierarchy[trajID] = std::vector<int>();
      }

      if (fRecursiveTrajHierarchy.find(parentID) ==
            fRecursiveTrajHierarchy.end()) {
        fRecursiveTrajHierarchy[parentID] = std::vector<int>();
      }

      fRecursiveTrajHierarchy[parentID].push_back(trajID);
    }

    if (thisTrajRel && parentsIrrel) {
      // This trajectory enters our subdetector, but none of its parents do -
      // we've found a top-level particle that we'll return truth-matching
      // results for.
      fRecursiveTrajList.push_back(trajID);
    }
  }
}

bool COMET::IShowerTruthInfo::AllParentsIrrelevant(
                                     COMET::IHandle<COMET::IG4Trajectory> traj,
                                     std::map<int, bool> &singleTrajRelevant) {

  const COMET::ICOMETEvent& event = *(COMET::IEventFolder::GetCurrentEvent());
  COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories =
    event.Get<COMET::IG4TrajectoryContainer> ("truth/G4Trajectories");

  int parentID = traj->GetParentId();
  COMET::IHandle<COMET::IG4Trajectory> parent =
    trajectories->GetTrajectory(parentID);

  if (parentID == 0 || !parent) {
    // We've gone as high as we can and only found irrelevant trajectories
    return true;
  } else if (singleTrajRelevant[parentID]) {
    // The parent is relevant
    return false;
  } else if (parentID == traj->GetTrackId()) {
    COMETWarn("Infinite loop detected in parent map: " <<
                parentID << "<-->" << traj->GetTrackId());
    return true;
  } else {
    // Recurse up and see if the grandparent is relevant
    return AllParentsIrrelevant(parent, singleTrajRelevant);
  }
}

bool COMET::IShowerTruthInfo::AllChildrenIrrelevant(
                                       COMET::IHandle<COMET::IG4Trajectory> traj,
                                       std::map<int, bool> &singleTrajRelevant,
                                       ParticleTree &parentChildMap) {
  if (!traj) {
    return true;
  }

  int trajID = traj->GetTrackId();

  if (singleTrajRelevant[trajID]) {
    // This trajectory is relevant - break out
    return false;
  }

  const COMET::ICOMETEvent& event = *(COMET::IEventFolder::GetCurrentEvent());
  COMET::IHandle<COMET::IG4TrajectoryContainer> trajectories =
    event.Get<COMET::IG4TrajectoryContainer> ("truth/G4Trajectories");
  std::vector<int> childIDs = parentChildMap[trajID];

  for (ParticleList::iterator it = childIDs.begin();
       it != childIDs.end();
       it++) {
    // Recurse down each of the children, seeing if any of them are relevant
    COMET::IHandle<COMET::IG4Trajectory> child = trajectories->GetTrajectory(*it);

    if (child && child->GetTrackId() == trajID) {
      COMETWarn("Infinite loop detected in finding children: " <<
                  child->GetTrackId() << "<-->" << trajID);
    } else if (!AllChildrenIrrelevant(child, singleTrajRelevant, parentChildMap)) {
        return false;
    }
  }

  return true;
}

void COMET::IShowerTruthInfo::FillSelectionTruthDetails() {
  FillTruthDetails(&(*fHits), fSelecDetails, false);
}

void COMET::IShowerTruthInfo::FillEventTruthDetails() {
  const COMET::ICOMETEvent& event = *(COMET::IEventFolder::GetCurrentEvent());
  COMET::IHandle<COMET::IDataVector> hitV = event.Get<COMET::IDataVector> ("hits");

  for (COMET::IDataVector::iterator hiter = hitV->begin();
       hiter != hitV->end();
       ++hiter) {
    COMET::IHitSelection* hitList = dynamic_cast<COMET::IHitSelection*> (*hiter);
    FillTruthDetails(hitList, fEventDetails, true);
  }
}

void COMET::IShowerTruthInfo::FillTruthDetails(COMET::IHitSelection* hits,
                                COMET::IShowerTruthInfo::ParticleDetails &details,
                                bool eventLevel) {
  for (COMET::IHitSelection::const_iterator hitIt = hits->begin();
       hitIt != hits->end();
       hitIt++) {

    COMET::IHandle<COMET::IHit> hit(*hitIt);
    int thisDet = TrackTruthInfo::HitToDetectorId(hit);

    if (!eventLevel) {
      // We're looking at a hit from the input hit selection - make a note of
      // which subdetectors are being used.
      fDetectorsUsed |= thisDet;
    } else if (!(fDetectorsUsed & thisDet)) {
      // We're looking at a hit from the event's hit selections. If it's not in
      // one of the subdetectors we're using we can skip analysing it.
      continue;
    }

    // If the hit is in the right orientation, extract the information from it.
    if ((fViews == kXY && (hit->IsXHit() && hit->IsYHit() && !hit->IsZHit()))
      || (fViews == kYZ && (!hit->IsXHit() && hit->IsYHit() && hit->IsZHit()))
      || (fViews == kXZ && (hit->IsXHit() && !hit->IsYHit() && hit->IsZHit()))
      || (fViews == kXYZ)) {

      COMET::IHandle<COMET::IComboHit> combohit(*hitIt);
      COMET::IHandle<COMET::IReconHit> recohit(*hitIt);

      if (combohit) {
        COMET::IHitSelection::const_iterator comboIt;
        for (comboIt = combohit->GetHits().begin();
             comboIt != combohit->GetHits().end();
             comboIt++) {
          COMET::IHandle<COMET::IHit> contHit(*comboIt);
          FillTruthDetails(contHit, details, eventLevel);
        }
      } else if (recohit) {
        for (int i = 0; i < recohit->GetContributorCount(); i++) {
          COMET::IHandle<COMET::IHit> contHit(recohit->GetContributor(i));
          FillTruthDetails(contHit, details, eventLevel);
        }
      } else {
        FillTruthDetails(hit, details, eventLevel);
      }
    }
  }
}

void COMET::IShowerTruthInfo::FillTruthDetails(COMET::IHandle<COMET::IHit> &hit,
                                COMET::IShowerTruthInfo::ParticleDetails &details,
                                bool eventLevel) {
  UInt_t channel = hit->GetChannelId(0).AsUInt();
  std::vector<COMET::IG4VHit*> g4Hits = HitTruthInfo::GetHitTruthInfo(hit);

  for (std::vector<COMET::IG4VHit*>::const_iterator g4HitIt = g4Hits.begin();
       g4HitIt != g4Hits.end();
       g4HitIt++) {
    COMET::IG4HitSegment* g4HitCast = dynamic_cast<COMET::IG4HitSegment*> (*g4HitIt);
    int primaryID = g4HitCast->GetPrimaryId();

    // A IG4HitSegment can have contributions from multiple particles. We don't
    // know how much each particle contributed, so just divide it equally. It's
    // not perfect, but a reasonable approximation, especially as this case
    // doesn't happen very often.
    const std::vector<int>& contributors = g4HitCast->GetContributors();
    double energyFraction = g4HitCast->GetEnergyDeposit() / contributors.size();

    for (std::vector<int>::const_iterator trajIDIter = contributors.begin();
         trajIDIter != contributors.end();
         trajIDIter++) {
      int trajID = (*trajIDIter);
      details[trajID][channel] += energyFraction;

      if (fEnableSingle && !eventLevel) {
        // Note that this trajecotry contributes to the input hit selection.
        MarkSingleTrajectory(trajID);
      }

      if (fEnablePrimary) {
        if (!eventLevel) {
          // Always add the primary->child mapping for children that contribute
          // to the input hit selection.
          MarkPrimaryTrajectory(trajID, primaryID);
        } else if (fPrimaryTrajHierarchy.find(primaryID) !=
                     fPrimaryTrajHierarchy.end()) {
          // Only add the primary->child mappings found in event-level hit
          // selections if we're already interested in the primary.
          MarkPrimaryTrajectory(trajID, primaryID);
        }
      }
    }
  }
}

double COMET::IShowerTruthInfo::GetG4EnergyOfHit(COMET::IHandle<COMET::IHit> hit) {
  double energy = 0;
  std::vector<COMET::IG4VHit*> g4Hits = HitTruthInfo::GetHitTruthInfo(hit);

  for (std::vector<COMET::IG4VHit*>::const_iterator g4HitIt = g4Hits.begin();
       g4HitIt != g4Hits.end();
       g4HitIt++) {
    COMET::IG4HitSegment* g4HitCast = dynamic_cast<COMET::IG4HitSegment*> (*g4HitIt);
    energy += g4HitCast->GetEnergyDeposit();
  }

  // Pure noise hits don't have any G4 energy associated with them. Approximate
  // a level instead.
  if (g4Hits.size() == 0) {
    energy = fDarkNoiseEnergy;
  }

  return energy;
}

void COMET::IShowerTruthInfo::GetReconStats(int &totalRecoHits,
                                         double &totalRecoEnergy) {
  for (COMET::IHitSelection::iterator it = fHits->begin();
       it != fHits->end();
       it++) {
    COMET::IHandle<COMET::IHit> hit = (*it);

    // If the hit is in the right orientation, increase the number of hits and
    // amount of energy the reconstructed object contains.
    if ((fViews == kXY && (hit->IsXHit() && hit->IsYHit() && !hit->IsZHit()))
      || (fViews == kYZ && (!hit->IsXHit() && hit->IsYHit() && hit->IsZHit()))
      || (fViews == kXZ && (hit->IsXHit() && !hit->IsYHit() && hit->IsZHit()))
      || (fViews == kXYZ)) {

      COMET::IHandle<COMET::IComboHit> combohit = hit;
      COMET::IHandle<COMET::IReconHit> recohit = hit;

      if (combohit) {
        COMET::IHitSelection::const_iterator comboIt;
        for (comboIt = combohit->GetHits().begin();
             comboIt != combohit->GetHits().end();
             comboIt++) {
          totalRecoHits++;
          totalRecoEnergy += GetG4EnergyOfHit(*comboIt);
        }
      } else if (recohit) {
        for (int i = 0; i < recohit->GetContributorCount(); i++) {
          totalRecoHits++;
          totalRecoEnergy += GetG4EnergyOfHit(recohit->GetContributor(i));
        }
      } else {
        totalRecoHits++;
        totalRecoEnergy += GetG4EnergyOfHit(hit);
      }
    }
  }
}

void COMET::IShowerTruthInfo::CalculateStatistics() {
  int totalRecoHits = 0;
  double totalRecoEnergy = 0.0;

  // Get the number of hits and amount of G4 energy in the reconstructed object
  // so we can later compute the cleanliness.
  GetReconStats(totalRecoHits, totalRecoEnergy);

  if (fEnableSingle) {
    fSingleTrajectoryResults = CalculateStatistics(totalRecoHits,
                                                   totalRecoEnergy,
                                                   kSingle);
  }

  if (fEnablePrimary) {
    fPrimaryTrajectoryResults = CalculateStatistics(totalRecoHits,
                                                    totalRecoEnergy,
                                                    kPrimary);
  }

  if (fEnableRecursive) {
    fRecursiveTrajectoryResults = CalculateStatistics(totalRecoHits,
                                                      totalRecoEnergy,
                                                      kRecursive);
  }
}

std::vector<COMET::IShowerTruthInfo::TruthInfo>
    COMET::IShowerTruthInfo::CalculateStatistics(double totalRecoHits,
                                        double totalRecoEnergy,
                                        COMET::IShowerTruthInfo::Algorithm type) {
  std::vector<TruthInfo> results;

  // Get the right list of particles to calculate statistics for.
  ParticleList list;

  if (type == kSingle) {
    list = fSingleTrajList;
  } else if (type == kPrimary) {
    list = fPrimaryTrajList;
  } else if (type == kRecursive) {
    list = fRecursiveTrajList;
  }

  // Work out the statistics for each particle we're returning a result for.
  for (ParticleList::iterator it = list.begin(); it != list.end(); it++) {
    TruthInfo info;
    info.NumHits = 0;
    info.G4Energy = 0.0;

    int id = (*it);
    int idealNumHits = 0;
    double idealEnergy = 0.0;
    info.ParticleID = id;

    if (type == kSingle) {
      // Number of hits/energy in the input object from this traj.
      GetSingleNumHitsEnergy(id, fSelecDetails, info.NumHits, info.G4Energy);

      // Number of hits/energy from this traj in all the hits in the
      // sudetectors used by the input object.
      GetSingleNumHitsEnergy(id, fEventDetails, idealNumHits, idealEnergy);
    } else if (type == kPrimary) {
      // Number of hits/energy in the input object due to this primary.
      GetPrimaryNumHitsEnergy(id, fSelecDetails, info.NumHits, info.G4Energy);

      // Number of hits/energy due to this primary in all the hits in the
      // sudetectors used by the input object.
      GetPrimaryNumHitsEnergy(id, fEventDetails, idealNumHits, idealEnergy);
    } else if (type == kRecursive) {
      // Number of hits/energy in the input object due to this top-level traj.
      GetRecursiveNumHitsEnergy(id, fSelecDetails, info.NumHits, info.G4Energy);

      // Number of hits/energy due to this top-level traj in all the hits in
      // the sudetectors used by the input object.
      GetRecursiveNumHitsEnergy(id, fEventDetails, idealNumHits, idealEnergy);
    }

    // Finally do the completeness and cleanliness calculations.
    info.Completeness_NumHits = (double) info.NumHits / (double) idealNumHits;
    info.Completeness_G4Energy = info.G4Energy / idealEnergy;

    info.Cleanliness_NumHits = (double) info.NumHits / (double) totalRecoHits;
    info.Cleanliness_G4Energy = info.G4Energy / totalRecoEnergy;

    results.push_back(info);
  }

  return results;
}

void COMET::IShowerTruthInfo::GetPrimaryNumHitsEnergy(int trajID,
                                COMET::IShowerTruthInfo::ParticleDetails &details,
                                int& numHits,
                                double& energy) {
  std::vector<UInt_t> usedChannels;
  std::vector<int> usedPrimaries;
  GetPrimaryNumHitsEnergy(
                trajID, details, numHits, energy, usedChannels, usedPrimaries);
}

void COMET::IShowerTruthInfo::GetPrimaryNumHitsEnergy(int trajID,
                                COMET::IShowerTruthInfo::ParticleDetails &details,
                                int& numHits,
                                double& energy,
                                std::vector<UInt_t>& usedChannels,
                                std::vector<int>& usedPrimaries) {
  usedPrimaries.push_back(trajID);

  // Loop over all the particles that have this particle as their primary
  for (ParticleList::iterator trajIt = fPrimaryTrajHierarchy[trajID].begin();
       trajIt != fPrimaryTrajHierarchy[trajID].end();
       trajIt++) {

    for (HitDetails::iterator hitIt = details[*trajIt].begin();
         hitIt != details[*trajIt].end();
         hitIt++) {

      UInt_t channel = hitIt->first;
      double channelEnergy = hitIt->second;

      // Make sure we haven't already used this channel for a previous particle
      if (std::find(usedChannels.begin(), usedChannels.end(), channel)
            == usedChannels.end()) {
        numHits++;
        usedChannels.push_back(channel);
      }

      // Energy was equally distributed amongst particles - always add share
      energy += channelEnergy;
    }

    // If this particle is also a primary, recurse downwards
    if (fPrimaryTrajHierarchy.find(*trajIt) != fPrimaryTrajHierarchy.end()
        && trajID != (*trajIt)) {
      // Make sure we aren't going to enter an infinite loop by recursing
      // downwards. There is an edge case where a primary contributes to a
      // hit segment, but the hit segment itself is associated to a different
      // primary which is at a LOWER level in the hierarchy.
      if (std::find(usedPrimaries.begin(), usedPrimaries.end(), *trajIt) !=
            usedPrimaries.end()) {
        COMETWarn("Infinite loop detected in primary trajectory map: " <<
                    *trajIt << "<-->" << trajID);
      } else {
        GetPrimaryNumHitsEnergy(
               *trajIt, details, numHits, energy, usedChannels, usedPrimaries);
      }
    }
  }
}

void COMET::IShowerTruthInfo::GetRecursiveNumHitsEnergy(int trajID,
                                COMET::IShowerTruthInfo::ParticleDetails &details,
                                int& numHits,
                                double& energy) {
  std::vector<UInt_t> usedChannels;
  GetRecursiveNumHitsEnergy(trajID, details, numHits, energy, usedChannels);
}

void COMET::IShowerTruthInfo::GetRecursiveNumHitsEnergy(int trajID,
                                COMET::IShowerTruthInfo::ParticleDetails &details,
                                int& numHits,
                                double& energy,
                                std::vector<UInt_t>& usedChannels) {
  for (HitDetails::iterator it = details[trajID].begin();
       it  != details[trajID].end();
       it++) {

    UInt_t channel = it->first;
    double channelEnergy = it->second;

    // See whether we already used this channel for a previous particle
    if (std::find(usedChannels.begin(), usedChannels.end(), channel)
          == usedChannels.end()) {
      numHits++;
      usedChannels.push_back(channel);
    }

    // Energy was equally distributed amongst particles - always add the share
    energy += channelEnergy;
  }

  // Recurse down to the daughters of this particle
  for (ParticleList::iterator it = fRecursiveTrajHierarchy[trajID].begin();
       it != fRecursiveTrajHierarchy[trajID].end();
       it++) {
    if (*it == trajID) {
      COMETWarn("Infinite loop detected in recursive trajectory hierarchy: " <<
                  *it << "<-->" << trajID);
    } else {
      GetRecursiveNumHitsEnergy(*it, details, numHits, energy, usedChannels);
    }
  }
}

void COMET::IShowerTruthInfo::GetSingleNumHitsEnergy(int trajID,
                                COMET::IShowerTruthInfo::ParticleDetails &details,
                                int& numHits,
                                double& energy) {
  for (HitDetails::iterator it = details[trajID].begin();
       it != details[trajID].end();
       it++) {
    numHits++;
    energy += it->second;
  }
}

COMET::IShowerTruthInfo::Views
    COMET::IShowerTruthInfo::GetViews(COMET::IHandle<COMET::IReconBase> object) {
  bool isX = false;
  bool isY = false;
  bool isZ = false;

  COMET::IHandle<COMET::IReconCluster> cl = object;
  COMET::IHandle<COMET::IReconTrack> tr = object;
  COMET::IHandle<COMET::IReconShower> sh = object;
  COMET::IHandle<COMET::IReconPID> pid = object;
  COMET::IHandle<COMET::IReconVertex> ve = object;

  if (cl) {
    isX = cl->IsXCluster();
    isY = cl->IsYCluster();
    isZ = cl->IsZCluster();
  } else if (tr) {
    isX = tr->IsXTrack();
    isY = tr->IsYTrack();
    isZ = tr->IsZTrack();
  } else if (sh) {
    isX = sh->IsXShower();
    isY = sh->IsYShower();
    isZ = sh->IsZShower();
  } else if (pid) {
    isX = pid->IsXPID();
    isY = pid->IsYPID();
    isZ = pid->IsZPID();
  } else if (ve) {
    isX = ve->IsXVertex();
    isY = ve->IsYVertex();
    isZ = ve->IsZVertex();
  }

  Views views = kViewNotSet;

  if (isX && isY && isZ) {
    views = kXYZ;
  } else if (isX && isY) {
    views = kXY;
  } else if (isX && isZ) {
    views = kXZ;
  } else if (isY && isZ) {
    views = kYZ;
  }

  return views;
}