#pragma once

#include <elisa/ComplementDist.hxx>
#include <elisa/Elisa.hxx>

namespace sd::disc
{

template <typename Trait, typename Score>
auto bin_pv_summary(Component<Trait> const& c,
                    const Config&           cfg,
                    Score&&                 fn,
                    double                  initial = 0)
{
    ndarray<typename Trait::float_type, 2> conf(initial, c.summary.size(), 1);

    auto pr = make_independence_dist(c, cfg);

#pragma omp parallel for schedule(dynamic, 1) // reduction(push_back : ranking)
    for (size_t j = 0; j < c.summary.size(); ++j) {
        if (is_singleton(c.summary.point(j))) continue;
        auto q     = c.frequency[j];
        conf(j, 0) = fn(pr, q, c.summary.point(j));
    }
    return conf;
}

template <typename Trait, typename Score>
auto bin_pv_summary_c(Component<Trait> const& c,
                      const Config&           cfg,
                      Score&&                 fn,
                      double                  initial = 0)
{
    ndarray<typename Trait::float_type, 2> conf(initial, c.summary.size(), 1);

#pragma omp parallel for schedule(dynamic, 1) // reduction(push_back : ranking)
    for (size_t j = 0; j < c.summary.size(); ++j) {
        if (is_singleton(c.summary.point(j))) continue;
        auto pr    = make_complement_dist(c, cfg, j);
        auto q     = c.frequency[j];
        conf(j, 0) = fn(pr, q, c.summary.point(j));
    }
    return conf;
}
template <typename Trait, typename Score>
auto bin_pv_summary(Composition<Trait> const& c,
                    Config const&             cfg,
                    Score&&                   fn,
                    double                    initial = 0)
{
    ndarray<typename Trait::float_type, 2> conf(
        initial, c.summary.size(), c.data.num_components());

    auto pr = make_independence_dist(c, cfg);

#pragma omp parallel for schedule(dynamic, 1)
    for (size_t j = 0; j < c.summary.size(); ++j) {
        if (is_singleton(c.summary.point(j))) continue;
        for (size_t i = 0; i < c.models.size(); ++i) {
            if (!c.assignment[i].contains(j)) continue;
            conf(j, i) = fn(pr[i], c.frequency(j, i), c.summary.point(j), i);
        }
    }
    return conf;
}
template <typename Trait, typename Score>
auto bin_pv_summary_c(Composition<Trait> const& c,
                      Config const&             cfg,
                      Score&&                   fn,
                      double                    initial = 0)
{
    ndarray<typename Trait::float_type, 2> conf(
        initial, c.summary.size(), c.data.num_components());

#pragma omp parallel for schedule(dynamic, 1)
    for (size_t j = 0; j < c.summary.size(); ++j) {
        if (is_singleton(c.summary.point(j))) continue;
        for (size_t i = 0; i < c.models.size(); ++i) {
            if (!c.assignment[i].contains(j)) continue;
            auto pr    = make_complement_dist(c, i, cfg, j);
            conf(j, i) = fn(pr, c.frequency(j, i), c.summary.point(j), i);
        }
    }
    return conf;
}

template <typename Trait>
auto binomial_pv_mat(Composition<Trait> const& c, const Config& cfg)
{
    return bin_pv_summary(c,
                          cfg,
                          [&](const auto& pr, const auto& q, const auto& x, size_t i) {
                              return binomial_pvalue(
                                  pr.expectation(x), q, c.data.subset(i).size());
                          },
                          1);
}
// template <typename Trait>
// auto binomial_pv_mat(Component<Trait> const& c, const Config& cfg)
// {
//     return bin_pv_summary(c, cfg, [&](const auto& pr, const auto& q, const auto& x) {
//         return binomial_pvalue(pr.expectation(x), q, c.data.size());
//     }, 1);
// }
// template <typename Trait>
// auto binomial_pv_mat_c(Composition<Trait> const& c, const Config& cfg)
// {
//     return bin_pv_summary_c(
//         c, cfg, [&](const auto& pr, const auto& q, const auto& x, size_t i) {
//             return binomial_pvalue(pr.expectation(x), q, c.data.subset(i).size());
//         }, 1);
// }
// template <typename Trait>
// auto binomial_pv_mat_c(Component<Trait> const& c, const Config& cfg)
// {
//     return bin_pv_summary_c(c, cfg, [&](const auto& pr, const auto& q, const auto& x) {
//         return binomial_pvalue(pr.expectation(x), q, c.data.size());
//     }, 1);
// }
// template <typename Trait>
// auto binomial_logpv_mat(Composition<Trait> const& c, const Config& cfg)
// {
//     return bin_pv_summary(c, cfg, [&](const auto& pr, const auto& q, const auto& x, size_t i)
//     {
//         return binomial_log_pvalue(pr.expectation(x), q, c.data.subset(i).size());
//     }, 1);
// }
// template <typename Trait>
// auto binomial_logpv_mat(Component<Trait> const& c, const Config& cfg)
// {
//     return bin_pv_summary(c, cfg, [&](const auto& pr, const auto& q, const auto& x) {
//         return binomial_log_pvalue(pr.expectation(x), q, c.data.size());
//     });
// }
// template <typename Trait>
// auto binomial_logpv_mat_c(Composition<Trait> const& c, const Config& cfg)
// {
//     return bin_pv_summary_c(
//         c, cfg, [&](const auto& pr, const auto& q, const auto& x, size_t i) {
//             return binomial_log_pvalue(pr.expectation(x), q, c.data.subset(i).size());
//         });
// }
// template <typename Trait>
// auto binomial_logpv_mat_c(Component<Trait> const& c, const Config& cfg)
// {
//     return bin_pv_summary_c(c, cfg, [&](const auto& pr, const auto& q, const auto& x) {
//         return binomial_log_pvalue(pr.expectation(x), q, c.data.size());
//     });
// }

} // namespace sd::disc