#pragma once

#include <desc/utilities/Explaination.hxx>
#include <ndarray/ndarray.hxx>

#if WITH_EIGEN
#include <pybind11/eigen.h>
#else
// #define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
// #include <Python.h>
// #include <numpy/arrayobject.h>
#endif

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

#include <algorithm>
#include <chrono>
#include <exception>
#include <optional>
#include <stdexcept>

namespace py = pybind11;

namespace sd::disc
{

using real_matrix = ndarray<double, 2>;

auto to_numpy(ndarray<double, 2>& x)
{
    return py::array_t<double>({x.extent(0), x.extent(1)}, x.data());
}

auto to_numpy(ndarray<double, 2> const& x)
{
    auto y = py::array_t<double>();
    y.resize({x.extent(0), x.extent(1)});
    auto target = static_cast<double*>(y.request().ptr);
    for (size_t i = 0; i < x.size(); ++i) {
        target[i] = static_cast<double>(x.data()[i]);
    }
    return y;
}

template <typename T = std::chrono::milliseconds, typename Fn>
auto timeit(Fn&& fn)
{
    namespace ch = std::chrono;
    // using namespace std::chrono_literals;

    auto before = ch::high_resolution_clock::now();
    std::forward<Fn>(fn)();
    auto after = ch::high_resolution_clock::now();
    // return (after - before) / 1ms;
    return ch::duration_cast<T>(after - before);
}

template <typename S>
Dataset<S> to_dataset(const py::list& xss)
{
    Dataset<S> out;

    if (xss.size() == 0) return out;
    //
    // if (!py::isinstance<py::list>(xss[0])) throw std::domain_error{"data is not list of
    // lists"};

    itemset<S> buffer;
    out.reserve(xss.size());
    for (const auto& xs : xss) {
        buffer.clear();
        for (const auto& i : xs) {
            buffer.insert(py::cast<size_t>(i));
        }
        // for (const auto& i : xs) {
        // buffer.insert(tr.convert_to_internal(py::cast<size_t>(i))); }
        out.insert(buffer);
        // out.insert(pylist_to_itemset(xs, buffer, tr));
    }
    return out;
}

template <typename S>
Dataset<S> to_dataset(const py::array_t<bool>& x)
{
    if (x.ndim() != 2) throw std::domain_error{"data is not a matrix"};

    size_t n = x.shape()[0];
    size_t m = x.shape()[1];

    itemset<S> t;
    Dataset<S> out;
    out.reserve(n);

    for (size_t i = 0; i < n; ++i) {
        t.clear();
        for (size_t j = 0; j < m; ++j) {
            if (*x.data(i, j)) t.insert(j);
        }
        out.insert(t);
    }

    return out;
}

template <typename S>
Dataset<S> to_dataset(const py::object& x) // BiMap tr
{
    if (py::isinstance<py::list>(x)) {
        return to_dataset<S>(py::cast<py::list>(x));
    }
    if (py::isinstance<py::array>(x)) {
        return to_dataset<S>(py::cast<py::array_t<bool>>(x));
    }
    throw std::runtime_error{"cannot represent the given python object as data"};
}

template <typename S, typename T>
auto to_dataset(Dataset<T> const& xss)
    -> std::conditional_t<std::is_same_v<S, T>, const Dataset<S>&, Dataset<S>>
{
    if constexpr (std::is_same_v<S, T>) {
        return xss;
    } // std::forward<Dataset<S>>(xss); }
    else
    {
        itemset<S> buffer;
        Dataset<S> out;
        out.reserve(xss.size());
        for (const auto & [xs] : xss) {
            buffer.clear();
            sd::for_each_item(xs, [&](size_t i) { buffer.insert(i); });
            out.insert(buffer);
        }
        return out;
    }
}

class ModelInterface : public Explaination
{
public:
    virtual std::vector<size_t> predict(Dataset<tag_sparse> const& x) const = 0;
    virtual std::vector<size_t> predict(Dataset<tag_dense> const& x) const  = 0;

    virtual std::vector<size_t> predict(py::object const& x) const                        = 0;
    virtual real_matrix         predict_probabilities(Dataset<tag_sparse> const& x) const = 0;
    virtual real_matrix         predict_probabilities(Dataset<tag_dense> const& x) const  = 0;
    virtual real_matrix         predict_probabilities(py::object const& x) const          = 0;
    virtual real_matrix predict_log_probabilities(Dataset<tag_sparse> const& x) const     = 0;
    virtual real_matrix predict_log_probabilities(Dataset<tag_dense> const& x) const      = 0;
    virtual real_matrix predict_log_probabilities(py::object const& x) const              = 0;

    virtual double log_likelihood(Dataset<tag_sparse> const& x,
                                  std::vector<size_t> const& y) const                      = 0;
    virtual double log_likelihood(Dataset<tag_dense> const&  x,
                                  std::vector<size_t> const& y) const                      = 0;
    virtual double log_likelihood(py::object const& x, std::vector<size_t> const& y) const = 0;

    // virtual std::vector<size_t> predict(binary_matrix const& x) const                     =
    // 0; virtual std::vector<size_t> predict(sparse_format const& x) const
    // = 0; virtual real_matrix         predict_probabilities(binary_matrix const& x) const
    // = 0; virtual real_matrix         predict_probabilities(sparse_format const& x) const
    // = 0; virtual real_matrix predict_log_probabilities(binary_matrix const& x) const
    // = 0; virtual real_matrix predict_log_probabilities(sparse_format const& x) const
    // = 0; virtual double log_likelihood(binary_matrix const&       x,
    //                               std::vector<size_t> const& y) const                      =
    //                               0;
    // virtual double log_likelihood(sparse_format const&       x,
    //                               std::vector<size_t> const& y) const                      =
    //                               0;

    virtual ~ModelInterface() = default;
};

class Interface : public ModelInterface
{
public:
    virtual void fit(Dataset<tag_sparse> const& x, std::vector<size_t> const& y) = 0;
    virtual void fit(Dataset<tag_dense> const& x, std::vector<size_t> const& y)  = 0;
    virtual void fit(py::object const& x, std::vector<size_t> const& y)          = 0;
    // virtual void fit(binary_matrix const& x, std::vector<size_t> const& y)       = 0;
    // virtual void fit(sparse_format const& x, std::vector<size_t> const& y)       = 0;
};

template <typename Trait, typename BaseClass>
class SingleClass : public BaseClass
{
public:
    using pattern_type = typename Trait::pattern_type;

    virtual void do_fit(Dataset<pattern_type> x) = 0;

    std::vector<size_t> do_predict(Dataset<pattern_type> const& x) const
    {
        return std::vector<size_t>(x.size());
    }
    real_matrix do_predict_probabilities(Dataset<pattern_type> const& x) const
    {
        real_matrix m;
        disc::predict_probabilities(x, component.model, m);
        return m;
        // return disc::predict_probabilities(x, component.model).template cast<double>();
    }
    real_matrix do_predict_log_probabilities(Dataset<pattern_type> const& x) const
    {
        real_matrix m;
        disc::predict_log_probabilities(x, component.model, m);
        return m;
        // return disc::predict_log_probabilities(x, component.model).template cast<double>();
    }
    auto do_log_likelihood(Dataset<pattern_type> const& x) const
    {
        return static_cast<double>(disc::log_likelihood(x, component.model));
    }
    //
    // Interface
    //

    void fit(Dataset<tag_sparse> const& x, std::vector<size_t> const&) override
    {
        do_fit(to_dataset<pattern_type>(x));
    }
    void fit(Dataset<tag_dense> const& x, std::vector<size_t> const&) override
    {
        do_fit(to_dataset<pattern_type>(x));
    }
    // void fit(binary_matrix const& x, std::vector<size_t> const&) override
    // {
    //     do_fit(to_dataset<pattern_type>(x));
    // }
    // void fit(sparse_format const& x, std::vector<size_t> const&) override
    // {
    //     do_fit(to_dataset<pattern_type>(x));
    // }
    void fit(py::object const& x, std::vector<size_t> const&) override
    {
        do_fit(to_dataset<pattern_type>(x));
    }
    std::vector<size_t> predict(Dataset<tag_sparse> const& x) const override
    {
        return do_predict(to_dataset<pattern_type>(x));
    }
    std::vector<size_t> predict(Dataset<tag_dense> const& x) const override
    {
        return do_predict(to_dataset<pattern_type>(x));
    }
    // std::vector<size_t> predict(binary_matrix const& x) const override
    // {
    //     return do_predict(to_dataset<pattern_type>(x));
    // }
    // std::vector<size_t> predict(sparse_format const& x) const override
    // {
    //     return do_predict(to_dataset<pattern_type>(x));
    // }
    std::vector<size_t> predict(py::object const& x) const override
    {
        return do_predict(to_dataset<pattern_type>(x));
    }
    real_matrix predict_probabilities(Dataset<tag_sparse> const& x) const override
    {
        return do_predict_probabilities(to_dataset<pattern_type>(x));
    }
    real_matrix predict_probabilities(Dataset<tag_dense> const& x) const override
    {
        return do_predict_probabilities(to_dataset<pattern_type>(x));
    }
    // real_matrix predict_probabilities(binary_matrix const& x) const override
    // {
    //     return do_predict_probabilities(to_dataset<pattern_type>(x));
    // }
    // real_matrix predict_probabilities(sparse_format const& x) const override
    // {
    //     return do_predict_probabilities(to_dataset<pattern_type>(x));
    // }
    real_matrix predict_probabilities(py::object const& x) const override
    {
        return do_predict_probabilities(to_dataset<pattern_type>(x));
    }
    real_matrix predict_log_probabilities(Dataset<tag_sparse> const& x) const override
    {
        return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    }
    real_matrix predict_log_probabilities(Dataset<tag_dense> const& x) const override
    {
        return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    }
    // real_matrix predict_log_probabilities(binary_matrix const& x) const override
    // {
    //     return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    // }
    // real_matrix predict_log_probabilities(sparse_format const& x) const override
    // {
    //     return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    // }
    real_matrix predict_log_probabilities(py::object const& x) const override
    {
        return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    }

    double log_likelihood(Dataset<tag_sparse> const& x,
                          std::vector<size_t> const&) const override
    {
        return do_log_likelihood(to_dataset<pattern_type>(x));
    }
    double log_likelihood(Dataset<tag_dense> const& x,
                          std::vector<size_t> const&) const override
    {
        return do_log_likelihood(to_dataset<pattern_type>(x));
    }
    // double log_likelihood(binary_matrix const& x, std::vector<size_t> const&) const override
    // {
    //     return do_log_likelihood(to_dataset<pattern_type>(x));
    // }
    // double log_likelihood(sparse_format const& x, std::vector<size_t> const&) const override
    // {
    //     return do_log_likelihood(to_dataset<pattern_type>(x));
    // }
    double log_likelihood(py::object const& x, std::vector<size_t> const&) const override
    {
        return do_log_likelihood(to_dataset<pattern_type>(x));
    }

    //
    // Disable Explanation
    //

    std::vector<size_t> contrasting(std::vector<size_t> const&,
                                    std::vector<size_t> const&) const override
    {
        throw std::runtime_error("A single class cannot contrast itself");
    };
    std::vector<size_t> emerging(std::vector<size_t> const&, std::vector<size_t> const&) const
    {
        throw std::runtime_error("A single class cannot contrast itself");
        // return {};
    }
    std::vector<size_t> common(std::vector<size_t> const&) const
    {
        throw std::runtime_error("A single class cannot contrast itself");
        //    return characteristic(0);
    }
    std::vector<size_t> characteristic(int) const
    {
        // std::vector<size_t> s(set.size());
        // std::iota(s.begin(), s.end(), 0);
        // return s;
        throw std::runtime_error("A single class cannot contrast itself");
    }
    std::vector<size_t> unique(std::vector<size_t> const&) const
    {
        throw std::runtime_error("A single class cannot contrast itself");
    }

protected:
    Component<Trait> component;
};

template <typename Trait, typename BaseClass>
class MultiClass : public BaseClass
{
public:
    using pattern_type = typename Trait::pattern_type;
    using float_type   = typename Trait::float_type;

    virtual void do_fit(Dataset<pattern_type> x, std::vector<size_t> const& y) = 0;

    std::vector<size_t> do_predict(Dataset<pattern_type> const& x) const
    {
        auto                k = composition.data.num_components();
        std::vector<size_t> labels(k, 0);
        for (size_t i = 0; i < k; ++i) {
            labels[i] = label(composition.data.subset(i)[0]);
        }
        auto y = predict_labels(x, composition.models);
        for (auto& l : y) l = labels[l];
        return y;
    }
    real_matrix do_predict_probabilities(Dataset<pattern_type> const& x) const
    {
        real_matrix m;
        disc::predict_probabilities(x, composition.models, m);
        return m;
        // return (disc::predict_probabilities(x, composition.models)).template cast<double>();
    }
    real_matrix do_predict_log_probabilities(Dataset<pattern_type> const& x) const
    {
        real_matrix m;
        disc::predict_log_probabilities(x, composition.models, m);
        return m;
    }
    double do_log_likelihood(Dataset<pattern_type> const& x, std::vector<size_t> const& y) const
    {
        return static_cast<double>(disc::log_likelihood(x, y, composition.models));
    }
    //
    // Interface
    //

    void fit(Dataset<tag_sparse> const& x, std::vector<size_t> const& y) override
    {
        do_fit(to_dataset<pattern_type>(x), y);
    }
    void fit(Dataset<tag_dense> const& x, std::vector<size_t> const& y) override
    {
        do_fit(to_dataset<pattern_type>(x), y);
    }
    // void fit(binary_matrix const& x, std::vector<size_t> const& y) override
    // {
    //     do_fit(to_dataset<pattern_type>(x), y);
    // }
    // void fit(sparse_format const& x, std::vector<size_t> const& y) override
    // {
    //     do_fit(to_dataset<pattern_type>(x), y);
    // }
    void fit(py::object const& x, std::vector<size_t> const& y) override
    {
        do_fit(to_dataset<pattern_type>(x), y);
    }
    std::vector<size_t> predict(Dataset<tag_sparse> const& x) const override
    {
        return do_predict(to_dataset<pattern_type>(x));
    }
    std::vector<size_t> predict(Dataset<tag_dense> const& x) const override
    {
        return do_predict(to_dataset<pattern_type>(x));
    }
    // std::vector<size_t> predict(binary_matrix const& x) const override
    // {
    //     return do_predict(to_dataset<pattern_type>(x));
    // }
    // std::vector<size_t> predict(sparse_format const& x) const override
    // {
    //     return do_predict(to_dataset<pattern_type>(x));
    // }
    std::vector<size_t> predict(py::object const& x) const override
    {
        return do_predict(to_dataset<pattern_type>(x));
    }
    real_matrix predict_probabilities(Dataset<tag_sparse> const& x) const override
    {
        return do_predict_probabilities(to_dataset<pattern_type>(x));
    }
    real_matrix predict_probabilities(Dataset<tag_dense> const& x) const override
    {
        return do_predict_probabilities(to_dataset<pattern_type>(x));
    }
    // real_matrix predict_probabilities(binary_matrix const& x) const override
    // {
    //     return do_predict_probabilities(to_dataset<pattern_type>(x));
    // }
    // real_matrix predict_probabilities(sparse_format const& x) const override
    // {
    //     return do_predict_probabilities(to_dataset<pattern_type>(x));
    // }
    real_matrix predict_probabilities(py::object const& x) const override
    {
        return do_predict_probabilities(to_dataset<pattern_type>(x));
    }
    real_matrix predict_log_probabilities(Dataset<tag_sparse> const& x) const override
    {
        return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    }
    real_matrix predict_log_probabilities(Dataset<tag_dense> const& x) const override
    {
        return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    }
    // real_matrix predict_log_probabilities(binary_matrix const& x) const override
    // {
    //     return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    // }
    // real_matrix predict_log_probabilities(sparse_format const& x) const override
    // {
    //     return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    // }
    real_matrix predict_log_probabilities(py::object const& x) const override
    {
        return do_predict_log_probabilities(to_dataset<pattern_type>(x));
    }
    double log_likelihood(Dataset<tag_sparse> const& x,
                          std::vector<size_t> const& y) const override
    {
        return do_log_likelihood(to_dataset<pattern_type>(x), y);
    }
    double log_likelihood(Dataset<tag_dense> const&  x,
                          std::vector<size_t> const& y) const override
    {
        return do_log_likelihood(to_dataset<pattern_type>(x), y);
    }
    // double log_likelihood(binary_matrix const& x, std::vector<size_t> const& y) const
    // override
    // {
    //     return do_log_likelihood(to_dataset<pattern_type>(x), y);
    // }
    // double log_likelihood(sparse_format const& x, std::vector<size_t> const& y) const
    // override
    // {
    //     return do_log_likelihood(to_dataset<pattern_type>(x), y);
    // }
    double log_likelihood(py::object const& x, std::vector<size_t> const& y) const override
    {
        return do_log_likelihood(to_dataset<pattern_type>(x), y);
    }

protected:
    Composition<Trait> composition;
};

template <typename S, typename T, typename Fn>
void select_dist_type(bool is_relaxed, Fn&& fn)
{
    if (is_relaxed) {
        using trait = Trait<S, T, RelEntDistribution<S, T>>;
        std::forward<Fn>(fn)(trait{});
    }
    else
    {
        using trait = Trait<S, T, MaxEntDistribution<S, T>>;
        std::forward<Fn>(fn)(trait{});
    }
}

// using precise_float_t = long double;

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

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

} // namespace sd::disc
