
#include <disc/interfaces/BoostMultiprecision.hxx>

#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>

#include <disc/disc/Desc.hxx>
#include <disc/disc/Disc.hxx>
#include <relaxent/Classifier.hxx>
#include <relaxent/DESCVariants.hxx>
#include <relaxent/IntersectingFactorization.hxx>

#include "Utilities.hxx"
#include <synth/PatternPlanter.hxx>

namespace py = pybind11;

using namespace sd::disc::pytools;
using namespace sd::viva;
using namespace sd::disc;

namespace sd::disc
{

// TODO merge/reuse PyDisc.cpp

template <typename S, typename T, typename Fn>
void select_dist_type(Fn&& fn, bool relaxed = true)
{
    if (relaxed)
    {
        using dist_t = IntersectingFactorsDistribution<S, T, viva::SuperSetcoverWeightFunction>;
        using trait  = Trait<S, T, dist_t>;
        std::forward<Fn>(fn)(trait{});
    }
    else
    {
        using trait = Trait<S, T, MEDistribution<S, T>>;
        std::forward<Fn>(fn)(trait{});
    }
}

template <typename S, typename Fn>
void select_real_type(bool is_precise, Fn&& fn, bool relaxed = true)
{
    if (is_precise)
    {
        select_dist_type<S, float_hp_t>(std::forward<Fn>(fn));
    }
    else
    {
        select_dist_type<S, double>(std::forward<Fn>(fn));
    }
}

template <typename Fn>
void build_trait(bool is_sparse, bool is_precise, Fn&& fn, bool relaxed = true)
{
    if (is_sparse)
    {
        select_real_type<tag_sparse>(is_precise, std::forward<Fn>(fn));
    }
    else
    {
        select_real_type<tag_dense>(is_precise, std::forward<Fn>(fn));
    }
}

} // namespace sd::disc

template <typename trait_type>
auto reap_discover_patternset_(const py::list& dataset, size_t min_support)
{
    using namespace sd::disc;

    PatternsetResult<trait_type> co;

    insert_pylists_into_data(dataset, co.data);

    DecompositionSettings cfg;
    cfg.min_support = min_support;
    cfg.use_bic     = true;

    sd::disc::discover_patterns_generic(co, cfg, IUCNDesc{});

    return translate_to_pydict(co);
}

auto reap_discover_patternset(const py::list& dataset,
                              size_t          min_support,
                              bool            is_sparse  = false,
                              bool            is_precise = false)
{
    py::dict r;
    sd::disc::build_trait(is_sparse, is_precise, [&](auto trait) {
        r = reap_discover_patternset_<decltype(trait)>(dataset, min_support);
    });
    return r;
}

template <typename trait_type>
auto reap_characterize_partitions_(const py::list& dataset,
                                   const py::list& labels,
                                   size_t          min_support)
{
    using namespace sd::disc;

    Composition<trait_type> co;
    insert_pylists_into_data(co.data, dataset, labels);

    DecompositionSettings cfg;
    cfg.min_support = min_support;
    cfg.use_bic     = true;

    initialize_composition(co, cfg);
    auto initial_encoding = co.encoding;
    discover_patterns_generic(co, cfg, IUCNDesc{});
    co.initial_encoding = initial_encoding;

    return translate_to_pydict(co);
}

auto reap_characterize_partitions(const py::list& dataset,
                                  const py::list& labels,
                                  size_t          min_support,
                                  bool            is_sparse  = false,
                                  bool            is_precise = false)
{
    py::dict r;
    sd::disc::build_trait(is_sparse, is_precise, [&](auto trait) {
        r = reap_characterize_partitions_<decltype(trait)>(dataset, labels, min_support);
    });
    return r;
}

template <typename trait_type>
auto reap_discover_composition_(const py::list& dataset, size_t min_support, double alpha)
{
    using namespace sd::disc;

    Composition<trait_type> co;

    insert_pylists_into_data(dataset, co.data);

    DecompositionSettings cfg;
    cfg.alpha       = alpha;
    cfg.min_support = min_support;
    cfg.use_bic     = true;

    initialize_composition(co, cfg);
    auto initial_encoding = co.encoding;

    auto pm = [](auto& c, const auto& g) { discover_patterns_generic(c, g, IUCNDesc{}); };
    decompose_until_summary_converges(co, cfg, pm, sd::EmptyCallback{});
    co.initial_encoding = initial_encoding;

    return translate_to_pydict(co);
}

auto reap_discover_composition(const py::list& dataset,
                               size_t          min_support,
                               double          alpha,
                               bool            is_sparse  = false,
                               bool            is_precise = false)
{
    py::dict r;
    sd::disc::build_trait(is_sparse, is_precise, [&](auto trait) {
        r = reap_discover_composition_<decltype(trait)>(dataset, min_support, alpha);
    });
    return r;
}

template <typename Trait>
py::dict reap_cross_validate_impl(const py::list& dataset,
                                  const py::list& labels,
                                  size_t          min_support,
                                  bool            is_relaxed)
{
    using namespace sd::disc;

    Composition<Trait> co;
    insert_pylists_into_data(co.data, dataset, labels);
    simplify_labels(co.data);

    DecompositionSettings cfg;
    cfg.min_support = min_support;
    cfg.use_bic     = true;

    const auto tt = is_relaxed ? TrainingType::UCDESC1 : TrainingType::CDESC1;

    auto [train_err, test_err, train_time, test_time] =
        sd::disc::cross_validate<Trait>(co.data, 0.8, 10, 5, tt, cfg);

    py::dict r;
    using Stats = sd::IncrementalDescription<double>;
    Stats stats;
    stats            = Stats(train_err);
    r["train_l1"]    = stats.mean();
    r["train_l1_sd"] = stats.sd();

    stats           = Stats(test_err);
    r["test_l1"]    = stats.mean();
    r["test_l1_sd"] = stats.sd();

    stats.reset();
    for (auto ms : train_time)
        stats += ms.count();
    r["train_ms"]    = stats.mean();
    r["train_ms_sd"] = stats.sd();
    stats.reset();

    for (auto ms : train_time)
        stats += ms.count();
    r["test_ms"]    = stats.mean();
    r["test_ms_sd"] = stats.sd();

    return r;
}

auto reap_cross_validate(const py::list& dataset,
                         const py::list& labels,
                         size_t          min_support,
                         bool            is_sparse,
                         bool            is_precise,
                         bool            is_relaxed)
{
    py::dict result;
    build_trait(is_sparse,
                is_precise,
                [&](auto trait) {
                    result = reap_cross_validate_impl<decltype(trait)>(
                        dataset, labels, min_support, is_relaxed);
                },
                is_relaxed);
    return result;
}

template <typename S, typename T>
struct PyMEDist
{
    using float_type   = T;
    using pattern_type = S;

    PyMEDist(size_t dim, size_t max_factor_size = 8, size_t max_factor_width = 12)
        : dist(dim, 1, max_factor_size, max_factor_width)
    {
        sd::disc::estimate_model(dist);
    }

    void insert(float_type label, const py::list& t)
    {
        dist.insert(label, pylist_to_itemset(t, buf), true);
    }

    void insert_batch(py::array_t<double>& ys, const py::list& ts)
    {
        if (static_cast<size_t>(ys.size()) != ts.size())
        {
            throw std::runtime_error("arguments of 'insert_batch' are of different length");
        }
        py::buffer_info nfo = ys.request();
        auto y = sd::slice<const double>(static_cast<const double*>(nfo.ptr), ys.size());
        for (size_t i = 0; i < ts.size(); ++i)
        {
            dist.insert(y[i], pylist_to_itemset(ts[i], buf), false);
        }
        sd::disc::estimate_model(dist);
    }

    float_type infer_generalized_itemset(const py::list& t) const
    {
        thread_local sd::disc::itemset<S> buf;
        return dist.expected_generalized_frequency(pylist_to_itemset(t, buf));
    }

    float_type infer(const py::list& t) const
    {
        thread_local sd::disc::itemset<S> buf;
        return dist.expected_frequency(pylist_to_itemset(t, buf));
    }

    static std::string type_name()
    {
        std::string name = "REDist_";
        name += storage_type_to_str<S>();
        name += "_";
        name += float_storage_type_to_str<T>();
        return name;
    }

private:
    using Weight = sd::viva::SuperSetcoverWeightFunction;
    sd::disc::IntersectingFactorsDistribution<S, T, Weight> dist;
    sd::disc::itemset<S>                                    buf;
};

template <typename Trait>
void declare_py_distribution(py::module& m)
{
    using Class = PyMEDist<typename Trait::pattern_type, typename Trait::float_type>;

    py::class_<Class>(m, Class::type_name().c_str())
        .def(py::init<size_t>())
        .def(py::init<size_t, size_t>())
        .def(py::init<size_t, size_t, size_t>())
        .def("infer", &Class::infer)
        .def("infer_generalized_itemset", &Class::infer_generalized_itemset)
        .def("insert_batch", &Class::insert_batch)
        .def("insert", &Class::insert);
}

namespace sd::disc
{

template <typename S = tag_sparse, typename T = double>
auto sample_transactions(bool   discrete_sampler = false,
                         size_t num_rows         = 32,
                         size_t dim              = 5,
                         size_t num_patterns     = 5,
                         size_t min_width        = 2,
                         size_t max_width        = 5,
                         size_t seed             = 5,
                         double noise            = 0.01)
{
    Dataset<S>           data;
    LabeledDataset<T, S> patterns;

    std::default_random_engine rng{seed};

    if (discrete_sampler)
    {
        generate_dataset(
            num_rows, dim, num_patterns, data, patterns, noise, rng, 0, min_width, max_width);
    }
    else
    {
        generate_dataset_discrete_sampler(
            num_rows, dim, num_patterns, data, patterns, noise, rng, 0, min_width, max_width);
    }
    return std::make_pair(data, patterns);
}

pybind11::dict sample_transactions_py(bool discrete_sampler = false,
                                      // bool        sparse           = false,
                                      size_t num_rows     = 32,
                                      size_t dim          = 5,
                                      size_t num_patterns = 5,
                                      size_t min_width    = 2,
                                      size_t max_width    = 5,
                                      size_t seed         = 5,
                                      double noise        = 0.01)
{

    auto [x, s] = sample_transactions<tag_sparse, double>(
        discrete_sampler, num_rows, dim, num_patterns, min_width, max_width, seed, noise);

    pybind11::list xy, sy;

    pytools::insert_data_into_pylist(x, xy);
    pytools::insert_data_into_pylist(s, sy);

    pybind11::dict r;
    r["data"]        = xy;
    r["patterns"]    = sy;
    r["frequencies"] = s.template col<0>();
    return r;
}

} // namespace sd::disc

PYBIND11_MODULE(reap, m)
{
    using namespace pybind11::literals;

    m.doc() = R"doc(    
              REAP: Relaxed maximum Entropy Accelerated Patternset miner
              ----------------------------------------------------------
              Discovers factors for the relaxed maximum entropy distribution in terms of informative patterns
            )doc";

    m.def("discover_patternset",
          &::reap_discover_patternset,
          "A function that discovers significant patterns for a dataset using the relaxed "
          "maximum "
          "entropy distribution",
          "dataset"_a,
          "min_support"_a = 1,
          "is_sparse"_a   = false,
          "is_precise"_a  = false);
    m.def("characterize_partitions",
          &reap_characterize_partitions,
          "A function that discovers significant patterns and characterizes multiple, given a "
          "partition using the relaxed maximum entropy distribution",
          "dataset"_a,
          "partition_label"_a,
          "min_support"_a = 1,
          "is_sparse"_a   = false,
          "is_precise"_a  = false);
    m.def("discover_composition",
          &reap_discover_composition,
          "A function that discovers differently distributed partitions of the dataset as well "
          "as significant patterns and characterizes the partitions using the relaxed maximum "
          "entropy distribution",
          "dataset"_a,
          "min_support"_a = 1,
          "alpha"_a       = 0.05,
          "is_sparse"_a   = false,
          "is_precise"_a  = false);

    m.def("cross_validate",
          &reap_cross_validate,
          "10-fold 80/20 cross validation using relaxed maxent distribution",
          "dataset"_a,
          "partition_label"_a,
          "min_support"_a = 1,
          "alpha"_a       = 0.05,
          "is_sparse"_a   = false,
          "is_precise"_a  = false);

    m.def("sample_transactions",
          &sd::disc::sample_transactions_py,
          "Samples patterns and transactional datasets",
          "discrete_sampler"_a = true,
          "num_rows"_a         = 32,
          "dim"_a              = 5,
          "num_patterns"_a     = 5,
          "min_width"_a        = 2,
          "max_width"_a        = 5,
          "seed"_a             = 5,
          "noise"_a            = 0.01);

    for (int i = 0; i <= 1; ++i)
    {
        for (int j = 0; j <= 1; ++j)
        {
            sd::disc::build_trait(
                i, j, [&](auto trait) { declare_py_distribution<decltype(trait)>(m); });
        }
    }

    m.attr("__version__") = "dev";
}
