#include "TrackTruthInfo.hxx"

#include <map>
#include <IEventFolder.hxx>
#include <IComboHit.hxx>
#include <IReconHit.hxx>
#include <IMultiHit.hxx>
#include <IMCDigit.hxx>
#include <IG4HitSegment.hxx>
#include <COMETGeomIdDef.hxx> 


#include "retrieveHitTruthInfo.hxx"

namespace TrackTruthInfo {


std::map<int, int> GetG4Trajectories(const COMET::IHitSelection &hits)
{
  std::map<int, int> counts;

  // then loop over everything and log the trajectory ids.
  for (COMET::IHitSelection::const_iterator hit = hits.begin();
       hit != hits.end(); hit++) {
    
    std::vector<COMET::IG4VHit*> contributors = HitTruthInfo::GetHitTruthInfo(*hit);

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

      for (unsigned int i = 0; i < g4HitCast->GetContributors().size(); i++) {
        int id = g4HitCast->GetContributors()[i];

        // a std::map initializes any unused elements to zero
        counts[id]++;
      }
    }
  }

  return counts;
}


std::map<int, double> GetG4TrajectoryFractions(const COMET::IHandle<COMET::IReconBase> track)
{
  COMET::IHandle<COMET::IHitSelection> hits = track->GetHits();
  if (!hits) {
    COMETWarn("TrackTruthInfo::GetG4TrajectoryID : Track has no hits!");
    return std::map<int,double>();
  }

  std::map<int, int> counts = GetG4Trajectories(*hits);

  int total = 0;
  for (std::map<int, int>::iterator t = counts.begin();
       t != counts.end(); t++)
    total += t->second;

  std::map<int, double> fractions;
  for (std::map<int,int>::iterator t = counts.begin();
       t != counts.end(); t++) 
    fractions[t->first] = ((double) t->second)/((double) total);

  return fractions;
}


int GetG4TrajIDHits(const COMET::IHitSelection &hits)
{
  std::map<int, int> counts = GetG4Trajectories(hits);

  int maxID = 0;
  int maxCount = 0;

  if (counts.empty()) {
    COMETWarn("TrackTruthInfo::GetG4TrajIDHits : None of the hits have associated MC info.");
    return 0;
  }

  for (std::map<int, int>::iterator tr = counts.begin();
       tr != counts.end(); tr++) {
    if (tr->second > maxCount) {
      maxCount = tr->second;
      maxID = tr->first;
    }
  }

  return maxID;
}


int GetG4TrajectoryID(const COMET::IHandle<COMET::IReconBase> track)
{
  COMET::IHandle<COMET::IHitSelection> hits = track->GetHits();
  if (!hits) {
    COMETWarn("TrackTruthInfo::GetG4TrajectoryID : Track has no hits!");
    return 0;
  }

  return GetG4TrajIDHits(*hits);
}


COMET::IHandle<COMET::IG4Trajectory> GetG4TrajHits(const COMET::IHitSelection &hits)
{
  int G4Id = GetG4TrajIDHits(hits);
      if (!G4Id){
              return COMET::IHandle<COMET::IG4Trajectory>();
      }

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

      return trajectories->GetTrajectory(G4Id);

}


COMET::IHandle<COMET::IG4Trajectory> GetG4Trajectory(const COMET::IHandle<COMET::IReconBase> track)
{
  int G4Id = GetG4TrajectoryID(track);
      if (!G4Id){
              return COMET::IHandle<COMET::IG4Trajectory>();
      }

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

      return trajectories->GetTrajectory(G4Id);
}

COMET::IHandle<COMET::IG4Trajectory> GetG4Trajectory(const COMET::IHandle<COMET::IReconBase> track, double &complete, double &clean)
{
  int G4Id = GetG4TrajectoryID(track, complete, clean);
      if (!G4Id){
              return COMET::IHandle<COMET::IG4Trajectory>();
      }

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

      return trajectories->GetTrajectory(G4Id);
}

bool HitMatchesTrajectory(COMET::IHandle<COMET::IHit> hit, int g4id)
{
  std::vector<COMET::IG4VHit*> contributors = HitTruthInfo::GetHitTruthInfo(hit);

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

    for (unsigned int i = 0; i < g4HitCast->GetContributors().size(); i++)
      if (g4HitCast->GetContributors()[i] == g4id) 
        return true;
  }
  return false;
}


COMET::IHandle<COMET::IHitSelection> GetTrueTrackHits(int g4id, COMET::IHandle<COMET::IHitSelection> universe) 
{
  COMET::IHandle<COMET::IHitSelection> trackHits(new COMET::IHitSelection());

  for (COMET::IHitSelection::iterator hit = universe->begin();
       hit != universe->end(); hit++) {
    if (HitMatchesTrajectory(*hit, g4id))
      trackHits->push_back(*hit);
  }

  return trackHits;
}

//*************************************************
std::map<int, int> GetG4Trajectories(COMET::IHandle<COMET::IHit> hit)
{
    std::map<int, int> counts;

    // then loop over everything and log the trajectory ids.

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

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

        for (unsigned int i = 0; i < g4HitCast->GetContributors().size(); i++) {
            int id = g4HitCast->GetContributors()[i];

            // a std::map initializes any unused elements to zero
            counts[id]++;
        }
    }

    return counts;
}

//*************************************************
int GetG4TrajIDHits(const COMET::IHitSelection &hits, double &complete, double &clean, std::string views)
{
    complete = -0xABCDEF;
    clean = -0xABCDEF;

    int TotalFoundPadsBars = 0;
    unsigned int DetectorUsed = 0;

    std::map<int,int> recoCount;

    const COMET::ICOMETEvent& event = *(COMET::IEventFolder::GetCurrentEvent());
    // Loop over everything and log the trajectory ids and the number of pads/bars for each id
    for (COMET::IHitSelection::const_iterator hit = hits.begin(); hit != hits.end(); hit++) {
        // Check if it is a IComboHit and expand
        COMET::IHandle<COMET::IComboHit> combohit(*hit);
        COMET::IHandle<COMET::IReconHit> recohit(*hit);
        if (combohit) {
            TotalFoundPadsBars += combohit->GetHits().size();
            for (COMET::IHitSelection::const_iterator hit2 = combohit->GetHits().begin();
                    hit2 != combohit->GetHits().end(); hit2++) {
                // use a recursive call to get the constituents of the combo hit
                std::map<int,int> tmpcount;
                tmpcount = GetG4Trajectories(*hit2);
                for (std::map<int,int>::iterator it = tmpcount.begin(); it != tmpcount.end(); it++)
                    recoCount[it->first]++;    // Just count the number of pads/bars, not the number of contributions
                DetectorUsed |= HitToDetectorId(*hit2);
                
            }
        }
        // If not a IComboHit, check if it's a IReconHit
        else if (recohit){
              TotalFoundPadsBars += recohit->GetContributorCount();    
            for ( int i=0; i < recohit->GetContributorCount(); i++){
                COMET::IHandle<COMET::IHit> hit3 = recohit->GetContributor(i);
                std::map<int,int> tmpcount;
                tmpcount = GetG4Trajectories(hit3);
                for (std::map<int,int>::iterator it = tmpcount.begin(); it != tmpcount.end(); it++)
                    recoCount[it->first]++;    // Just count the number of pads/bars, not the number of contributions
                DetectorUsed |= HitToDetectorId(hit3);
            }
        }
        else{
            COMET::IHandle<COMET::IHit> handleHit(*hit);
            TotalFoundPadsBars++;
            std::map<int,int> tmpcount = GetG4Trajectories(handleHit);
            for (std::map<int,int>::iterator it = tmpcount.begin(); it != tmpcount.end(); it++)
                recoCount[it->first]++;    // Just count the number of pads/bars, not the number of contributions
            DetectorUsed |= HitToDetectorId(*hit);
        }
        // If not a IComboHit, just get the constituents
    }


    int maxID = 0;
    int maxCount = 0;

    if (recoCount.empty()) {
        COMETWarn("TrackTruthInfo::GetG4TrajIDHits : None of the hits have associated MC info.");
        return 0;
    }

    for (std::map<int, int>::iterator tr = recoCount.begin();
            tr != recoCount.end(); tr++) {
        if (tr->second > maxCount) {
            maxCount = tr->second;
            maxID = tr->first;
        }
    }

    if ( ! DetectorUsed)    // Don't know where the hits are from, exit
        return maxID;

    // Cleanliness
    clean = (double)maxCount / (double)TotalFoundPadsBars;

    // Count the number of pads/bars from this true particle
    int TotalTruePadsBars = 0;
    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);
        for (COMET::IHitSelection::const_iterator hit = hitList->begin(); hit != hitList->end(); hit++) {
            if ( ! (DetectorUsed & HitToDetectorId(*hit)) )
                continue;
            COMET::IHandle<COMET::IHit> thit = *hit;
             if ( (views.find_first_of('3') != std::string::npos) ||
                 (views == "XY" && (thit->IsXHit() &&  thit->IsYHit() && !thit->IsZHit()) ) ||
                 (views == "YZ" && (!thit->IsXHit() &&  thit->IsYHit() && thit->IsZHit()) ) ||
                 (views == "XZ" && (thit->IsXHit() &&  !thit->IsYHit() && thit->IsZHit()) ) )
                if ( HitMatchesTrajectory(*hit, maxID) )
                    TotalTruePadsBars++;
        }
    }

    // Completeness
    if (TotalTruePadsBars)
        complete = (double)maxCount / (double)TotalTruePadsBars;
    else
        COMETWarn("TrackTruthInfo::GetG4TrajIDHits : Missing true hits for efficiency calculation.");

    return maxID;
}

//*************************************************
int GetG4TrajectoryID(const COMET::IHandle<COMET::IReconBase> track, double &complete, double &clean)
{
    COMET::IHandle<COMET::IHitSelection> hits = track->GetHits();
    if (!hits) {
        COMETWarn("TrackTruthInfo::GetG4TrajectoryID : Track has no hits!");
        return 0;
    }
    if (hits->size()<1) {
        COMETWarn("TrackTruthInfo::GetG4TrajectoryID : Track has no hits!");
        return 0;
    }

    std::string views = "33";
    COMET::IHandle<COMET::IReconTrack> rtrack = track;
    if(rtrack){
       if ( rtrack->IsXTrack() && rtrack->IsYTrack() && !rtrack->IsZTrack() ) views = "XY";
       else if ( rtrack->IsXTrack() && !rtrack->IsYTrack() && rtrack->IsZTrack() ) views = "XZ";
       else if ( !rtrack->IsXTrack() && rtrack->IsYTrack() && rtrack->IsZTrack() ) views = "YZ";
    }

    int G4Id = GetG4TrajIDHits(*hits, complete, clean, views);
    return G4Id;
}

//*************************************************
int HitToDetectorId(COMET::IHandle<COMET::IHit> hit)
{
  COMET::IGeometryId geomId = hit->GetGeomId();
  return GeomIdToDetectorId(geomId);
}

//*************************************************
int GeomIdToDetectorId(COMET::IGeometryId geomId) {
  // Use things in ND namespace
  using COMET::IReconBase;
  namespace GeomId = COMET::GeomId;
  
  // Due to circular a dependency, we can't use oaGeomInfo.
  // We must do this by hand. 
  //Start by doing all the detectors that are determined
  // uniquely by their subsystem ID : P0D, DsECal and SMRD
  
  int ssid = geomId.GetSubsystemId();
  if (ssid == GeomId::Def::kCDC)      return IReconBase::kCDC;
  if (ssid == GeomId::Def::kCTH)      return IReconBase::kCTH;
  if (ssid == GeomId::Def::kECAL)     return IReconBase::kECAL;
  if (ssid == GeomId::Def::kStrawTrk) return IReconBase::kStrawTrk;

  return 0;
}

//*************************************************
int GetDetectorCrossedByG4Traj(COMET::IHandle<COMET::IG4Trajectory> g4traj)
{
      return (GetDetectorCrossedByG4Traj(g4traj->GetTrackId()));
}

//*************************************************
int GetDetectorCrossedByG4Traj(int g4id)
{
    int DetectorCrossed = 0;
    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);
        for (COMET::IHitSelection::const_iterator hit = hitList->begin(); hit != hitList->end(); hit++) {
            if ( HitMatchesTrajectory(*hit, g4id)){
                DetectorCrossed = DetectorCrossed | HitToDetectorId(*hit);
              }
        }
    }
    return DetectorCrossed;
}

//*************************************************
COMET::IG4TrajectoryPoint GetCloserG4Point(COMET::IHandle<COMET::IG4Trajectory> traj, TVector3 *Position, double &dist){

  std::vector<COMET::IG4TrajectoryPoint> points = traj->GetTrajectoryPoints();

  double x0 = Position->X();
  double y0 = Position->Y();
  double z0 = Position->Z();

  double dmin = DBL_MAX;
  int icloser = -1;

  for(unsigned int i = 0 ; i < points.size() ; i++ ) {
    double deltaz = (points[i].GetPosition().Z()-z0);
    double deltay = (points[i].GetPosition().Y()-y0);
    double deltax = (points[i].GetPosition().X()-x0);
    double distance = TMath::Sqrt(deltax*deltax+deltay*deltay+deltaz*deltaz); 
    if( distance < dmin ) {
      dmin = distance;
      icloser = i;
    }
  }

  if( icloser >= 0 ) {
    dist = dmin; 
    return points[icloser];
  }
  else {
    COMETWarn("TrackTruthInfo::GetCloserG4Point : None of the true points are close to the hit.");
    return COMET::IG4TrajectoryPoint();
  }
}

}