#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 TH1Params {
        std::string label_x;
        int nbins;
        double low;
        double high;
        std::string label_y;

        static TH1Params 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!",
                        true);
            } else {
                auto params = hist_params[param_key];
                return TH1Params({params["label_x"].as<std::string>(),
                                  params["nbins"].as<int>(),
                                  params["low"].as<double>(),
                                  params["high"].as<double>(),
                                  params["label_y"].as<std::string>()
                                 });

            }
        }
    };

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

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

        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);
            }
        }

        GenContainer *clone_as(const std::string &new_name) {
            return new ContainerTH1<V>(new_name, this->title, this->params);
        }
    };

    struct TH2Params {
        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;

        static TH2Params 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!",
                        true);
            } else {
                auto params = hist_params[param_key];
                return TH2Params({params["label_x"].as<std::string>(),
                                  params["nbins_x"].as<int>(),
                                  params["low_x"].as<double>(),
                                  params["high_x"].as<double>(),
                                  params["label_y"].as<std::string>(),
                                  params["nbins_y"].as<int>(),
                                  params["low_y"].as<double>(),
                                  params["high_y"].as<double>()
                                 });

            }
        }
    };

    template<typename V>
    class ContainerTH2 : public Container<TH2, V> {
    private:
        std::string title;
        TH2Params params;
    public:
        ContainerTH2(const std::string &name, const std::string &title, TH2Params 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());
        }

        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;
        }

        GenContainer *clone_as(const std::string &new_name) {
            return new ContainerTGraph<V>(new_name, this->title);
        }

        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);
        }

        GenContainer *clone_as(const std::string &new_name) {
            return new Vector<V>(new_name);
        }

        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]++;
        }

        GenContainer *clone_as(const std::string &new_name) {
            return new Counter<V>(new_name);
        }
    };


    template<typename T>
    class EfficiencyContainer : public Container<TH1F, std::vector<T>> {
    private:
        std::function<bool(T)> *selector;  // Selects whether object is up for consideration
        std::function<bool(T)> *predicate; // Says whether the object passes the efficiency criteria
        std::function<float(T)> *get_val;  // Returns a floating point value from the object that is actually
        // used in the histogram
        bool write_num_den;
        TH1F num;
        TH1F den;
        TH1F eff;

        TH1Params params;

        void _fill() {
            for (auto &obj : this->value->get_value()) {
                if (selector == nullptr or (*selector)(obj)) {
                    float val = (*get_val)(obj);
                    den.Fill(val);
                    if ((*predicate)(obj)) {
                        num.Fill(val);
                    }
                }
            }
        }

    public:
        EfficiencyContainer(const std::string &name, Value<std::vector<T>> *value, TH1Params params,
                            std::function<bool(T)> *selector, std::function<bool(T)> *predicate,
                            std::function<float(T)> *get_val, bool write_num_den = false)
                : Container<TH1F, std::vector<T>>(name, value),
                  num{(name + "_num").c_str(), (name + "_num").c_str(), params.nbins, params.low, params.high},
                  den{(name + "_den").c_str(), (name + "_den").c_str(), params.nbins, params.low, params.high},
                  eff{name.c_str(), name.c_str(), params.nbins, params.low, params.high},
                  selector(selector), predicate(predicate), get_val(get_val), params(params),
                  write_num_den(write_num_den) {
            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;
        }

        GenContainer *clone_as(const std::string &new_name) {
            return new EfficiencyContainer<T>(new_name, this->value, this->params, selector, predicate, get_val);
        }

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

}
#endif // root_container_hpp