#pragma once

#include <disc/disc/Composition.hxx>
#include <disc/disc/Encoding.hxx>
#include <disc/disc/Settings.hxx>

namespace sd
{
namespace disc
{
// MEDistribution<Dist...>
template <typename C, typename Dist, typename Candidate, typename F>
auto assignment_score(
    const C& c, size_t c_index, const Dist& pr, Candidate& x, F fr, bool use_bic)
{
    size_t support = fr * c.data.subset(c_index).size();
    auto gain =
        support == 0 ? 0 : support * std::log2(fr / pr.expected_frequency(x));

    if (use_bic)
    {
        gain -= additional_cost_bic(x, c.data.size(), pr);
    }
    else
    {
        gain -= additional_cost_mdl(support, pr) + constant_mdl_cost(c, x);
    }

    return gain;
}

template <typename Trait, typename Pattern, typename F>
bool is_candidate_assignment_good(const Composition<Trait>& c,
                                  const Pattern&            x,
                                  const size_t              c_index,
                                  const F                   fr,
                                  const bool                bic)
{
    if (fr == 0 || !c.models[c_index].is_itemset_allowed(x))
    {
        return false;
    }
    else
    {
        return assignment_score(c, c_index, c.models[c_index], x, fr, bic) > 0;
    }
}

template <typename Trait, typename Pattern, typename Float>
bool find_assignment_if(Composition<Trait>& c,
                        const size_t        x_index,
                        const Pattern&      x,
                        const size_t        c_index,
                        const Float&        fr,
                        const bool          bic,
                        const bool          estimate_coefficients = false)
{
    if (is_candidate_assignment_good(c, x, c_index, fr, bic))
    {
        c.assignment[c_index].insert(x_index);
        c.models[c_index].insert(fr, x, estimate_coefficients);
        
        return true;
    }
    else
    {
        return false;
    }
}

template <typename Trait, typename Candidate>
bool find_assignment_impl(Composition<Trait>& c, const Candidate& candidate, bool bic)
{
    using float_type      = typename Trait::float_type;
    const auto overall_fr = static_cast<float_type>(candidate.support) / c.data.size();
    return find_assignment_if(c, c.summary.size(), candidate.pattern, 0, overall_fr, bic, true);
}

template <typename Trait, typename Candidate, typename Masks>
size_t find_assignment_impl(Composition<Trait>& c,
                            const Candidate&    candidate,
                            const Masks&        masks,
                            bool                bic)
{
    using float_type = typename Trait::float_type;
    assert(masks.size() == c.assignment.size());
    size_t assignment_count = 0;

    for (size_t i = 0; i < c.assignment.size(); ++i)
    {
        auto s = size_of_intersection(candidate.row_ids, masks[i]);
        auto q = static_cast<float_type>(s) / c.data.subset(i).size();

        if (find_assignment_if(c, c.summary.size(), candidate.pattern, i, q, bic, true))
        {
            assignment_count += 1;
        }
    }
    return assignment_count;
}

template <typename Trait, typename Candidate, typename Masks>
void insert_pattern_to_summary(Composition<Trait>& c,
                               const Candidate&    candidate,
                               const Masks&        masks)
{
    using float_type = typename Trait::float_type;

    assert(c.frequency.extent(1) > 0);

    const auto glob_frequency = static_cast<float_type>(candidate.support) / c.data.size();

    c.summary.insert(glob_frequency, candidate.pattern);
    c.frequency.push_back();
    auto new_q = c.frequency[c.frequency.length() - 1];

    for (size_t j = 0; j < c.data.num_components(); ++j)
    {
        auto s = size_of_intersection(candidate.row_ids, masks[j]);
        // auto q   = static_cast<float_type>(s) / c.data.size();
        auto q   = static_cast<float_type>(s) / c.data.subset(j).size();
        new_q(j) = q;
    }
    new_q.back() = glob_frequency;
}

template <typename Trait, typename Candidate, typename Masks>
bool find_assignment(Composition<Trait>& c,
                     const Candidate&    candidate,
                     const Masks&        masks,
                     bool                bic)
{
    size_t assignment_count = 0;

    if (c.data.num_components() == 1)
    {
        assignment_count += find_assignment_impl(c, candidate, bic);
    }
    else
    {
        assignment_count += find_assignment_impl(c, candidate, masks, bic);
    }

    if (assignment_count > 0)
    {
        insert_pattern_to_summary(c, candidate, masks);
    }

    return assignment_count > 0;
}

template <typename Trait>
auto construct_component_masks(const Composition<Trait>& c)
{
    using tid_container = long_storage_container<typename Trait::pattern_type>;

    assert(check_invariant(c));

    std::vector<tid_container> masks;
    masks.reserve(16);
    if (c.data.size() > 1)
    {
        masks.resize(c.data.num_components(), tid_container{c.data.size()});
        for (size_t s = 0, row = 0, n = c.data.num_components(); s < n; ++s)
        {
            for ([[maybe_unused]] const auto& x : c.data.subset(s))
            {
                masks[s].insert(row);
                ++row;
            }
        }
    }
    return masks;
}

} // namespace disc
} // namespace sd