filval/root/include/root_dataset.hpp

/**
 * @file
 * @author  Caleb Fangmeier <caleb@fangmeier.tech>
 * @version 0.1
 *
 * @section LICENSE
 *
 *
 * MIT License
 *
 * Copyright (c) 2017 Caleb Fangmeier
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy
 * of this software and associated documentation files (the "Software"), to deal
 * in the Software without restriction, including without limitation the rights
 * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 * copies of the Software, and to permit persons to whom the Software is
 * furnished to do so, subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 * SOFTWARE.
 */
#ifndef root_dataset_h
#define root_dataset_h

#include <string>
#include <tuple>

#include "TChain.h"

#include "filval.hpp"
#include "root_filval.hpp"


namespace fv_root {
    using namespace std;
    using namespace fv;
    using namespace fv_util;

    template<typename TREE_CLASS>
    class TreeDataSet : public DataSet {
    private:
        // Maps filenames to data label, eg. "TTTT", or "TTZ"
        std::map<std::string, std::string> input_labels;
        // Maps filenames to data category. Either "signal" or "background"
        std::map<std::string, std::string> input_categories;
        std::string output_filename;
        std::map<std::string, DataFileDescriptor> file_descriptors;
        TFile *output_file;
        TREE_CLASS *tree_obj;
        long next_entry;
        long nentries;

        int get_events() {
            if (this->max_events) return std::min(this->max_events, nentries);
            else return nentries;
        }

        fv_util::DataFileDescriptor &get_current_file() {
            TFile *file = tree_obj->fChain->GetCurrentFile();
            std::string filename = file->GetName();
            return file_descriptors.at(filename);
        }

        bool load_next() {
            if (next_entry >= get_events()) return false;
            tree_obj->LoadTree(next_entry);
            tree_obj->GetEntry(next_entry);
            ++next_entry;
            return true;
        }


        /* void save_event_count(){ */
        /*     std::map<std::string,int> event_counts; */
        /*     std::map<std::string,float> xsecs; */
        /*     string fname, label; */
        /*     for(auto& p : input_labels){ */
        /*         std::tie(fname, label) = p; */
        /*         TFile f(fname.c_str()); */
        /*         TH1D* count = (TH1D*)f.Get("Count"); */
        /*         event_counts[label] = (int)count->GetBinContent(1); */

        /*         TTree* tree = (TTree*)f.Get("tree"); */
        /*         TBranch* b = tree->GetBranch("xsec"); */
        /*         float xsec; */
        /*         b->SetAddress(&xsec); */
        /*         b->GetEntry(1); */
        /*         xsecs[label] = xsec; */
        /*     } */
        /*     output_file->cd(); */
        /*     gDirectory->WriteObjectAny(&event_counts, "std::map<std::string,int>", "_event_counts"); */
        /*     gDirectory->WriteObjectAny(&xsecs, "std::map<std::string,float>", "_xsecs"); */
        /* } */

    public:

        TreeDataSet(const std::string &output_filename, const std::vector<DataFileDescriptor> &dfds,
                    const std::string &tree_name)
                : DataSet(),
                  output_filename(output_filename),
                  next_entry(0) {
            TChain *chain = new TChain(tree_name.c_str());
            int cnt = 0;
            for (const DataFileDescriptor &dfd : dfds) {
                file_descriptors[dfd.filename] = dfd;
                /* file_descriptors.insert(std::make_pair(dfd.filename, dfd)); */
                /* file_descriptors.emplace(std::make_pair(dfd.filename, dfd)); */
                file_descriptors[dfd.filename].file_number = cnt++;
                chain->Add(dfd.filename.c_str());
                input_labels[dfd.filename] = dfd.label;
                input_categories[dfd.filename] = dfd.category;
            }
            tree_obj = new TREE_CLASS(chain);
            nentries = tree_obj->fChain->GetEntries();
            output_file = TFile::Open(output_filename.c_str(), "RECREATE");
            set_branch_status("*", false);
        }

        ~TreeDataSet() {
            /* save_event_count_and_xsection(); */
            output_file->Close();
        }

        int get_current_event() {
            return next_entry - 1;
        }

        const std::string &get_current_event_category() const {
            TFile *file = tree_obj->fChain->GetCurrentFile();
            std::string filename = file->GetName();
            return input_categories.at(filename);
        }

        const std::string &get_current_event_label() const {
            TFile *file = tree_obj->fChain->GetCurrentFile();
            std::string filename = file->GetName();
            return input_labels.at(filename);
        }

        template<typename T>
        Value<T> *track_branch(const std::string &bname) {
            TBranch *branch = tree_obj->fChain->GetBranch(bname.c_str());
            if (branch == nullptr) {
                CRITICAL("Branch: " << bname << " does not exist in input tree.");
            }
            T *bref = (T *) branch->GetAddress();
            set_branch_status(bname, true);
            INFO("Registering branch \"" << bname
                                         << "\" with address " << bref
                                         << " and type " << typeid(bref).name());
            return new ObservedValue<T>(bname, bref);
        }

        template<typename T>
        Value<T *> *track_branch_ptr(const std::string &bname) {
            TBranch *branch = tree_obj->fChain->GetBranch(bname.c_str());
            if (branch == nullptr) {
                CRITICAL("Branch: " << bname << " does not exist in input tree.");
            }
            T *bref = (T *) branch->GetAddress();
            set_branch_status(bname, true);
            INFO("Registering pointer branch \"" << bname
                                                 << "\" with address " << bref
                                                 << " and type " << typeid(bref).name());
            return new PointerValue<T>(bname, bref);
        }

        void set_branch_status(const std::string &bname, bool status) {
            unsigned int found = 0;
            tree_obj->fChain->SetBranchStatus(bname.c_str(), true, &found);
            if (!found) {
                DataFileDescriptor &dfd = get_current_file();
                CRITICAL("Branch: " << bname << " does not exist in input tree of file "
                                    << dfd.filename);
            }

        }

        template<typename T>
        Value<T> *track_branch_obj(const std::string &bname) {
            TBranch *branch = tree_obj->fChain->GetBranch(bname.c_str());
            if (branch == nullptr) {
                CRITICAL("Branch: " << bname << " does not exist in input tree.");
            }
            T **bref = (T **) branch->GetAddress();
            set_branch_status(bname, true);
            INFO("Registering object branch \"" << bname
                                                << "\" with address " << bref
                                                << " and type " << typeid(bref).name());
            return new ObjectValue<T>(bname, bref);
        }

        void save_config() {
            if (the_config == nullptr) return;
            TObjString config(the_config->as_string().c_str());
            output_file->cd();
            config.Write("_config");
        }

        void save_all() {
            output_file->cd();

            // Save the value names for each container to enable looking up
            // what values are plotted
            std::map<string, string> value_lookup = this->get_container_name_value_map();
            gDirectory->WriteObjectAny(&value_lookup, "std::map<std::string,std::string>", "_value_lookup");

            for (auto container : containers) {
                container.second->save_as("outfile", SaveOption::ROOT);
            }
            save_config();
        }
    };
}
#endif // root_dataset_h