#include <cmath>
#include <algorithm>
#include <iostream>

#include <tut.h>
#include <TGeoManager.h>

#include "tutRunSimulation.hxx"
#include "tutTestHits.hxx"

#include "HEPUnits.hxx"
#include "ICOMETEvent.hxx"
#include "IOADatabase.hxx"
#include "IGeomIdManager.hxx"
#include "IG4HitSegment.hxx"
#include "IG4Trajectory.hxx"

using namespace COMET;

namespace tut {
    struct baseHitSegments {
        baseHitSegments() {
            // Run before each test.
            GenerateSIMG4Events("HitSegments");
        }
        ~baseHitSegments() {
            // Run after each test.
        }
    };

    // Declare the test
    typedef test_group<baseHitSegments>::object testHitSegments;
    test_group<baseHitSegments> groupHitSegments("HitSegments");

    // Check that the first event sample has enough hit containers
    template<> template<>
    void testHitSegments::test<1> () {
        int evtC = TestHits::CountHitContainers<IG4HitSegment>(SimG4::StartEvents[0],
                                                                      SimG4::EndEvents[0],
                                                                      SimG4::Events);
        ensure_equals("Events with hit segments exist",
                      evtC, SimG4::EndEvents[0]-SimG4::StartEvents[0]);
    }

    // Check that all hit contributors are in the trajectory container.
    template<> template<>
    void testHitSegments::test<2> () {
        for (SimG4::EventVector::iterator e = SimG4::Events.begin() + SimG4::StartEvents[0];
             e != SimG4::Events.begin() + SimG4::EndEvents[0]; 
             ++e) {
            IHandle<IDataVector> dv=(*e)->Get<IDataVector>("truth/g4Hits");
            IHandle<IG4TrajectoryContainer> trajCon 
                = (*e)->Get<IG4TrajectoryContainer>("truth/G4Trajectories");
            for (IDataVector::iterator d = dv->begin();
                 d != dv->end();
                 ++d) {
                IHandle<IG4HitContainer> hitCon
                    = (*d)->Get<IG4HitContainer>(".");
                for (IG4HitContainer::iterator h = hitCon->begin();
                     h != hitCon->end();
                     ++h) {
                    const IG4HitSegment* hit 
                        = dynamic_cast<const IG4HitSegment*>(*h);
                    if (hit) {
                        for (std::vector<int>::const_iterator c
                                 = hit->GetContributors().begin();
                             c != hit->GetContributors().end();
                             ++c) {
                            IHandle<IG4Trajectory> contrib
                                = trajCon->GetTrajectory(*c);
                            ensure("Contributing trajectory exists", contrib);
                        }
                    }
                }
            }
        }
    }

    // Check that there is energy deposited and track length for each hit.
    template<> template<>
    void testHitSegments::test<3> () {
        TestHits::HitChecker checkHits = TestHits::HitChecker(SimG4::Events, 
                                                        SimG4::StartEvents[0],
                                                        SimG4::EndEvents[0]);
        for (TestHits::HitChecker::const_iterator hitIter = checkHits.begin();
                                                  hitIter != checkHits.end();
                                                ++hitIter){
            const IG4HitSegment* hit 
                = dynamic_cast<const IG4HitSegment*>(&hitIter);
            if (hit) {
                ensure_greaterthan("Hit has energy deposit",hit->GetEnergyDeposit(),0);
                ensure_greaterthan("Hit has track length", hit->GetTrackLength(),0);
            }
        }
    }

    // Check that the hit track length is OK.
    template<> template<>
    void testHitSegments::test<4> () {
        int event = 0;
        int longCount = 0;
        // EG: Inhereted this from ND280...
        // We don't set a stepping limit for volumes outside of
        // the TPC, so we can have very long steps (limited by the
        // actual geometry).  The longest single active geometric
        // object in the COMET is about 3 meters long.
        double maxSegmentLength = 300*unit::cm;
        TestHits::HitChecker checkHits = TestHits::HitChecker(SimG4::Events, 
                                                        SimG4::StartEvents[0],
                                                        SimG4::EndEvents[0]);
        for (TestHits::HitChecker::const_iterator hitIter = checkHits.begin();
                                                  hitIter != checkHits.end();
                                                ++hitIter){
            const IG4HitSegment* hit 
                = dynamic_cast<const IG4HitSegment*>(&hitIter);
            if (hit) {
                TVector3 start(hit->GetStartX(),
                               hit->GetStartY(),
                               hit->GetStartZ());
                TVector3 stop(hit->GetStopX(),
                              hit->GetStopY(), 
                              hit->GetStopZ());
                double segmentLength = (stop-start).Mag();
                // Skip comparing segment and track length for short tracks
                bool shortTrack = hit->GetTrackLength() < unit::mm &&
                                  segmentLength < unit::mm;
                // CW: According to the current definition of hit segment (step separate not defined)
                //     it's possible to have a track that leaves the volume first and then comes back and merged as the same hit segment.
                //     In this case, the length between the step point leaving the volume and the step point entering the volume contributes nothing to this hit segment.
                //     However, track length is accumulated only when the step is fully inside the volume. So the segmentLength we calculated here might be larger than track length.
                //if (!shortTrack){
                //    ensure_greaterthan("Track longer than segment length.",
                //                   hit->GetTrackLength(),
                //                   segmentLength - 0.1*unit::mm);
                //}
                if (segmentLength>maxSegmentLength) {
                    ++longCount;
                    TVector3 dir;
                    dir = dir*(1.0/segmentLength);
                    std::cout << "Long Node Volumes in event: "
                        << event << std::endl;
                    std::cout << "  Energy: " 
                        << hit->GetEnergyDeposit() << std::endl;
                    std::cout << "  Length: " 
                        << hit->GetTrackLength() << std::endl;
                    IOADatabase::Get().GeomId().FindNodeAndEnter(start);
                    std::cout << "  Start Name: " 
                        << gGeoManager->GetCurrentNode()->GetName()
                        << std::endl;
                    for (double step=segmentLength/10; 
                            step<0.95*segmentLength; 
                            step += segmentLength/10) {
                        TVector3 pos = start + dir*step;
                        IOADatabase::Get().GeomId().FindNodeAndEnter(pos);
                        std::cout << "  Step " << step << " " 
                            << (gGeoManager
                                    ->GetCurrentNode()->GetName())
                            << std::endl;
                    }
                    IOADatabase::Get().GeomId().FindNodeAndEnter(stop);
                    std::cout << "  Stop Name: " 
                        << gGeoManager->GetCurrentNode()->GetName()
                        << std::endl;
                }
            }
        }
        ensure("No Long Segments", longCount==0);
    }

    // Check that end points of the hit segment are in the same volume.
    template<> template<>
    void testHitSegments::test<5> () {
        TestHits::HitChecker checkHits = TestHits::HitChecker(SimG4::Events, 
                                                        SimG4::StartEvents[0],
                                                        SimG4::EndEvents[0]);
        for (TestHits::HitChecker::const_iterator hitIter = checkHits.begin();
                                                  hitIter != checkHits.end();
                                                ++hitIter){
            const IG4HitSegment* hit 
                = dynamic_cast<const IG4HitSegment*>(&hitIter);
         
            if (hit) {
                TVector3 start(hit->GetStartX(),
                               hit->GetStartY(),
                               hit->GetStartZ()); 
                TVector3 stop(hit->GetStopX(),
                              hit->GetStopY(), 
                              hit->GetStopZ());
                TVector3 dir = (stop-start).Unit();
                double segmentLength = (stop-start).Mag();
                // Skip really short segments.
                if (segmentLength<0.3*unit::mm) continue;
                // Figure out how far to step to get away from roundoff
                // errors.  This needs to be a relatively large step since
                // some of the hits "scrape a wall".
                double safetyStep = std::min(0.1*unit::mm,
                        segmentLength/4);
                TVector3 safeStart = start + safetyStep*dir;
                IOADatabase::Get().GeomId().FindNodeAndEnter(safeStart);
                int startNodeId = gGeoManager->GetCurrentNodeId();
                ensure("Start node id was found",(startNodeId>=0));
                std::string startName
                    = gGeoManager->GetCurrentNode()->GetName();
                TVector3 safeStop = stop - safetyStep*dir;
                IOADatabase::Get().GeomId().FindNodeAndEnter(safeStop);
                int stopNodeId = gGeoManager->GetCurrentNodeId();
                std::string stopName
                    = gGeoManager->GetCurrentNode()->GetName();
                ensure("Stop node id was found",(stopNodeId>=0));
                if (startNodeId != stopNodeId) {
                    std::cout << "Node Id mismatch" << std::endl;
                    std::cout << "  Length " << segmentLength << std::endl;
                    IOADatabase::Get().GeomId().FindNodeAndEnter(start);
                    std::cout << "  Start Name: " 
                        << gGeoManager->GetCurrentNode()->GetName()
                        << std::endl;
                    for (double step=safetyStep; 
                            step<10*safetyStep && step<segmentLength/2; 
                            step += safetyStep) {
                        TVector3 pos = start + dir*step;
                        IOADatabase::Get().GeomId().FindNodeAndEnter(pos);
                        std::cout << "  Step " << step << " " 
                            << (gGeoManager
                                    ->GetCurrentNode()->GetName())
                            << std::endl;
                    }
                    std::cout << "..." << std::endl;
                    for (double step
                            = std::max(segmentLength-10*safetyStep,
                                segmentLength/2);
                            step<segmentLength; 
                            step += safetyStep) {
                        TVector3 pos = start + dir*step;
                        IOADatabase::Get().GeomId().FindNodeAndEnter(pos);
                        std::cout << "  Step " << step << " " 
                            << (gGeoManager
                                    ->GetCurrentNode()->GetName())
                            << std::endl;
                    }
                    IOADatabase::Get().GeomId().FindNodeAndEnter(stop);
                    std::cout << "  Stop Name: " << segmentLength
                        << " " 
                        << gGeoManager->GetCurrentNode()->GetName()
                        << std::endl;
                }
                ensure("Start and stop nodes are the same",
                        (startNodeId == stopNodeId));
            }
        }
    }

    // Check that secondaries create their own hit instances and their
    // trajectories are stored in the trajectory container.
    template<> template<>
    void testHitSegments::test<6> () {
        // Event 8 is a 1 GeV positron with a 50MeV gamma cut.
        ICOMETEvent* event = SimG4::Events[8];

        auto trajs = event->Get<IG4TrajectoryContainer>("truth/G4Trajectories");

        auto hits = event->Get<IG4HitContainer>("truth/g4Hits/TestTestSurface_0");
        ensure("TestSurface has a hit container", hits);

        // Get a list of track IDs which produced hits
        std::list<int> contributors;

        for (IG4VHit* vhit: *hits) {
            int id = vhit->GetPrimaryId();
            if (std::find(contributors.begin(), contributors.end(), id) == contributors.end())
                contributors.push_back(id);
        }

        size_t n_contributors = contributors.size();
        size_t n_trajs = trajs->size();
        ensure("The number of hit contributors is equal to or less than the number of "
                "saved trajectories.",
                n_contributors <= n_trajs);
	ensure_lessthan("Any difference between contributors and saved trajectories "
                        "is below the tolerance level",
			float(n_trajs-n_contributors)/n_trajs, 0.005);
    }
};