IUnpacker.cxx

#include "IUnpacker.hxx"
#include "args.hxx"
#include "ICOMETEvent.hxx"
#include "IG4Trajectory.hxx"
#include "IG4VHit.hxx"
#include "IHitSelection.hxx"
#include "IDigitContainer.hxx"
#include <TFile.h>
#include <TTree.h>

using namespace COMET;

IUnpacker::IUnpacker() {

}

IUnpacker::~IUnpacker() {
    // Free memory associated with args::Flags
    // allocated in RegisterArgs.
    for (auto& it: fFlags) {
        delete it.second.arg;
        if (it.second.oaEvent_obj_ptr)
            delete it.second.oaEvent_obj_ptr;
    }
}

void IUnpacker::SetOutputFile(TFile* f) {
    fOut = f;
}

void IUnpacker::AddArgFlag(const std::string& flag_name,
                           const std::string& flag_description)
{
    flag_info info;
    info.description = flag_description;
    info.tree = nullptr;
    info.oaEvent_obj_ptr = nullptr;
    fFlags[flag_name] = std::move(info);
    fOrder.push_back(flag_name);
}

void IUnpacker::RegisterArgs(args::Group& group) {
    /// Register all our args in the order they were added.
    for (const auto flag_name: fOrder) {
        auto& info = fFlags[flag_name];
        // Prefix will go before all arg commands, e.g. for 'cdc-hits',
        // the result will be '--[prefix]cdc-hits'.
        std::string prefix = ""; //"save-";

        // Assign to the command-line flag its name and description
        // This will register the command-line argument, which will
        // appear when we execute the program, and show its description
        // when we specify the -h or --help flag.
        info.arg = new args::Flag(group, flag_name,
                info.description, // (description)
                {prefix + flag_name} // (actual command-line argument)
                );
    }
}

bool IUnpacker::CheckAllFlagsRegistered() {
    for (const auto it: fFlags) {
        // To check if a function was registered, we can check the 
        // value of the container name. If it is empty, then the user did
        // not call IUnpacker::RegisterFunction.
        const auto container_name = it.second.container_name; 
        if (container_name == "") {
            std::cout << "No unpacking function registered for flag '" <<
                it.first << "'.\n";
            return false;
        }
    }
    return true;
}

void IUnpacker::ProcessEvent(const ICOMETEvent& event) {
    /// We need to do some work to get the right containers
    /// out of the event and then pass them to the unpacking functions.
    /// The main difficulty here, is to evaluate the type of 
    /// container that we're accessing, so that we can properly
    /// obtain its elements.

    /// The four main container types (IG4TrajectoryContainer,
    /// IG4HitContainer, IHitSelection and IDigitContainer) are actually not
    /// compatible with each other. We cannot use them without explicitly
    /// knowing their type. Hence, this method needs to infer their type from
    /// their name, which makes it probably the least flexible method
    /// in the whole program. 
    
    /// If new containers are introduced that don't respect the current
    /// naming conventions, this method will not be able to infer their type
    /// and someone will need to come in and hard-code them in here.

    // Iterate through all flags and apply functions to containers
    for (const auto it: fFlags) {
        const auto info = it.second;

        // Break if the flag wasn't activated (tree is nullptr)
        if (info.tree == nullptr) continue;

        // Container name
        const auto name = info.container_name;
        TTree* tree = info.tree;
        // Unpacking function
        const auto f = fFunctions[name];

        if (name.find("G4Trajectories") != std::string::npos) {
            // We're dealing with the trajectory container
            IHandle<IG4TrajectoryContainer> tc =
               event.Get<IG4TrajectoryContainer>(name); 

            if (tc) {
                for (const auto traj: *tc) {
                    // Set pointer to oaEvent object to the address of
                    // object inside the container
                    *info.oaEvent_obj_ptr = &(traj.second);
                    // Execute the bound function, which will do the extraction
                    // and place the results inside the object that we gave ROOT
                    // when we branched
                    bool save = f();
                    // And now we can happily fill the tree.
                    if (save)
                        tree->Fill();
                }
            }
        } 
        else if (name.find("g4Hits/") != std::string::npos) {
            // We're dealing with a (truth) hit container
            IHandle<IG4HitContainer> hc =
               event.Get<IG4HitContainer>(name); 

            if (hc) {
                for (const IG4VHit* hit: *hc) {
                    *info.oaEvent_obj_ptr = hit;
                    bool save = f();
                    if (save)
                        tree->Fill();
                }
            }
        }
        else if (name.find("hits/") != std::string::npos) {
            // We're dealing with a hit selection
            IHandle<IHitSelection> hs =
               event.Get<IHitSelection>(name); 

            if (hs) {
                for (const IHandle<IHit> hit: *hs) {
                    *info.oaEvent_obj_ptr = &(*hit);
                    bool save = f();
                    if (save)
                        tree->Fill();
                }
            }
        }
        else if (name.find("digits/") != std::string::npos) {
            // We're dealing with a digits container
            IHandle<IDigitContainer> dc =
               event.Get<IDigitContainer>(name); 

            if (dc) {
                for (const IDigit* digit: *dc) {
                    *info.oaEvent_obj_ptr = digit;
                    bool save = f();
                    if (save)
                        tree->Fill();
                }
            }
        }
    }
}

void IUnpacker::Write() {
    for (const auto it: fFlags) {
        const auto info = it.second;
        if (info.tree)
            info.tree->Write();
    }
}