filval/root/include/root_container.hpp

#ifndef root_container_hpp
#define root_container_hpp

#include <iostream>
#include <utility>
#include <map>

#include "TROOT.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TGraph.h"
#include "TH1.h"
#include "TH2.h"

#include "TMVA/Factory.h"
#include "TMVA/DataLoader.h"
#include "TMVA/DataSetInfo.h"

#include "container.hpp"

namespace fv_root_util {
    using namespace fv;
    /**
     * Save a TObject. The TObject will typically be a Histogram or Graph object,
     * but can really be any TObject. The SaveOption can be used to specify how to
     * save the file.
     */
    void save_as(TObject *container, const std::string &fname, const SaveOption &option = SaveOption::PNG) {

        auto save_img = [](TObject *container, const std::string &fname) {
            TCanvas *c1 = new TCanvas("c1");
            container->Draw();
            c1->Draw();
            c1->SaveAs(fname.c_str());
            delete c1;
        };

        auto save_bin = [](TObject *container) {
            INFO("Saving object: " << container->GetName() << " into file " << gDirectory->GetName());
            container->Write(container->GetName(), TObject::kOverwrite);
        };

        switch (option) {
            case SaveOption::PNG:
                save_img(container, fname + ".png");
                break;
            case SaveOption::PDF:
                save_img(container, fname + ".pdf");
                break;
            case SaveOption::ROOT:
                save_bin(container);
                break;
            default:
                break;
        }
    }


    /**
     * Saves an STL container into a ROOT file. ROOT knows how to serialize STL
     * containers, but it needs the *name* of the type of the container, eg.
     * std::map<int,int> to be able to do this. In order to generate this name at
     * run-time, the fv_util::get_type_name function uses RTTI to get type info
     * and use it to look up the proper name.
     *
     * For nexted containers, it is necessary to generate the CLING dictionaries
     * for each type at compile time to enable serialization. To do this, add the
     * type definition into the LinkDef.hpp header file.
     */
    void save_as_stl(void *container, const std::string &type_name,
                     const std::string &obj_name,
                     const SaveOption &option = SaveOption::PNG) {
        switch (option) {
            case SaveOption::PNG:
                INFO("Cannot save STL container " << type_name << " as png");
                break;
            case SaveOption::PDF:
                INFO("Cannot save STL container " << type_name << " as pdf");
                break;
            case SaveOption::ROOT:
                /* DEBUG("Writing object \"" << obj_name << "\" of type \"" << type_name << "\"\n"); */
                gDirectory->WriteObjectAny(container, type_name.c_str(), obj_name.c_str());
                break;
            default:
                break;
        }
    }
}

namespace fv_root {
    using namespace fv;

    struct THParams {
        std::string label_x;
        int nbins_x;
        double low_x;
        double high_x;
        std::string label_y;
        int nbins_y;
        double low_y;
        double high_y;

        THParams(int nbins_x, double low_x, double high_x, int nbins_y=0, double low_y=0, double high_y=0)
          :nbins_x(nbins_x), low_x(low_x), high_x(high_x), nbins_y(nbins_y), low_y(low_y), high_y(high_y){}

        THParams(const std::string label_x, int nbins_x, double low_x, double high_x,
                const std::string label_y = "", int nbins_y=0, double low_y=0, double high_y=0)
          :label_x(label_x), nbins_x(nbins_x), low_x(low_x), high_x(high_x),
           label_y(label_y), nbins_y(nbins_y), low_y(low_y), high_y(high_y){}

        THParams() {
            label_x = "";
            nbins_x = 20;
            low_x = 0;
            high_x = 1;
            label_y = "";
            nbins_y = 20;
            low_y = 0;
            high_y = 1;
        }

        static THParams lookup(const std::string &&param_key) {
            auto hist_params = fv_util::the_config->get("hist-params");
            if (!hist_params[param_key]) {
                CRITICAL("Key \"" << param_key
                                  << "\" does not exist under hist-params in supplied config file. Add it!");
                std::string empty{""};
                return {};
            } else {
                auto params = hist_params[param_key];
                return {params["label_x"].as<std::string>(""),
                        params["nbins_x"].as<int>(20),
                        params["low_x"].as<double>(0),
                        params["high_x"].as<double>(1),
                        params["label_y"].as<std::string>(""),
                        params["nbins_y"].as<int>(20),
                        params["low_y"].as<double>(0),
                        params["high_y"].as<double>(1)
                       };

            }
        }
    };

    template<typename V>
    class ContainerTH1 : public Container<TH1, V> {
    protected:
        std::string title;
        THParams params;

    public:
        ContainerTH1(const std::string &name, const std::string &title, const THParams &params)
                : Container<TH1, V>(name), title(title), params(params) {
            this->container = new TH1D(this->get_name().c_str(), this->title.c_str(),
                                       params.nbins_x, params.low_x, params.high_x);
            this->container->SetXTitle(params.label_x.c_str());
            this->container->SetYTitle(params.label_y.c_str());
        }

        ContainerTH1(const std::string &name, const THParams &params)
                : ContainerTH1<V>(name, name, params) { }

        void fill(const V &v) {
            this->container->Fill(v);
        }

        void fill_weight(const V &v, const float &weight) {
            this->container->Fill(v, weight);
        }

        void fill_many(const std::vector<V> &vs) {
            for (const V &v : vs) {
                this->container->Fill(v);
            }
        }

        void save_as(const std::string &fname, const SaveOption &option = SaveOption::PNG) {
            if (this->get_container() == nullptr) {
                WARNING("Container " << this->get_name() << " never filled, cowardly refusing to save");
            } else {
                fv_root_util::save_as(this->get_container(), fname, option);
            }
        }
    };

    template<typename V>
    class ContainerTH2 : public Container<TH2, V> {
    private:
        std::string title;
        THParams params;
    public:
        ContainerTH2(const std::string &name, const std::string &title, THParams params)
                : Container<TH2, V>(name), title(title), params(params) {
            this->container = new TH2D(this->get_name().c_str(), this->title.c_str(),
                                       params.nbins_x, params.low_x, params.high_x,
                                       params.nbins_y, params.low_y, params.high_y);
            this->container->SetXTitle(params.label_x.c_str());
            this->container->SetYTitle(params.label_y.c_str());
        }

        ContainerTH2(const std::string &name, const THParams &params)
                : ContainerTH2<V>(name, name, params) { }

        void fill(const V& x, const V& y) {
            this->container->Fill(x, y);
        }

        void fill_weight(const V& x, const V& y, const V& weight) {
            this->container->Fill(x, y, weight);
        }

        void save_as(const std::string &fname, const SaveOption &option = SaveOption::PNG) {
            fv_root_util::save_as(this->get_container(), fname, option);
        }

    };

    template<typename V>
    class ContainerTGraph : public Container<TGraph, V> {
    private:
        std::vector<V> x_data;
        std::vector<V> y_data;
        std::string title;
        bool data_modified;


    public:
        ContainerTGraph(const std::string &name, const std::string &title)
                : Container<TGraph, V>(name), data_modified(false) {
            this->container = new TGraph();
        }

        TGraph *get_container() {
            if (data_modified) {
                delete this->container;
                this->container = new TGraph(x_data.size(), x_data.data(), y_data.data());
                this->container->SetName(this->get_name().c_str());
                this->container->SetTitle(title.c_str());
                data_modified = false;
            }
            return this->container;
        }

        void save_as(const std::string &fname, const SaveOption &option = SaveOption::PNG) {
            fv_root_util::save_as(get_container(), fname, option);
        }

        void fill(const V& x, const V& y) {
            x_data.push_back(x);
            y_data.push_back(y);
            data_modified = true;
        }
    };

    template<typename V>
    class Vector : public Container<std::vector<V>, V> {
    public:
        Vector(const std::string &name)
                : Container<std::vector<V>, V>(name) {
            this->container = new std::vector<V>;
        }

        void fill(const V& v) {
            this->container->push_back(v);
        }

        void save_as(const std::string &fname, const SaveOption &option = SaveOption::PNG) {
            std::string type_name = "std::vector<" + fv_util::get_type_name(typeid(V)) + ">";
            fv_root_util::save_as_stl(this->get_container(), type_name, this->get_name(), option);
        }
    };

/**
 * A Counter that keeps a mapping of the number of occurances of each input
 * value.
 */
    template<typename V>
    class Counter : public Container<std::map<V, size_t>, V> {
    public:
        Counter(const std::string &name)
                : Container<std::map<V, size_t>, V>(name) {
            this->container = new std::map<V, size_t>;
        }

        void save_as(const std::string &fname, const SaveOption &option = SaveOption::PNG) {
            std::string type_name = "std::map<" + fv_util::get_type_name(typeid(V)) + ",int>";
            fv_root_util::save_as_stl(this->get_container(), type_name, this->get_name(), option);
        }

        void fill(const V& v) {
            (*this->container)[v]++;
        }
    };


    template<typename V>
    class EfficiencyContainer : public Container<TH1F, V> {
    private:
        TH1F num;
        TH1F den;
        TH1F eff;

        THParams params;

    public:
        EfficiencyContainer(const std::string &name, THParams params)
                : Container<TH1F, V>(name),
                  num{(name + "_num").c_str(), (name + "_num").c_str(), params.nbins_x, params.low_x, params.high_x},
                  den{(name + "_den").c_str(), (name + "_den").c_str(), params.nbins_x, params.low_x, params.high_x},
                  eff{name.c_str(), name.c_str(), params.nbins_x, params.low_x, params.high_x}, params(params) {
            num.SetXTitle(params.label_x.c_str());
            num.SetYTitle(params.label_y.c_str());
            den.SetXTitle(params.label_x.c_str());
            den.SetYTitle(params.label_y.c_str());
            eff.SetXTitle(params.label_x.c_str());
            eff.SetYTitle(params.label_y.c_str());
            this->container = &eff;
        }

        TH1F *get_container() {
            eff.Sumw2();
            eff.Divide(&num, &den, 1, 1, "B");
            return this->container;
        }

        void save_as(const std::string &fname, const SaveOption &option = SaveOption::PNG) {
            fv_root_util::save_as(this->get_container(), fname, option);
            fv_root_util::save_as(&num, fname, option);
            fv_root_util::save_as(&den, fname, option);
        }

        void fill(const V& v, bool pass) {
            den.Fill(v);
            if (pass) num.Fill(v);
        }
    };

    template<typename V>
    class EfficiencyContainer2D : public Container<TH2F, V> {
    private:
        TH2F num;
        TH2F den;
        TH2F eff;

        THParams params;

    public:
        EfficiencyContainer2D(const std::string &name, THParams params)
                : Container<TH2F, V>(name),
                  num{(name + "_num").c_str(), (name + "_num").c_str(),
                      params.nbins_x, params.low_x, params.high_x, params.nbins_y, params.low_y, params.high_y},
                  den{(name + "_den").c_str(), (name + "_den").c_str(),
                      params.nbins_x, params.low_x, params.high_x, params.nbins_y, params.low_y, params.high_y},
                  eff{name.c_str(), name.c_str(),
                      params.nbins_x, params.low_x, params.high_x, params.nbins_y, params.low_y, params.high_y},
                  params(params) {
            num.SetXTitle(params.label_x.c_str());
            num.SetYTitle(params.label_y.c_str());
            den.SetXTitle(params.label_x.c_str());
            den.SetYTitle(params.label_y.c_str());
            eff.SetXTitle(params.label_x.c_str());
            eff.SetYTitle(params.label_y.c_str());
            this->container = &eff;
        }

        TH2F *get_container() {
            eff.Sumw2();
            eff.Divide(&num, &den, 1, 1, "B");
            return this->container;
        }

        void save_as(const std::string &fname, const SaveOption &option = SaveOption::PNG) {
            fv_root_util::save_as(this->get_container(), fname, option);
            fv_root_util::save_as(&num, fname, option);
            fv_root_util::save_as(&den, fname, option);
        }

        void fill(const V& x,const V& y, bool pass) {
            den.Fill(x, y);
            if (pass) num.Fill(x, y);
        }
    };

}
#endif // root_container_hpp