#include <algorithm>
#include <iostream>
#include <tut.h>

#include "tutRunSimulation.hxx"

#include <TGeoManager.h>

#include <HEPUnits.hxx>
#include <ICOMETEvent.hxx>
#include <IG4VHit.hxx>
#include <IG4Trajectory.hxx>
#include <IG4PrimaryVertex.hxx>

#include "IG4HitSegment.hxx"

using namespace COMET;

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


    // Declare the test
    typedef test_group<basePersistencyManager>::object testPersistencyManager;
    test_group<basePersistencyManager> groupPersistencyManager("PersistencyManager");

    /// Check that all trajectories associated with hits are saved
    template<> template<>
    void testPersistencyManager::test<1> () {
        // 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);

        for (IG4VHit* vhit: *hits) {
            auto hit = dynamic_cast<IG4HitSegment*>(vhit);
            if (!hit) continue;

            // Verify that the hit's trajectory exists in the trajectory container
            auto traj_it = trajs->find(hit->GetPrimaryId());
            ensure("Trajectory of HitSegment exists in trajectory container",
                    traj_it != trajs->end());
        }
    }

    /// Check that all trajectories are saved in a "saveAllTraj" run
    template<> template<>
    void testPersistencyManager::test<2>() {
        // Sample 9 has "/db/set/saveAllTraj true" on a lead surface
        // Event 8 is a 1 GeV positron
        ICOMETEvent* event = SimG4::Events[SimG4::StartEvents[9] + 8];

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

        // To determine whether all trajectories were saved, we check that
        // trajectory IDs are all present between the smallest one (1) and 
        // the largest one.
        std::vector<int> track_ids;
        for (auto traj: *trajs) {
            track_ids.push_back(traj.first);
        }

        // Find max ID
        int max_id = 0;
        for (int id: track_ids) {
            if (id > max_id)
                max_id = id;
        }

        // Ensure that every integer between 1 and max_id is present in track_ids
        for (int i=1; i<max_id; ++i) {
            ensure("All tracks are present in the trajectory container "
                    "in a 'saveAllTraj' run",
                    std::find(track_ids.begin(), track_ids.end(), i) != track_ids.end());
        }
    }

    /// Check that the primary we persisted corresponds to the initial
    /// conditions we set in tutRunSimulation.
    template<> template<>
    void testPersistencyManager::test<3>() {
        // Event 8 is a 1 GeV positron
        ICOMETEvent* event = SimG4::Events[8];
        auto prims = event->Get<IG4PrimaryVertexContainer>("truth/G4PrimVertex00");
        ensure("Event has a primary vertex container", prims);
        ensure("Event only has one primary vertex", prims->size() == 1);

        auto& prim = (*prims)[0];
        int pid = prim.GetPrimaryParticles()[0].GetPDGCode();
        ensure("Primary particle is a positron", 
                pid == -11);

        double E = prim.GetPrimaryParticles()[0].GetMomentum().E();
        ensure("Primary particle has 1 GeV energy",
                 fabs(E - 1000 * unit::MeV) < 10 * unit::eV);
    }
};