#pragma once

#include <desc/Desc.hxx>

namespace sd::disc
{

template <typename Trait>
auto make_independence_dist(Component<Trait> const& c, Config const& cfg)
{
    auto pr = make_distribution(c, cfg);
    for (size_t i = 0; i < c.summary.size(); ++i) {
        const auto& x = c.summary.point(i);
        if (is_singleton(x)) {
            pr.insert_singleton(c.frequency[i], x, false);
        }
    }
    estimate_model(pr);
    return pr;
}

template <typename Trait>
auto make_independence_dist(Composition<Trait> const& c, Config const& cfg)
{
    std::decay_t<decltype(c.models)> prs(c.models.size(), make_distribution(c, cfg));

#pragma omp parallel for schedule(dynamic, 1)
    for (size_t i = 0; i < c.models.size(); ++i) {
        for (size_t j = 0; j < c.summary.size(); ++j) {
            if (is_singleton(c.summary.point(j))) {
                prs[i].insert_singleton(c.frequency(j, i), c.summary.point(j), false);
            }
        }
        estimate_model(prs[i]);
    }

    return prs;
}

template <typename Trait>
auto make_complement_dist(Component<Trait> const& c, Config const& cfg, size_t except)
{
    auto pr = make_distribution(c, cfg);
    for (size_t j = 0; j < c.summary.size(); ++j) {
        if (is_singleton(c.summary.point(j)) && except != j)
            pr.insert_singleton(c.summary.label(j), c.summary.point(j));
    }
    estimate_model(pr);

    for (size_t j = 0; j < c.summary.size(); ++j) {
        if (!is_singleton(c.summary.point(j)) && except != j &&
            pr.is_allowed(c.summary.point(j)))
            pr.insert(c.frequency[j], c.summary.point(j), true);
    }
    return pr;
}

template <typename Trait>
auto make_complement_dist(Composition<Trait> const& c,
                          size_t                    component,
                          Config const&             cfg,
                          size_t                    except)
{
    auto pr = make_distribution(c, cfg);
    for (size_t j = 0; j < c.summary.size(); ++j) {
        if (is_singleton(c.summary.point(j)) && except != j)
            pr.insert_singleton(c.summary.label(j), c.summary.point(j));
    }
    for_each_item(c.assignment[component], [&](size_t j) {
        if (!is_singleton(c.summary.point(j)) && except != j &&
            pr.is_allowed(c.summary.point(j)))
            pr.insert(c.summary.label(j), c.summary.point(j), false);
    });

    estimate_model(pr);
    return pr;
}
}