#pragma once

#include <disc/disc/Desc.hxx>

namespace sd::disc
{

template <typename S, typename T, typename Fn>
auto sum_over_par(const bit_view<S>& s, T init, Fn&& fn)
{
    if (s.count() > 64)
    {
        auto end = s.length();
#pragma omp parallel for reduction(+ : init)
        for (size_t i = 0; i < end; ++i)
        {
            if (s.contains(i))
            {
                init += fn(i);
            }
        }
    }
    else
    {
        foreach(s, [&](size_t j) { init += fn(j); });
    }
    return init;
}

template <typename S, typename T, typename Fn>
auto sum_over_par(const sparse_bit_view<S>& s, T init, Fn&& fn)
{
    auto end = s.container.size();
#pragma omp parallel for reduction(+ : init) if (s.container.size() > 64)
    for (size_t i = 0; i < end; ++i)
    {
        init += fn(s.container[i]);
    }
    return init;
}

template <typename D, typename M, typename RowIDS>
auto compute_gain_impl(const D& d, const M& pr_before, const M& pr, const RowIDS& rows)
{
    using float_type = typename M::float_type;

    const auto l = float_type(1) / (pr.model.dimension() + d.size());
    const auto u = float_type(1) - l;

    return sum_over_par(rows, float_type{0}, [&](size_t j) -> float_type {
        auto b = pr_before.expected_generalized_frequency(point(d[j]));
        auto p = pr.expected_generalized_frequency(point(d[j]));

        if (!(b > 0) || !(b < 1))
            b = std::clamp(b, l, u);
        if (!(p > 0) || !(p < 1))
            p = std::clamp(p, l, u);

        auto r = -std::log2(b / p);

        if (std::isinf(r) || std::isnan(r) || r != r)
        {
            return 0;
        }

        return r;
    });
}

template <typename X, typename D, typename M, typename RowIDS>
auto compute_gain(const X& x, const D& d, const M& m, const RowIDS& rows)
{
    auto pr = m;
    pr.insert(float(x.support) / d.size(), x.pattern, true);
    return compute_gain_impl(d, m, pr, rows);
}

template <typename Trait, typename X, typename Config>
auto compute_gain(const Composition<Trait>& c, const X& x, const Config& cfg)
{
    using float_type = typename Trait::float_type;

    auto                                  enc = c.encoding;
    itemset<typename Trait::pattern_type> row_subset;
    for (size_t i = 0; i < c.data.num_components(); ++i)
    {
        const auto& d = c.data.subset(i);

        intersection(x.row_ids, c.masks[i], row_subset);
        auto si = count(row_subset);
        // auto        si     = size_of_intersection(x.row_ids, c.masks[i]);
        auto fr            = float_type(si) / d.size();
        auto do_assignment = is_candidate_assignment_good(c, x.pattern, i, fr, cfg.use_bic);

        if (do_assignment)
        {

            auto next = c.models[i];
            next.insert(float_type(x.support) / d.size(), x.pattern, true);

            enc.of_data -= compute_gain_impl(c.data, c.models[i], next, row_subset);
            // /*thread_local*/ dist_type next = c.models[i];
            // next.insert(fr, x.pattern, true);
            // enc.of_data -= c.subset_encodings[i];
            // enc.of_data += log_likelihood(next, d);
            enc.of_model += additional_cost_bic(x.pattern, d.size(), next);
        }
    }
    return c.encoding.objective() - enc.objective();
}

template <typename Trait, typename X, typename Config>
auto compute_gain(const PatternsetResult<Trait>& c, const X& x, const Config&)
{
    auto enc  = c.encoding;
    auto next = c.model;
    next.insert(double(x.support) / c.data.size(), x.pattern, true);
    enc.of_data -= compute_gain_impl(c.data, c.model, next, x.row_ids);
    enc.of_model += additional_cost_bic(x.pattern, c.data.size(), next);
    return c.encoding.objective() - enc.objective();
}

template <typename C, typename Candidate, typename Config>
decltype(auto) actual_gain(C&& c, Candidate& x, const Config& cfg)
{
    return compute_gain(c, x, cfg);
}

template <typename Trait, typename Candidate, typename Config>
auto true_greedy_slow(const PatternsetResult<Trait>& c,
                                   const Candidate&               x,
                                   const Config&)
{
    using float_type = typename Trait::float_type;
    using dist_type  = typename Trait::distribution_type;

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

    /*thread_local*/ dist_type next;
    next = c.model;                  // expensive
    next.insert(fr, x.pattern, true);

    auto enc    = c.encoding;
    enc.of_data = log_likelihood(next, c.data);
    enc.of_model += additional_cost_bic(x.pattern, c.data.size(), next);

    return c.encoding.objective() - enc.objective();
}

template <typename Trait, typename Candidate, typename Config>
auto true_greedy_slow(const Composition<Trait>& c,
                                   const Candidate&          x,
                                   const Config&             cfg)
{
    using float_type = typename Trait::float_type;
    using dist_type  = typename Trait::distribution_type;

    auto enc = c.encoding;

    for (size_t i = 0; i < c.data.num_components(); ++i)
    {
        const auto& d      = c.data.subset(i);
        auto        si     = size_of_intersection(x.row_ids, c.masks[i]);
        auto        fr     = float_type(si) / d.size();
        auto do_assignment = is_candidate_assignment_good(c, x.pattern, i, fr, cfg.use_bic);

        if (do_assignment)
        {
            /*thread_local*/ dist_type next = c.models[i];
            next.insert(fr, x.pattern, true);
            enc.of_data -= c.subset_encodings[i];
            enc.of_data += log_likelihood(next, d);
            enc.of_model += additional_cost_bic(x.pattern, d.size(), next);
        }
    }
    return c.encoding.objective() - enc.objective();
}


} // namespace sd::disc