#include <functional>
#include <algorithm>
#include <cassert>
#include <string>

#include <TString.h>
#include <TVector3.h>

#include <IAlgorithmResult.hxx>
#include <IDatum.hxx>
#include <IG4PrimaryVertex.hxx>
#include <IGeomInfo.hxx>
#include <IGeometryId.hxx>
#include <IHandle.hxx>
#include <IHitSelection.hxx>
#include <ICOMETEvent.hxx>
#include <ICOMETLog.hxx>
#include <IPIDState.hxx>
#include <IP0DGeom.hxx>
#include <IRealDatum.hxx>
#include <IIntegerDatum.hxx>
#include <IReconCluster.hxx>
#include <IReconNode.hxx>
#include <IReconPID.hxx>
#include <IReconShower.hxx>
#include <IReconTrack.hxx>
#include <IReconVertex.hxx>
#include <IReconState.hxx>
#include <IShowerState.hxx>
#include <ITrackState.hxx>
#include <IVertexState.hxx>
#include <IG4HitSegment.hxx>
#include <ITFBChannelId.hxx>
#include <retrieveHitTruthInfo.hxx>

#include "IP0DReconModule.hxx"

using namespace COMET;

ClassImp(COMET::IP0DReconModule);

ClassImp(COMET::IP0DReconModule::IP0DAlgoRes);
COMET::IP0DReconModule::IP0DAlgoRes::~IP0DAlgoRes(){}

ClassImp(COMET::IP0DReconModule::IP0DVertex);
COMET::IP0DReconModule::IP0DVertex::~IP0DVertex(){}

ClassImp(COMET::IP0DReconModule::IP0DParticle);
COMET::IP0DReconModule::IP0DParticle::~IP0DParticle() {}

ClassImp(COMET::IP0DReconModule::IP0DTrack);
COMET::IP0DReconModule::IP0DTrack::~IP0DTrack(){}

ClassImp(COMET::IP0DReconModule::IP0DShower);
COMET::IP0DReconModule::IP0DShower::~IP0DShower(){}

ClassImp(COMET::IP0DReconModule::IP0DNode);
COMET::IP0DReconModule::IP0DNode::~IP0DNode(){}

ClassImp(COMET::IP0DReconModule::IP0DCluster);
COMET::IP0DReconModule::IP0DCluster::~IP0DCluster() {}

ClassImp(COMET::IP0DReconModule::IP0DHit);

#define CVSTAG "\
  $Name: v5r19 $"
#define CVSID  "\
  $Id: IP0DReconModule.cxx,v 1.46 2012/08/18 01:17:05 gilje Exp $"

COMET::IP0DReconModule::IP0DReconModule(const char *name,
                                     const char *title) {
  SetNameTitle(name, title);

  // Enable this module by default:
  fIsEnabled = kTRUE;
  fDescription = "P0D Recon Output";
  fCVSTagName = CVSTAG;
  fCVSID = CVSID;

  // Tree variables
  fNAlgoResults = 0;
  fNVertices = 0;
  fNParticles = 0;
  fNShowers = 0;
  //  Commented out this line to remove extraneous oaAnalysis output
  //  std::cout << "fNShowers:\tSta\t" << fNShowers << std::endl;
  fNTracks = 0;
  fNNodes = 0;
  fNClusters = 0;
  fNHits = 0;

  fAlgoResults = new TClonesArray("COMET::IP0DReconModule::IP0DAlgoRes", 200);
  fVertices    = new TClonesArray("COMET::IP0DReconModule::IP0DVertex", 200);
  fParticles   = new TClonesArray("COMET::IP0DReconModule::IP0DParticle", 200);
  fShowers     = new TClonesArray("COMET::IP0DReconModule::IP0DShower", 200);
  fTracks      = new TClonesArray("COMET::IP0DReconModule::IP0DTrack", 200);
  fNodes       = new TClonesArray("COMET::IP0DReconModule::IP0DNode", 200);
  fHits        = new TClonesArray("COMET::IP0DReconModule::IP0DHit", 200);
  fClusters    = new TClonesArray("COMET::IP0DReconModule::IP0DCluster", 200);

  fRejectAlgoResultList.clear();
  fRejectAlgoResultList.push_back(TPRegexp("IP0DTrack$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0DExpandSeeds$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0D3DTrackMatching$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0DPairwiseVertex$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0D3DShowerMatching$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0DShareCharge$"));
  fRejectAlgoResultList.push_back(TPRegexp("-Expand$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0DPairwiseVertexSingle$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0DTrackVertex$"));
  fRejectAlgoResultList.push_back(TPRegexp("TP0DTrackPID$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0DShowerInputs$"));
  fRejectAlgoResultList.push_back(TPRegexp("IP0DShowerVertex$"));
  fRejectAlgoResultList.push_back(TPRegexp("TP0DShowerPID$"));
}


COMET::IP0DReconModule::~IP0DReconModule() {}

Bool_t COMET::IP0DReconModule::ProcessFirstEvent(COMET::ICOMETEvent&) {
  return true;
}


void COMET::IP0DReconModule::InitializeModule() {}

void COMET::IP0DReconModule::InitializeBranches() {

  // The number of objects that have been saved.
  fOutputTree->Branch("NAlgoResults", &fNAlgoResults, "NAlgoResults/S", fBufferSize);
  fOutputTree->Branch("NVertices",    &fNVertices,    "NVertices/S",    fBufferSize);
  fOutputTree->Branch("NParticles",   &fNParticles,   "NParticles/S",   fBufferSize);
  fOutputTree->Branch("NShowers",     &fNShowers,     "NShowers/S",     fBufferSize);
  fOutputTree->Branch("NTracks",      &fNTracks,      "NTracks/S",      fBufferSize);
  fOutputTree->Branch("NNodes",       &fNNodes,       "NNodes/S",       fBufferSize);
  fOutputTree->Branch("NHits",        &fNHits,        "NHits/S",        fBufferSize);
  fOutputTree->Branch("NClusters",    &fNClusters,    "NClusters/S",    fBufferSize);

  // A branch for the clone array of P0D object information.
  fOutputTree->Branch("AlgoResults", &fAlgoResults, fBufferSize, fSplitLevel);
  fOutputTree->Branch("Vertices",    &fVertices,    fBufferSize, fSplitLevel);
  fOutputTree->Branch("Particles",   &fParticles,   fBufferSize, fSplitLevel);
  fOutputTree->Branch("Showers",     &fShowers,     fBufferSize, fSplitLevel);
  fOutputTree->Branch("Tracks",      &fTracks,      fBufferSize, fSplitLevel);
  fOutputTree->Branch("Nodes",       &fNodes,       fBufferSize, fSplitLevel);
  fOutputTree->Branch("Hits",        &fHits,        fBufferSize, fSplitLevel);
  fOutputTree->Branch("Clusters",    &fClusters,    fBufferSize, fSplitLevel);
}


bool COMET::IP0DReconModule::FillTree(COMET::ICOMETEvent& event) {
  // Clear variables
  fNAlgoResults = 0;
  fNVertices    = 0;
  fNParticles   = 0;
  fNShowers     = 0;
  fNTracks      = 0;
  fNNodes       = 0;
  fNClusters    = 0;
  fNHits        = 0;

  fAlgoResults->Clear();
  fVertices->Clear();
  fParticles->Clear();
  fShowers->Clear();
  fTracks->Clear();
  fNodes->Clear();
  fClusters->Clear();
  fHits->Clear();

  fTempHitMap.clear();

  //###############################################################

  // Get the main result to save.
  COMET::IHandle<COMET::IAlgorithmResult> ReconP0D = event.GetFit("ReconP0D");

  if (!ReconP0D) {
    COMETInfo("No ReconP0D in Event");
    return true;
  }

  FillAlgorithmResult(ReconP0D, -1);

  return true;
}

template <class T>
void IP0DReconModule::FillBaseObject(T basePtr,
                                     COMET::IHandle<COMET::IReconBase> baseObject,
                                     bool saveHits){
  // Algorithm Name
  basePtr->AlgorithmName = baseObject->GetAlgorithmName();
  
  // Hit Based Information
  COMET::IHandle<COMET::IHitSelection> baseHits = baseObject->GetHits();
  
  if(baseHits){
    basePtr->Cycle = IP0DReconModule::GetCycle(baseHits);
    basePtr->NHits = (short)baseHits->size();
    
    if(saveHits){
      COMET::IHitSelection::const_iterator hit;
      for(hit = baseHits->begin();
          hit != baseHits->end(); ++hit){
        short hitId = FillHit(*hit);
        if(hitId >= 0)
          basePtr->Hits.push_back(hitId);
      }
    }
  } else {
    basePtr->Cycle = -1;
    basePtr->NHits = -1;
  }
  
  // Unique ID
  basePtr->UniqueID      = baseObject->GetUniqueID();
  
  // Daughter objects
  COMET::IHandle<COMET::IReconObjectContainer> objCont = baseObject->GetConstituents();
  if(objCont) {
    COMET::IReconObjectContainer::const_iterator constituent;
    for (constituent = objCont->begin();
         constituent != objCont->end();
         constituent++){
      ObjectID conID = FillReconObject(*constituent, saveHits);
      if(conID.type == ObjectID::kVertex){
        basePtr->Vertices.push_back(conID.id);}
      else if(conID.type == ObjectID::kParticle){
        basePtr->Particles.push_back(conID.id);}
      else if(conID.type == ObjectID::kTrack){
        basePtr->Tracks.push_back(conID.id);}
      else if(conID.type == ObjectID::kShower){
        basePtr->Showers.push_back(conID.id);}
      else if(conID.type == ObjectID::kCluster){
        basePtr->Clusters.push_back(conID.id);}
    }
  }
  
  // Nodes (probably only used by tracks, showers and particles, but
  // technically valid for all IReconBase objects.
  const COMET::IReconNodeContainer& nodeCont = baseObject->GetNodes();
  COMET::IReconNodeContainer::const_iterator node;
  for (node = nodeCont.begin();
       node != nodeCont.end();
       node++){
    int nodeID = FillNode(*node, saveHits);
    basePtr->Nodes.push_back(nodeID);
  }
  
  // Muon Decay Clusters
  for (IDataVector::iterator d = baseObject->begin();
       d != baseObject->end(); ++d) {
    if (!(*d)) continue;
    std::string sub_class_name = (*d)->ClassName();
    std::string fullname = (*d)->GetFullName().Data();
    
    if (sub_class_name.find("COMET::IReconObjectContainer") !=  std::string::npos &&
        fullname.find("//unnamed/decays") != std::string::npos) {
      COMET::IHandle<COMET::IReconObjectContainer> clusters
        = (*d)->Get<COMET::IReconObjectContainer>(".");
      if(!clusters) continue;
      
      COMET::IReconObjectContainer::const_iterator clus;
      for (clus = clusters->begin();
           clus != clusters->end();
           clus++){
        ObjectID conID = FillReconObject(*clus, saveHits);
        if(conID.type == ObjectID::kCluster){
          basePtr->Clusters.push_back(conID.id);}
        else
          COMETError("Muon Decay object not a cluster");
      }
    }
  }

};

short 
COMET::IP0DReconModule::FillAlgorithmResult(const COMET::IHandle<COMET::IAlgorithmResult> algoRes,
                                         short parent){

  TString fullName = algoRes->GetFullName();
  TString baseName = fullName;
  baseName.Remove(0,baseName.Last('/') + 1);

  // To save space, we're not going to save a number of the lower
  // algorithm results.  It is important to remember, that as soon as
  // a result is not saved, none of it's daughter results will be
  // saved either.
  for(std::vector<TPRegexp>::iterator cut = fRejectAlgoResultList.begin();
      cut != fRejectAlgoResultList.end(); ++cut){
    if(cut->MatchB(fullName, "", 0, 1)){
       COMETVerbose("Rejecting:\t" << baseName << " (" << cut->GetPattern() << ")");
       return -1;
    }
  }

  COMETVerbose("Fill Algorithm Result:\t" << baseName);

  int algoResID = fNAlgoResults;

  IP0DAlgoRes* p0dAlgoResult =
    new ((*fAlgoResults)[fNAlgoResults++]) IP0DAlgoRes;

  p0dAlgoResult->FullName      = fullName;
  p0dAlgoResult->AlgorithmName = baseName;
  p0dAlgoResult->Parent        = parent;

  // Get the used hits from the algo result
  COMET::IHandle<COMET::IHitSelection> usedHits = algoRes->GetHitSelection("used");

  // Save the used hits to a cluster
  if(usedHits && (!usedHits->empty())){
    COMET::IHitSelection::const_iterator hit;
    for(hit = usedHits->begin();
        hit != usedHits->end(); ++hit){
      short hitId = FillHit(*hit);
      if(hitId >= 0)
        p0dAlgoResult->Hits.push_back(hitId);
    }
    COMET::IHandle<COMET::IReconCluster> usedCluster(new COMET::IReconCluster);
    usedCluster->FillFromHits((baseName+"Used").Data(), *usedHits);

    ObjectID clusterID = FillReconObject(usedCluster, true);
    p0dAlgoResult->UsedHitCluster = clusterID.id;
  } 
  else p0dAlgoResult->UsedHitCluster     = -1;

  // Get the unused hits from the algo result
  COMET::IHandle<COMET::IHitSelection> unusedHits = algoRes->GetHitSelection("unused");

  // Save the unused hits to a cluster
  if(unusedHits && (!unusedHits->empty())){
    COMET::IHandle<COMET::IReconCluster> unusedCluster(new COMET::IReconCluster);
    unusedCluster->FillFromHits((baseName+"Unused").Data(), *unusedHits);

    // We don't want to save the unused hits. The 'false' prevents the 
    ObjectID clusterID = FillReconObject(unusedCluster, false);
    p0dAlgoResult->UnusedHitCluster = clusterID.id;
  } 
  else p0dAlgoResult->UnusedHitCluster     = -1;

  // Get the final recon object container
  COMET::IHandle<COMET::IReconObjectContainer> final
    = algoRes->GetResultsContainer("final");
  if(final){
    COMET::IReconObjectContainer::const_iterator reconObject;
    for (reconObject = final->begin(); reconObject != final->end(); reconObject++){
      ObjectID objID = FillReconObject(*reconObject, true);
      if(objID.type == ObjectID::kVertex){
  p0dAlgoResult->Vertices.push_back(objID.id);}
      else if(objID.type == ObjectID::kParticle){
  p0dAlgoResult->Particles.push_back(objID.id);}
      else if(objID.type == ObjectID::kTrack){
  p0dAlgoResult->Tracks.push_back(objID.id);}
      else if(objID.type == ObjectID::kShower){
  p0dAlgoResult->Showers.push_back(objID.id);}
      else if(objID.type == ObjectID::kCluster){
        p0dAlgoResult->Clusters.push_back(objID.id);}
    }
  }

  for (IDataVector::iterator d = algoRes->begin();
       d != algoRes->end(); ++d) {
    if (!(*d)) continue;
    std::string class_name = (*d)->ClassName();

    if (class_name.find("COMET::IAlgorithmResult") != std::string::npos) {
      COMET::IHandle<COMET::IAlgorithmResult> subResult
        = (*d)->Get<COMET::IAlgorithmResult>(".");
      if(!subResult) continue;
      int daughterID = FillAlgorithmResult(subResult, algoResID);
      if(daughterID != -1)
        p0dAlgoResult->AlgoResults.push_back(daughterID);
    }
  }

  return algoResID;
}

COMET::IP0DReconModule::ObjectID
COMET::IP0DReconModule::FillReconObject(const COMET::IHandle<COMET::IReconBase> reconObject,
                                     bool saveHits = true){

  std::string reco_class_name = reconObject->ClassName();

  if(reco_class_name.find("COMET::IReconVertex") != std::string::npos){
    COMETVerbose("Fill Vertex");
    return FillVertexObject(reconObject, saveHits);}
  else if(reco_class_name.find("COMET::IReconPID") != std::string::npos){
    COMETVerbose("Fill PID");
    return FillParticleObject(reconObject, saveHits);}
  else if(reco_class_name.find("COMET::IReconTrack") != std::string::npos){
    COMETVerbose("Fill Track");
    return FillTrackObject(reconObject, saveHits);}
  else if(reco_class_name.find("COMET::IReconShower") != std::string::npos){
    COMETVerbose("Fill Shower");
    return FillShowerObject(reconObject, saveHits);}
  else if(reco_class_name.find("COMET::IReconCluster") != std::string::npos){
    COMETVerbose("Fill Cluster");
    return FillClusterObject(reconObject, saveHits);}

  COMETError("Unknown object passed to IP0DReconModule");
  ObjectID blank;
  return blank;
}

COMET::IP0DReconModule::ObjectID
COMET::IP0DReconModule::FillVertexObject(const COMET::IHandle<COMET::IReconVertex> vertex,
                                      bool saveHits = true){
  ObjectID vertexID;
  vertexID.id   = fNVertices;
  vertexID.type = ObjectID::kVertex;

  IP0DVertex* p0dVertex =
    new ((*fVertices)[fNVertices++]) IP0DVertex;

  FillBaseObject(p0dVertex, vertex, saveHits);

  COMET::IHandle<COMET::IVertexState> vertexState = vertex->GetState();

  p0dVertex->Status        = vertex->GetStatus();
  p0dVertex->Quality       = vertex->GetQuality();
  p0dVertex->NDOF          = vertex->GetNDOF();
  p0dVertex->Position      = vertexState->GetPosition();
  p0dVertex->PosVariance   = vertexState->GetPositionVariance();
  p0dVertex->Fiducial
    = COMET::IGeomInfo::P0D().WaterFiducialDistance(p0dVertex->Position.X(),
             p0dVertex->Position.Y(),
             p0dVertex->Position.Z());

  p0dVertex->ValidDimensions =
    CountValidDimensions(p0dVertex->PosVariance.Vect());

  p0dVertex->Truth_TrajIDs.clear();
  p0dVertex->Truth_HitCount.clear();
  p0dVertex->Truth_ChargeShare.clear();

  COMET::IHandle<COMET::IHitSelection> vertexHits = vertex->GetHits();
  if (vertexHits){
    std::map< int, std::pair<int, float> > hitInfo = HitTruthInfo(vertexHits);

    std::map< int, std::pair< int, float > >::iterator parentID;
    for(parentID = hitInfo.begin(); parentID != hitInfo.end(); ++parentID){
      p0dVertex->Truth_TrajIDs.push_back((*parentID).first);
      p0dVertex->Truth_HitCount.push_back((*parentID).second.first);
      p0dVertex->Truth_ChargeShare.push_back((*parentID).second.second);
    }

    p0dVertex->Truth_PrimaryTrajIDs = HitTruthPrimaryInfo(vertexHits);
  }

  return vertexID;
}

COMET::IP0DReconModule::ObjectID
COMET::IP0DReconModule::FillParticleObject(const COMET::IHandle<COMET::IReconPID> particle,
                                        bool saveHits = true){
  ObjectID particleID;
  particleID.id   = fNParticles;
  particleID.type = ObjectID::kParticle;

  IP0DParticle* p0dParticle =
    new ((*fParticles)[fNParticles++]) IP0DParticle;

  FillBaseObject(p0dParticle, particle, saveHits);

  COMET::IHandle<COMET::IPIDState> particleState = particle->GetState();

  p0dParticle->Status        = particle->GetStatus();
  p0dParticle->Quality       = particle->GetQuality();
  p0dParticle->NDOF          = particle->GetNDOF();
  p0dParticle->Position      = particleState->GetPosition();
  p0dParticle->PosVariance   = particleState->GetPositionVariance();
  p0dParticle->Direction     = particleState->GetDirection();
  p0dParticle->DirVariance   = particleState->GetDirectionVariance();
  p0dParticle->Momentum      = particle->GetMomentum();
  p0dParticle->Charge        = particle->GetCharge();

  p0dParticle->ValidDimensions =
    CountValidDimensions(p0dParticle->PosVariance.Vect());

  const COMET::IHandle<COMET::IHitSelection> particleHits = particle->GetHits();
  if (particleHits){
    std::map< int, std::pair<int, float> > hitInfo = HitTruthInfo(particleHits);

    std::map< int, std::pair< int, float > >::iterator parentID;
    for(parentID = hitInfo.begin(); parentID != hitInfo.end(); ++parentID){
      p0dParticle->Truth_TrajIDs.push_back((*parentID).first);
      p0dParticle->Truth_HitCount.push_back((*parentID).second.first);
      p0dParticle->Truth_ChargeShare.push_back((*parentID).second.second);
    }

    p0dParticle->Truth_PrimaryTrajIDs = HitTruthPrimaryInfo(particleHits);
  }

  p0dParticle->SideDeposit = 0.0;
  p0dParticle->EndDeposit  = 0.0;

  COMET::IHitSelection::const_iterator hit;
  for (hit = particleHits->begin();
       hit != particleHits->end();
       ++hit) {

    COMET::IGeometryId geomid = (*hit)->GetGeomId();
    if (!COMET::GeomId::P0D::IsP0D(geomid)) continue;
    int p0dule = COMET::IGeomInfo::P0D().GetBar(geomid).GetP0Dule();
    int layer  = COMET::IGeomInfo::P0D().GetBar(geomid).GetLayer();
    int bar    = COMET::IGeomInfo::P0D().GetBar(geomid).GetNumber();
    if(p0dule > 38){
      p0dParticle->EndDeposit += (*hit)->GetCharge();
    }
    if(bar < 4 || (layer == 0 && bar > 121) || (layer == 1 && bar > 129)){
      p0dParticle->SideDeposit += (*hit)->GetCharge();
    }

  }

  // The  IReconPID  may  have  extra  PID  variables  stored  as  TReal-  and
  // TIntegerData.  Check to  see if they exist, and store  the values if they
  // do.

  // std::copy is required when the input vectors are doubles, but we
  // only store floats.  It is not required for integers, but it is
  // easier to keep the code the same.

  // Currently we only have TRealDatums as inputs (although some could
  // be stored as TIntegerDatums), but the code is included for
  // completeness.

  for(COMET::IDataVector::iterator dvIter = particle->begin();
      dvIter != particle->end(); ++dvIter){
    COMET::IRealDatum *rDatum = dynamic_cast<COMET::IRealDatum*>(*dvIter);
    if(rDatum){
      // Save the datum name, but remove the //unnamed/ at the beginning.
      TString datumName = rDatum->GetFullName();
      Ssiz_t lastSlash = datumName.Last('/');
      TSubString pidAlgoName = datumName(lastSlash + 1, datumName.Length());

      std::vector<float> tempVector;
      std::copy(rDatum->GetVector().begin(),
                rDatum->GetVector().end(),
                std::back_inserter(tempVector));

      p0dParticle->realPIDNames.push_back(pidAlgoName.Data());
      p0dParticle->realPIDValues.push_back(tempVector);
    }


    COMET::IIntegerDatum *iDatum = dynamic_cast<COMET::IIntegerDatum*>(*dvIter);
    if(iDatum){
      // Save the datum name, but remove the //unnamed/ at the beginning.
      TString datumName = iDatum->GetFullName();
      Ssiz_t lastSlash = datumName.Last('/');
      TSubString pidAlgoName = datumName(lastSlash + 1, datumName.Length());

      std::vector<short> tempVector;
      std::copy(iDatum->GetVector().begin(),
                iDatum->GetVector().end(),
                std::back_inserter(tempVector));

      p0dParticle->integerPIDNames.push_back(pidAlgoName.Data());
      p0dParticle->integerPIDValues.push_back(tempVector);
    }
  }

  // The most  likely PID is  available through particle->GetParticle,
  // but it is also stored as one of the alternates.  These are sorted
  // in order of PID, so copy them directly into the TTree.
  COMET::IReconObjectContainer::const_iterator alternate;

  for(alternate = particle->GetAlternates().begin();
      alternate != particle->GetAlternates().end();
      alternate++){
    COMET::IHandle<COMET::IReconPID> alt = (*alternate);
    if(alt){
      p0dParticle->PID.push_back(alt->GetParticleId());
      p0dParticle->PID_weight.push_back(alt->GetPIDWeight());
    }
    else COMETWarn("IReconPID alternate is not a IReconPID - Ignored");
  }

  return particleID;
}

COMET::IP0DReconModule::ObjectID
COMET::IP0DReconModule::FillTrackObject(const COMET::IHandle<COMET::IReconTrack> track,
                                     bool saveHits = true){
  ObjectID trackID;
  trackID.id   = fNTracks;
  trackID.type = ObjectID::kTrack;

  IP0DTrack* p0dTrack =
    new ((*fTracks)[fNTracks++]) IP0DTrack;

  FillBaseObject(p0dTrack, track, saveHits);

  COMET::IHandle<COMET::ITrackState> trackState = track->GetState();

  p0dTrack->Status        = track->GetStatus();
  p0dTrack->Quality       = track->GetQuality();
  p0dTrack->NDOF          = track->GetNDOF();
  p0dTrack->Position      = trackState->GetPosition();
  p0dTrack->PosVariance   = trackState->GetPositionVariance();
  p0dTrack->Direction     = trackState->GetDirection();
  p0dTrack->DirVariance   = trackState->GetDirectionVariance();
  p0dTrack->EDeposit      = trackState->GetEDeposit();

  p0dTrack->ValidDimensions =
    CountValidDimensions(p0dTrack->PosVariance.Vect());

  const COMET::IHandle<COMET::IHitSelection> trackHits = track->GetHits();
  if (trackHits) {
    std::map< int, std::pair<int, float> > hitInfo = HitTruthInfo(trackHits);

    std::map< int, std::pair< int, float > >::iterator parentID;
    for(parentID = hitInfo.begin(); parentID != hitInfo.end(); ++parentID){
      p0dTrack->Truth_TrajIDs.push_back((*parentID).first);
      p0dTrack->Truth_HitCount.push_back((*parentID).second.first);
      p0dTrack->Truth_ChargeShare.push_back((*parentID).second.second);
    }

    p0dTrack->Truth_PrimaryTrajIDs = HitTruthPrimaryInfo(trackHits);
  }

  p0dTrack->SideDeposit = 0.0;
  p0dTrack->EndDeposit  = 0.0;

  COMET::IHitSelection::const_iterator hit;
  for (hit = trackHits->begin();
       hit != trackHits->end();
       ++hit) {

    COMET::IGeometryId geomid = (*hit)->GetGeomId();
    if (!COMET::GeomId::P0D::IsP0D(geomid)) continue;
    int p0dule = COMET::IGeomInfo::P0D().GetBar(geomid).GetP0Dule();
    int layer  = COMET::IGeomInfo::P0D().GetBar(geomid).GetLayer();
    int bar    = COMET::IGeomInfo::P0D().GetBar(geomid).GetNumber();
    if(p0dule > 38){
      p0dTrack->EndDeposit += (*hit)->GetCharge();
    }
    if(bar < 4 || (layer == 0 && bar > 121) || (layer == 1 && bar > 129)){
      p0dTrack->SideDeposit += (*hit)->GetCharge();
    }

  }

  const COMET::IReconNodeContainer& nodeCont = track->GetNodes();

  // Get the track length from the nodes
  p0dTrack->Length = 0.0;
  COMET::IReconNodeContainer::const_iterator node = nodeCont.begin();
  COMET::IHandle<COMET::ITrackState> nodeState = (*node)->GetState();
  TVector3 frontPos = nodeState->GetPosition().Vect();
  ++node;

  while(node != nodeCont.end()){
    nodeState = (*node)->GetState();

    p0dTrack->Length += (frontPos - nodeState->GetPosition().Vect()).Mag();
    frontPos = nodeState->GetPosition().Vect();

    ++node;
  }

  return trackID;
}

COMET::IP0DReconModule::ObjectID
COMET::IP0DReconModule::FillShowerObject(const COMET::IHandle<COMET::IReconShower> shower,
                                      bool saveHits = true){
  ObjectID showerID;
  showerID.id   = fNShowers;
  showerID.type = ObjectID::kShower;

  IP0DShower* p0dShower =
    new ((*fShowers)[fNShowers++]) IP0DShower;

  FillBaseObject(p0dShower, shower, saveHits);

  // For Shower objects, the cluster information of the nodes is vital for the
  // pi0 analysis.  However, the cluster information for most nodes is
  // useless.  Therefore, rather than adding a "FillClusterObject" to each
  // "FillNode", we will do a parallel fill, within the
  // FillShowerObject.
  // N.B. FillNode has already been run within FillBaseObject.

  const COMET::IReconNodeContainer& nodeCont = shower->GetNodes();
  COMET::IReconNodeContainer::const_iterator node;
  for (node = nodeCont.begin();
       node != nodeCont.end();
       node++){
    COMET::IHandle<COMET::IReconCluster> cluster = (*node)->GetObject();
    COMETVerbose("Fill Cluster - Shower");
    ObjectID clusterID = FillClusterObject(cluster, saveHits);
    p0dShower->Clusters.push_back(clusterID.id);
  }

  COMET::IHandle<COMET::IShowerState> showerState = shower->GetState();

  p0dShower->Status        = shower->GetStatus();
  p0dShower->Quality       = shower->GetQuality();
  p0dShower->NDOF          = shower->GetNDOF();
  p0dShower->Position      = showerState->GetPosition();
  p0dShower->PosVariance   = showerState->GetPositionVariance();
  p0dShower->Direction     = showerState->GetDirection();
  p0dShower->DirVariance   = showerState->GetDirectionVariance();
  p0dShower->EDeposit      = showerState->GetEDeposit();
  p0dShower->Cone          = showerState->GetCone()[0];

  p0dShower->ValidDimensions =
    CountValidDimensions(p0dShower->PosVariance.Vect());

  p0dShower->Width         = 0.0;
  p0dShower->Length        = 0.0;
  p0dShower->SideDeposit   = 0.0;
  p0dShower->EndDeposit    = 0.0;

  const COMET::IHandle<COMET::IHitSelection> showerHits = shower->GetHits();

  double totalCharge = 0.0;
  float avgLen = 0.0;
  float avgLen2 = 0.0;

  if (showerHits) {
    std::map< int, std::pair<int, float> > hitInfo = HitTruthInfo(showerHits);

    std::map< int, std::pair< int, float > >::iterator parentID;
    for(parentID = hitInfo.begin(); parentID != hitInfo.end(); ++parentID){
      p0dShower->Truth_TrajIDs.push_back((*parentID).first);
      p0dShower->Truth_HitCount.push_back((*parentID).second.first);
      p0dShower->Truth_ChargeShare.push_back((*parentID).second.second);
    }

    p0dShower->Truth_PrimaryTrajIDs = HitTruthPrimaryInfo(showerHits);

    COMET::IHitSelection::const_iterator hit;

    for (hit = showerHits->begin();
   hit != showerHits->end();
   ++hit) {

      // The distance vector between the hit and the shower position
      TVector3 diff = (*hit)->GetPosition() - p0dShower->Position.Vect();

      // The distance vector along the shower direction
      TVector3 para = (diff * p0dShower->Direction) * p0dShower->Direction;

      // The distance vector normal to the shower direction
      TVector3 perp = diff - para;

      if ((*hit)->IsXHit()) {
        diff.SetY(0.0);
        perp.SetY(0.0);
      }
      else {
        diff.SetX(0.0);
        perp.SetX(0.0);
      }

      float q =  (*hit)->GetCharge();
      totalCharge += q;

      p0dShower->Width += q * perp.Mag();

      avgLen  += q * para.Mag();
      avgLen2 += q * (para * para);

      // Get the side and end charge
      COMET::IGeometryId geomid = (*hit)->GetGeomId();
      if (!COMET::GeomId::P0D::IsP0D(geomid)) continue;
      int p0dule = COMET::IGeomInfo::P0D().GetBar(geomid).GetP0Dule();
      int layer  = COMET::IGeomInfo::P0D().GetBar(geomid).GetLayer();
      int bar    = COMET::IGeomInfo::P0D().GetBar(geomid).GetNumber();
      if(p0dule > 38){
        p0dShower->EndDeposit += (*hit)->GetCharge();
      }
      if(bar < 4 || (layer == 0 && bar > 121) || (layer == 1 && bar > 129)){
        p0dShower->SideDeposit += (*hit)->GetCharge();
      }

    }
  }

  p0dShower->Width /= std::max(totalCharge,1.0);

  avgLen  /= std::max(totalCharge,1.0);
  avgLen2 /= std::max(totalCharge,1.0);
  double len2 = avgLen2 - avgLen*avgLen;

  p0dShower->Length = std::sqrt(std::max(0.0, len2));

  return showerID;
}

COMET::IP0DReconModule::ObjectID 
COMET::IP0DReconModule::FillClusterObject(const COMET::IHandle<COMET::IReconCluster> cluster,
                                       bool saveHits = true){

  ObjectID clusterID;
  clusterID.id   = fNClusters;
  clusterID.type = ObjectID::kCluster;

  IP0DCluster* p0dCluster =
    new ((*fClusters)[fNClusters++]) IP0DCluster;

  FillBaseObject(p0dCluster,cluster, saveHits);

  p0dCluster->NFiducialHits = -1;
  p0dCluster->EDeposit      = cluster->GetEDeposit();
  p0dCluster->Position      = cluster->GetPosition();
  p0dCluster->PosVariance   = cluster->GetPositionVariance();
  
  p0dCluster->ValidDimensions =
    CountValidDimensions(p0dCluster->PosVariance.Vect());

  TMatrixF moments = cluster->GetMoments();
  if(p0dCluster->arraySize == (moments.GetNrows() *
                               moments.GetNcols())){
    std::copy(moments.GetMatrixArray(),
              moments.GetMatrixArray()+p0dCluster->arraySize,
              p0dCluster->Moments);
  }

  const COMET::IHandle<COMET::IHitSelection> clusterHits = cluster->GetHits();
  if (clusterHits){
    std::map< int, std::pair<int, float> > hitInfo = HitTruthInfo(clusterHits);

    std::map< int, std::pair< int, float > >::iterator parentID;
    for(parentID = hitInfo.begin(); parentID != hitInfo.end(); ++parentID){
      p0dCluster->Truth_TrajIDs.push_back((*parentID).first);
      p0dCluster->Truth_HitCount.push_back((*parentID).second.first);
      p0dCluster->Truth_ChargeShare.push_back((*parentID).second.second);
    }

    p0dCluster->Truth_PrimaryTrajIDs = HitTruthPrimaryInfo(clusterHits);

    p0dCluster->NFiducialHits = 0;

    for (COMET::IHitSelection::const_iterator hit = clusterHits->begin();
   hit != clusterHits->end(); ++hit){
      TVector3 hitPosition = (*hit)->GetPosition();
      float fiducial = COMET::IGeomInfo::P0D().WaterFiducialDistance(hitPosition.X(),
                                                                  hitPosition.Y(),
                                                                  hitPosition.Z());

      if (fiducial > 0) p0dCluster->NFiducialHits += 1;

    }
  }

  return clusterID;
}

short COMET::IP0DReconModule::FillNode(const COMET::IHandle<COMET::IReconNode> node, bool saveHits){
  short nodeID = fNNodes;
  COMETVerbose("Fill Node");

  IP0DNode* p0dNode =
    new ((*fNodes)[fNNodes++]) IP0DNode;

  // The node state may be a PIDState, TrackState or ShowerState, but
  // they are all TMPositionDirectionStates, so we access the position
  // and direction regardless.
  COMET::IHandle<COMET::IReconState> s1 = node->GetState();
  if (s1) {
    const COMET::IMPositionDirectionState* posDirState = NULL;
    posDirState =
      dynamic_cast<const COMET::IMPositionDirectionState*>(COMET::GetPointer(s1));
    if(posDirState){
      p0dNode->Position      = posDirState->GetPosition();
      p0dNode->PosVariance   = posDirState->GetPositionVariance();
      p0dNode->Direction     = posDirState->GetDirection();
      p0dNode->DirVariance   = posDirState->GetDirectionVariance();

      p0dNode->ValidDimensions =
        CountValidDimensions(p0dNode->PosVariance.Vect());
    }
  }

  // Every node should have an object

  COMET::IHandle<COMET::IReconBase> nodeObject = node->GetObject();
  if(nodeObject){
    COMET::IHandle<COMET::IReconState> objState = nodeObject->GetState();
    if(objState){
      const COMET::IMEDepositState* eDepState = NULL;
      eDepState =
      dynamic_cast<const COMET::IMEDepositState*>(COMET::GetPointer(s1));
      if(eDepState){
        p0dNode->EDeposit = eDepState->GetEDeposit();
      }
    }

    const COMET::IHandle<COMET::IHitSelection> nodeHits = nodeObject->GetHits();
    if (nodeHits){
      COMET::IHitSelection::const_iterator hit;
      for(hit = nodeHits->begin();
    hit != nodeHits->end(); ++hit){
  short hitId = FillHit(*hit);
  if(hitId >= 0)
    p0dNode->Hits.push_back(hitId);
      }
      std::map< int, std::pair<int, float> > hitInfo = HitTruthInfo(nodeHits);

      std::map< int, std::pair< int, float > >::iterator parentID;
      for(parentID = hitInfo.begin(); parentID != hitInfo.end(); ++parentID){
        p0dNode->Truth_TrajIDs.push_back((*parentID).first);
        p0dNode->Truth_HitCount.push_back((*parentID).second.first);
        p0dNode->Truth_ChargeShare.push_back((*parentID).second.second);
      }

      p0dNode->Truth_PrimaryTrajIDs = HitTruthPrimaryInfo(nodeHits);
    }
  }

  return nodeID;
}

short 
COMET::IP0DReconModule::FillHit(const COMET::IHandle<COMET::IHit> hit){
  if (!hit) return -1;
  COMET::IHandle<COMET::ISingleHit> sh = hit;
  if (!sh) return -1;  

  // Check if this Hit has already been filled, if so pass back the
  // reference, instead of creating a new one.
  UInt_t chanIdUInt = sh->GetChannelId().AsUInt();
  std::map<UInt_t, short>::const_iterator mapIter;
  mapIter = fTempHitMap.find(chanIdUInt);
  if(mapIter != fTempHitMap.end())
    return mapIter->second;

  // We only want to save P0D hits at this time.  If a hit is from
  // outside the P0D, skip and return -1.
  if(sh->GetChannelId().GetSubDetector() != COMET::IChannelId::kP0D)
      return -1;

  // Otherwise, make a new entry and save it in the TempHitMap
  short hitID = fNHits;
  fTempHitMap[chanIdUInt] = hitID;

  IP0DHit* p0dHit =
    new ((*fHits)[fNHits++]) IP0DHit;

  p0dHit->ChanID = sh->GetChannelId().AsUInt();
  p0dHit->GeomID = sh->GetGeomId().AsInt();
  p0dHit->Charge = sh->GetCharge();
  p0dHit->Time   = sh->GetTime();

  return hitID;
}

std::map< int, std::pair<int, float> >
COMET::IP0DReconModule::HitTruthInfo(const COMET::IHandle<COMET::IHitSelection> hits){
  std::map< int, std::pair<int, float> > result;

  // then loop over everything and log the trajectory ids.
  for (COMET::IHitSelection::const_iterator hit = hits->begin();
       hit != hits->end(); hit++) {
    float hitCharge = (*hit)->GetCharge();

    std::map< int, int > hitCount;
    std::map< int, float > chargeShare;

    std::vector<COMET::IG4VHit*> hitContribs = HitTruthInfo::GetHitTruthInfo(*hit);

    for (std::vector<COMET::IG4VHit*>::const_iterator g4Hit = hitContribs.begin();
         g4Hit != hitContribs.end(); g4Hit++) {
      COMET::IG4HitSegment* g4HitCast = dynamic_cast<COMET::IG4HitSegment*>(*g4Hit);

      // For each contributing trajectory, increase the hit count and
      // give it an equal share of the energy deposit (the share may
      // not be equal, but that information is not currently
      // available).
      std::vector<int>::const_iterator trajIDIter;
      for (trajIDIter  = g4HitCast->GetContributors().begin();
           trajIDIter != g4HitCast->GetContributors().end();
           ++trajIDIter){
        hitCount[(*trajIDIter)]++;
        chargeShare[(*trajIDIter)] += g4HitCast->GetEnergyDeposit() / g4HitCast->GetContributors().size();
      }
    }

    for (std::map< int, int >::const_iterator hC = hitCount.begin();
         hC != hitCount.end(); ++hC){
      // If this hit  had a contribution from a given  ID, then increment that
      // ID's hitcount.
      if( (*hC).second > 0) result[(*hC).first].first++;
    }

    // Count up the  contributed charge (this won't necessarily  match the hit
    // Charge).
    float totalChargeShare = 0;
    for (std::map< int, float >::const_iterator cS = chargeShare.begin();
         cS != chargeShare.end(); ++cS){
      totalChargeShare += (*cS).second;
    }

    // For each entry, determine the contributed share.
    for (std::map< int, float >::const_iterator cS = chargeShare.begin();
         cS != chargeShare.end(); ++cS){
      result[(*cS).first].second += hitCharge * (*cS).second / totalChargeShare;
    }
  }

  return result;
}

std::vector<int>
COMET::IP0DReconModule::HitTruthPrimaryInfo(const COMET::IHandle<COMET::IHitSelection> hits){
  std::vector<int> result;

  // Loop over all hits, storing each primary ID which contributed
  for (COMET::IHitSelection::const_iterator hit = hits->begin();
       hit != hits->end(); hit++) {
    std::vector<COMET::IG4VHit*> hitContribs = HitTruthInfo::GetHitTruthInfo(*hit);

    for (std::vector<COMET::IG4VHit*>::const_iterator g4Hit = hitContribs.begin();
         g4Hit != hitContribs.end(); g4Hit++) {
      COMET::IG4HitSegment* g4HitCast = dynamic_cast<COMET::IG4HitSegment*>(*g4Hit);

      result.push_back(g4HitCast->GetPrimaryId());
    }
  }

  // Sort the list of primary IDs
  std::sort(result.begin(), result.end());

  // Remove duplicate IDs
  std::vector<int>::iterator newEnd =
    std::unique(result.begin(), result.end());
  result.erase(newEnd, result.end());

  return result;
}

short
COMET::IP0DReconModule::GetCycle(const COMET::IHandle<COMET::IHitSelection> hits){
  // As of July 2012, the ReconP0D results may contain objects from
  // outside the P0D.  Currently, this is only due to the Incremental
  // matching with the TPC.  Therefore, we consider all potential
  // hits, looking for a TFB hit.

  // Access the channel ID for each hit
  COMET::ITFBChannelId chanID;//(hits->front()->GetChannelId());
  COMET::IHitSelection::const_iterator hit;

  for(hit = hits->begin(); hit != hits->end(); ++hit){
    COMET::IHandle<COMET::ISingleHit> sh = (*hit);
    if (!sh) continue;
    chanID = COMET::ITFBChannelId(sh->GetChannelId());
    if(chanID.IsValid() && chanID.IsTFBChannel() && chanID.GetSubDetector() == COMET::IChannelId::kP0D)
      break;
  }

  // If the ID is invalid, return -1 as the cycle
  if(!chanID.IsValid() || !chanID.IsTFBChannel()){
    //COMETError("Invalid P0D Event containing no cycles!");
    return -1;
  }

  // Otherwise, return the cycle of the channel ID
  short cycle = (short) chanID.GetCapacitor();
  return cycle;
}

short
COMET::IP0DReconModule::CountValidDimensions(TVector3 posVar){
  short valid = 0;
  if (std::abs(posVar.X()) < 1E+6) valid += 1;
  if (std::abs(posVar.Y()) < 1E+6) valid += 2;
  if (std::abs(posVar.Z()) < 1E+6) valid += 4;
  return valid;
}