#pragma once

#include <desc/Component.hxx>
#include <desc/Composition.hxx>
#include <desc/utilities/BoostMultiprecision.hxx>
#include <ndarray/ndarray.hxx>

namespace sd::disc
{

struct Explaination
{
public:
    AssignmentMatrix                 A;
    std::vector<std::vector<size_t>> set;
    std::chrono::milliseconds        elapsed;
    Config                           cfg;
    ndarray<double, 2>               confidence_matrix;
    ndarray<double, 2>               frequency_matrix;

    virtual ~Explaination() {}
    virtual const std::vector<std::vector<size_t>>& patterns() const { return set; }
    virtual std::vector<size_t>                     contrasting(const std::vector<size_t>& I,
                                                                const std::vector<size_t>& J) const
    {
        itemset<tag_sparse> a, b;
        for (auto i : I) a.insert(A.at(i));
        for (auto j : J) b.insert(A.at(j));

        std::vector<size_t> c;
        std::set_symmetric_difference(
            a.begin(), a.end(), b.begin(), b.end(), std::back_inserter(c));
        return c;
    }
    virtual std::vector<size_t> emerging(const std::vector<size_t>& I,
                                         const std::vector<size_t>& J) const
    {
        itemset<tag_sparse> a, b;
        for (auto i : I) a.insert(A.at(i));
        for (auto j : J) setminus(a, A.at(j));
        return {a.container.begin(), a.container.end()};
    }
    virtual std::vector<size_t> characteristic(size_t i) const { return A.at(i).container; }
    virtual std::vector<size_t> common(const std::vector<size_t>& I) const
    {
        if (I.empty()) return {};

        itemset<tag_sparse> a;

        a.insert(A.at(I.front()));
        for (auto i : I) intersection(A.at(i), a);
        return {a.container.begin(), a.container.end()};
    }
    virtual std::vector<size_t> unique(const std::vector<size_t>& I) const
    {
        itemset<tag_sparse> a;
        for (auto i : I) a.insert(A.at(i));
        for (size_t i = 0; i < A.size(); ++i)
        {
            if (std::find(I.begin(), I.end(), i) == I.end()) { setminus(a, A.at(i)); }
        }
        return {a.container.begin(), a.container.end()};
    }
    // std::vector<std::vector<size_t>> patterns(const std::vector<size_t>& I) const
    // {
    //     std::vector<std::vector<size_t>> result(I.size());
    //     for (auto i : I)
    //     {
    //         if (!set.at(i).empty()) result.push_back(set[i]);
    //     };
    //     return result;
    // }

    template <typename T>
    void extract_explanation(Composition<T> const& c)
    {
        A.resize(c.data.num_components());
        for (auto& a : A) a.clear();

        size_t num_patterns = 0;
        for (size_t i = 0, k = 0; i < c.summary.size(); ++i)
        {
            num_patterns += !is_singleton(c.summary.point(i));
        }
        confidence_matrix.resize(num_patterns, c.data.num_components());
        frequency_matrix.resize(num_patterns, c.data.num_components());

        if (num_patterns == 0) return;
        set.reserve(num_patterns);

        for (size_t i = 0, k = 0; i < c.summary.size(); ++i)
        {
            const auto& x   = c.summary.point(i);
            auto        cnt = count(x);
            if (cnt > 1)
            {
                auto& next = set.emplace_back();
                next.reserve(cnt);
                for_each_item(x, [&](size_t index) { next.push_back(index); });

                for (size_t j = 0; j < c.data.num_components(); ++j)
                {
                    if (c.confidence(i, j) > 0) A[j].insert(k);
                    confidence_matrix(k, j) = (double)c.confidence(i, j);
                    frequency_matrix(k, j)  = (double)c.frequency(i, j);
                }
                ++k;
            }
        }
        // // this->frequency = copy_to_numpy(rm_singletons_mat(c.frequency, c.data.dim));

        // for (size_t j = 0; j < c.data.num_components(); ++j)
        // {
        //     const auto& aj = c.assignment[j];
        //     for (size_t i = 0; i < aj.container.size(); ++i)
        //     {
        //         if (aj.container[i] >= c.data.dim) A.at(j).insert(aj.container[i] -
        //         c.data.dim);
        //     }
        // }

        // // confidence_matrix.clear();
        // confidence_matrix.resize(c.summary.size() - c.data.dim, c.data.num_components());

        // for (size_t i = c.data.dim; i < c.confidence.extent(0); ++i)
        //     for (size_t j = 0; j < c.confidence.extent(1); ++j)
        //     {
        //         confidence_matrix(i - c.data.dim, j) = (double)c.confidence(i, j);
        //     }
    }
    template <typename T>
    void extract_explanation(Component<T> const& c)
    {
        A.resize(1);
        A[0].clear();

        size_t num_patterns = 0;
        for (size_t i = 0, k = 0; i < c.summary.size(); ++i)
        {
            num_patterns += !is_singleton(c.summary.point(i));
        }
        confidence_matrix.resize(num_patterns, 1);
        frequency_matrix.resize(num_patterns, 1);

        if (num_patterns == 0) return;
        set.reserve(num_patterns);

        for (size_t i = 0, k = 0; i < c.summary.size(); ++i)
        {
            const auto& x   = c.summary.point(i);
            auto        cnt = count(x);
            if (cnt > 1)
            {
                auto& next = set.emplace_back();
                next.reserve(cnt);
                for_each_item(x, [&](size_t index) { next.push_back(index); });

                if (c.confidence[i] > 0) A[0].insert(k);
                confidence_matrix(k, 0) = (double)c.confidence[i];
                frequency_matrix(k, 0)  = (double)c.frequency[i];
                ++k;
            }
        }
    }
    //     template <typename T>
    //     void extract_explanation(Component<T> const& c)
    //     {
    //         A.resize(1);
    //         A[0].clear();

    //         if (c.summary.size() > c.data.dim)
    //         {
    //             set.reserve(c.summary.size() - c.data.dim);

    //             for (size_t i = c.data.dim; i < c.summary.size(); ++i)
    //             {
    //                 if (count(c.summary.point(i)) > 0)
    //                 {
    //                     auto& next = set.emplace_back();
    //                     next.reserve(count(c.summary.point(i)));
    //                     for_each_item(c.summary.point(i),
    //                                   [&](size_t index) { next.push_back(index); });
    //                 }
    //             }
    //             // this->frequency.resize({set.size()});
    //             // auto target = static_cast<double*>(this->frequency.request().ptr);
    //             // std::copy(c.summary.template col<0>().begin() + c.data.dim,
    //             //           c.summary.template col<0>().end(),
    //             //           target);
    //             A.back().container.resize(c.summary.size() - c.data.dim, 1);
    //         }

    //         // confidence_matrix.clear();
    //         confidence_matrix.resize(c.summary.size() - c.data.dim, 1);

    //         for (size_t i = c.data.dim; i < c.confidence.size(); ++i)
    //         {
    //             confidence_matrix(i - c.data.dim) = (double)c.confidence[i];
    //         }
    //     }
    };

    template <typename T, typename D>
    void predict_labels(Dataset<T> const& x, std::vector<D> const& pr, std::vector<size_t>& y)
    {
        y.resize(x.size());
        using float_type = std::decay_t<decltype(pr[0].expectation(x.point(0)))>;
        const auto k     = pr.size();
        for (size_t i = 0; i < x.size(); ++i)
        {
            std::pair<size_t, float_type> y_i = {};
            for (size_t j = 0; j < k; ++j)
            {
                auto p = pr[j].expectation(x.point(i));
                if (y_i.second < p) { y_i = {j, p}; }
            }
            y[i] = y_i.first;
        }
    }
    template <typename T, typename D>
    std::vector<size_t> predict_labels(Dataset<T> const& x, std::vector<D> const& pr)
    {
        std::vector<size_t> y(x.size());
        predict_labels(x, pr, y);
        return y;
    }

    template <typename T, typename D, typename Matrix>
    void predict_probabilities(Dataset<T> const& x, std::vector<D> const& pr, Matrix& y)
    {
        // y.clear();
        y.resize(x.size(), pr.size());
        using float_type = std::decay_t<decltype(y(0, 0))>;

#pragma omp parallel for schedule(dynamic, 1) collapse(2)
        for (size_t i = 0; i < x.size(); ++i)
        {
            for (size_t j = 0; j < pr.size(); ++j)
            {
                y(i, j) = (float_type)pr[j].expectation(x.point(i));
            }
        }
    }
    template <typename T, typename D, typename Matrix>
    void predict_log_probabilities(Dataset<T> const& x, std::vector<D> const& pr, Matrix& y)
    {
        // y.clear();
        y.resize(x.size(), pr.size());
#pragma omp parallel for schedule(dynamic, 1) collapse(2)
        for (size_t i = 0; i < x.size(); ++i)
        {
            for (size_t j = 0; j < pr.size(); ++j)
            {
                using float_type = std::decay_t<decltype(y(0, 0))>;
                y(i, j)          = (float_type)pr[j].log_expectation(x.point(i));
            }
        }
    }
    template <typename T, typename D, typename Matrix>
    void predict_probabilities(Dataset<T> const& x, D const& pr, Matrix& y)
    {
        // y.clear();
        y.resize(x.size(), 1);
        using float_type = std::decay_t<decltype(y(0, 0))>;

#pragma omp parallel for schedule(dynamic, 1)
        for (size_t i = 0; i < x.size(); ++i)
        {
            y(i, 0) = (float_type)pr.expectation(x.point(i));
        }
    }

    template <typename T, typename D, typename Matrix>
    void predict_log_probabilities(Dataset<T> const& x, D const& pr, Matrix& y)
    {
        // y.clear();
        y.resize(x.size(), 1);
        using float_type = std::decay_t<decltype(y(0, 0))>;

#pragma omp parallel for schedule(dynamic, 1)
        for (size_t i = 0; i < x.size(); ++i)
        {
            y(i, 0) = (float_type)pr.log_expectation(x.point(i));
        }
    }

    template <typename T, typename D>
    auto log_likelihood(Dataset<T> const& x, D const& pr)
    {
        using float_type = std::decay_t<decltype(pr.expectation(x.point(0)))>;
        float_type l     = 0;
#pragma omp parallel for schedule(dynamic, 1)
        for (size_t i = 0; i < x.size(); ++i) { l += pr.log_expectation(x.point(i)); }
        return l;
    }
    template <typename T, typename D>
    auto
    log_likelihood(Dataset<T> const& x, std::vector<size_t> const& y, std::vector<D> const& pr)
    {
        using float_type = std::decay_t<decltype(pr[0].log_expectation(x.point(0)))>;
        float_type l     = 0;
#pragma omp parallel for schedule(dynamic, 1)
        for (size_t i = 0; i < x.size(); ++i) { l += pr[y[i]].log_expectation(x.point(i)); }
        return l;
    }

    template <typename Trait, typename Matrix, typename Fn>
    void calc_pvalues_indep(Component<Trait> const& c, const Config& cfg, Matrix& m, Fn&& fn)
    {
        m.resize(c.summary.size() - c.data.dim, 1);
        auto pr          = make_independence_dist(c, cfg);
        using float_type = std::decay_t<decltype(m(0, 0))>;

#pragma omp parallel for schedule(dynamic, 1) // reduction(push_back : ranking)
        for (size_t i = c.data.dim; i < c.summary.size(); ++i)
        {
            m(i - c.data.dim, 0) = (float_type)fn(
                pr.expectation(c.summary.point(i)), c.frequency[i], c.data.size());
        }
    }
    template <typename Trait, typename Matrix, typename Fn>
    void calc_pvalues_indep(Composition<Trait> const& c, Config const& cfg, Matrix& m, Fn&& fn)
    {
        auto pr          = make_independence_dist(c, cfg);
        auto d           = c.data.dim;
        using float_type = std::decay_t<decltype(m(0, 0))>;

        m.resize(c.summary.size() - c.data.dim, c.data.num_components());

#pragma omp parallel for schedule(dynamic, 1) collapse(2)
        for (size_t i = d; i < c.summary.size(); ++i)
            for (size_t j = 0; j < c.data.num_components(); ++j)
            {
                m(i - d, j) = (float_type)fn(pr[j].expectation(c.summary.point(i)),
                                             c.frequency(i, j),
                                             c.data.subset(j).size());
            }
    }

} // namespace sd::disc