#include "ExampleAnaElemTool.h"

#include "G4Event.hh"
#include "G4THitsCollection.hh"

#include "DD4hep/Detector.h"
#include "DD4hep/Plugins.h"
#include "DDG4/Geant4Converter.h"
#include "DDG4/Geant4Mapping.h"
#include "DDG4/Geant4HitCollection.h"
#include "DDG4/Geant4Data.h"
#include "DDG4/Geant4Hits.h"

DECLARE_COMPONENT(ExampleAnaElemTool)

void
ExampleAnaElemTool::BeginOfRunAction(const G4Run*) {
    G4cout << "Begin Run of detector simultion..." << G4endl;
}

void
ExampleAnaElemTool::EndOfRunAction(const G4Run*) {
    G4cout << "End Run of detector simultion..." << G4endl;
}

void
ExampleAnaElemTool::BeginOfEventAction(const G4Event* anEvent) {
    msg() << "Event " << anEvent->GetEventID() << endmsg;
}

void
ExampleAnaElemTool::EndOfEventAction(const G4Event* anEvent) {

    // save all data

    // create collections.
    auto trackercols = m_trackerCol.createAndPut();
    auto vxdcols = m_VXDCol.createAndPut();

    // readout defined in DD4hep
    auto lcdd = &(dd4hep::Detector::getInstance());
    auto allReadouts = lcdd->readouts();

    for (auto& readout : allReadouts) {
        info() << "Readout " << readout.first << endmsg;
    }

    // retrieve the hit collections
    G4HCofThisEvent* collections = anEvent->GetHCofThisEvent();
    if (!collections) {
        warning() << "No collections found. " << endmsg;
        return;
    }
    int Ncol = collections->GetNumberOfCollections();
    for (int icol = 0; icol < Ncol; ++icol) {
        G4VHitsCollection* collect = collections->GetHC(icol);
        if (!collect) {
            warning() << "Collection iCol " << icol << " is missing" << endmsg;
            continue;
        }
        size_t nhits = collect->GetSize();
        info() << "Collection " << collect->GetName()
               << " #" << icol
               << " has " << nhits << " hits."
               << endmsg;
        if (nhits==0) {
            // just skip this collection.
            continue;
        }
        // There are different types (new and old)

        dd4hep::sim::Geant4HitCollection* coll = dynamic_cast<dd4hep::sim::Geant4HitCollection*>(collect);
        if (coll) {
            info() << " cast to dd4hep::sim::Geant4HitCollection. " << endmsg;
            for(size_t i=0; i<nhits; ++i)   {

                dd4hep::sim::Geant4HitData* h = coll->hit(i);

                dd4hep::sim::Geant4Tracker::Hit* trk_hit = dynamic_cast<dd4hep::sim::Geant4Tracker::Hit*>(h);
                if ( 0 != trk_hit )   {
                    dd4hep::sim::Geant4HitData::Contribution& t = trk_hit->truth;
                    int trackID = t.trackID;
                    // t.trackID = m_truth->particleID(trackID);
                }
                // Geant4Calorimeter::Hit* cal_hit = dynamic_cast<Geant4Calorimeter::Hit*>(h);
                // if ( 0 != cal_hit )   {
                //     Geant4HitData::Contributions& c = cal_hit->truth;
                //     for(Geant4HitData::Contributions::iterator j=c.begin(); j!=c.end(); ++j)  {
                //         Geant4HitData::Contribution& t = *j;
                //         int trackID = t.trackID;
                //         // t.trackID = m_truth->particleID(trackID);
                //     }
                // }
            }
            continue;
        }

        typedef G4THitsCollection<dd4hep::sim::Geant4Hit> HitCollection;
        HitCollection* coll2 = dynamic_cast<HitCollection*>(collect);

        if (coll2) {
            info() << " cast to G4THitsCollection<dd4hep::sim::Geant4Hit>. " << endmsg;

            for(size_t i=0; i<nhits; ++i)   {
                dd4hep::sim::Geant4Hit* h = dynamic_cast<dd4hep::sim::Geant4Hit*>(coll2->GetHit(i));
                if (!h) {
                    warning() << "Failed to cast to dd4hep::sim::Geant4Hit. " << endmsg;
                    continue;
                }

                dd4hep::sim::Geant4TrackerHit* trk_hit = dynamic_cast<dd4hep::sim::Geant4TrackerHit*>(h);
                if (trk_hit) {
                    info() << " cast to dd4hep::sim::Geant4TrackerHit. " << endmsg;

                    plcio::SimTrackerHitCollection* colptr = nullptr;

                    if (collect->GetName() == "VXDCollection") {
                        colptr = vxdcols;
                    } else {
                        colptr = trackercols;
                    }

                    // auto edm_trk_hit = trackercols->create();
                    auto edm_trk_hit = (*colptr)->create();

                    // Refer to: ./DDG4/lcio/LCIOConversions.cpp
                    edm_trk_hit->setCellID0((trk_hit->cellID >>    0         ) & 0xFFFFFFFF);
                    edm_trk_hit->setCellID1((trk_hit->cellID >> sizeof(int)*8) & 0xFFFFFFFF);
                    edm_trk_hit->setEDep(trk_hit->energyDeposit/CLHEP::GeV);
                    edm_trk_hit->setTime(trk_hit->truth.time/CLHEP::ns);
                    edm_trk_hit->setPathLength(trk_hit->length/CLHEP::mm);
                    // lc_hit->setMCParticle(lc_mcp);
                    double pos[3] = {trk_hit->position.x()/CLHEP::mm,
                                     trk_hit->position.y()/CLHEP::mm,
                                     trk_hit->position.z()/CLHEP::mm};
                    edm_trk_hit->setPosition(plcio::DoubleThree(pos));

                    float mom[3] = {trk_hit->momentum.x()/CLHEP::GeV,
                                    trk_hit->momentum.y()/CLHEP::GeV,
                                    trk_hit->momentum.z()/CLHEP::GeV};
                    edm_trk_hit->setMomentum(plcio::FloatThree(mom));
                }

                dd4hep::sim::Geant4CalorimeterHit* cal_hit = dynamic_cast<dd4hep::sim::Geant4CalorimeterHit*>(h);
                if (cal_hit) {
                    info() << " cast to dd4hep::sim::Geant4CalorimeterHit. " << endmsg;
                }

            }

            continue;
        }

        warning() << "Failed to convert to collection "
                  << collect->GetName()
                  << endmsg;
        
    }
}

void
ExampleAnaElemTool::PreUserTrackingAction(const G4Track*) {

}

void
ExampleAnaElemTool::PostUserTrackingAction(const G4Track*) {

}

void
ExampleAnaElemTool::UserSteppingAction(const G4Step*) {

}

StatusCode
ExampleAnaElemTool::initialize() {
    StatusCode sc;

    return sc;
}

StatusCode
ExampleAnaElemTool::finalize() {
    StatusCode sc;

    return sc;
}