#pragma once

#include <relaxent/IntersectingFactorization.hxx>
#include <relaxent/ModelCost.hxx>
#include <relaxent/PatternsetPruning.hxx>

#include <disc/disc/Desc.hxx>

namespace sd::disc
{

template <typename D, typename M, typename RowIDS, typename buffer_type>
size_t usage_last_factor(const D& d, const M& m, buffer_type& buf)
{
    size_t usage = 0;
    for (size_t j = 0; j < d.size(); ++j)
    {
        buf.assign(point(d[j]));
        viva::factorize_model_itemsets(m.model, m.model.dim, buf, [&](size_t k) {
            if (k == m.model.factors.size() - 1)
            {
                ++usage;
                buf.clear(); // early stopping
            }
        });
    }
    return usage;
}

template <typename D, typename M, typename RowIDS, typename buffer_type>
size_t usage_last_factor(const D& d, const M& m, const RowIDS& rows, buffer_type& buf)
{
    size_t usage = 0;
    foreach(rows, [&](size_t j) {
        buf.assign(point(d[j]));
        viva::factorize_model_itemsets(m.model, m.model.dim, buf, [&](size_t k) {
            if (k == m.model.factors.size() - 1)
            {
                ++usage;
                buf.clear(); // early stopping
            }
        });
    });
    return usage;
}

template <typename X, typename D, typename M, typename RowIDS, typename buffer_type>
size_t
usage_of_new_factor(const X& x, const D& d, const M& dist, const RowIDS& rows, buffer_type& buf)
{
    thread_local std::decay_t<M> m;
    m = dist;
    // dist.model.factors.emplace_back(m.model.dim);

    // training is not necessary
    m.model.factors.emplace_back(m.model.dim);
    m.model.factors.back().range = x;
    m.model.factors.back().factor.itemsets.insert(0.5, x);
    m.model.factors.back().factor.itemsets.num_singletons = count(x);
    // m.model.factors.back().factor.singletons.set.resize(count(x));

    return usage_last_factor(d, m, rows, buf);
    // dist.model.factors.pop_back();

}

template <typename X, typename D, typename M, typename RowIDS, typename buffer_type>
auto usage_gain(const X&      pattern,
                size_t        support,
                const D&      d,
                const M&      m,
                const RowIDS& rows,
                buffer_type&  buf)
{
    using float_type = typename M::float_type;

    const auto fr = static_cast<float_type>(support) / d.size();
    const auto p  = m.expected_frequency(pattern);
    const auto g  = std::log2(fr / p);

    if (m.model.factors.size() <= m.model.dim)
    {
        return std::pair{support, support * g};
    }
    else
    {
        const auto usage = g > 0 ? usage_of_new_factor(pattern, d, m, rows, buf) : support;
        return std::pair{usage, g * usage};
    }
}

template <typename Trait, typename Pattern>
auto heuristic_usage(const PatternsetResult<Trait>& c, const Pattern& x)
{
    using float_type   = typename Trait::float_type;
    using pattern_type = typename Trait::pattern_type;

    float_type gain = -additional_cost_bic(x.pattern, c.data.size(), c.model);

    if constexpr (is_append_only_dist<std::decay_t<typename Trait::distribution_type>>())
    {
        thread_local itemset<pattern_type> buffer;
        gain += usage_gain(x.pattern, x.support, c.data, c.model, x.row_ids, buffer).second;
    }
    else
    {
        const auto fr = static_cast<float_type>(x.support) / c.data.size();
        const auto p  = c.model.expected_frequency(x.pattern);
        gain += x.support * std::log2(fr / p);
    }

    return gain;
}

template <typename Trait, typename Pattern>
auto heuristic_usage(const Composition<Trait>& c, const Pattern& x)
{
    using float_type   = typename Trait::float_type;
    using pattern_type = typename Trait::pattern_type;

    float_type acc = 0;

    for (size_t i = 0, k = c.data.num_components(); i < k; ++i)
    {
        const auto& d = c.data.subset(i);

        if constexpr (is_append_only_dist<std::decay_t<typename Trait::distribution_type>>())
        {
            thread_local itemset<pattern_type> rows, buffer;
            rows.clear();

            intersection(x.row_ids, c.masks[i], rows);
            acc += usage_gain(x.pattern, x.support, c.data, c.models[i], rows, buffer).second;
        }
        else
        {
            auto support  = size_of_intersection(x.row_ids, c.masks[i]);
            auto estimate = c.models[i].expected_frequency(x.pattern);
            auto fr       = static_cast<float_type>(support) / d.size();
            acc += support * std::log2(fr / estimate);
        }

        acc -= additional_cost_bic(x.pattern, d.size(), c.models[i]);
    }
    return acc;
}

} // namespace sd::disc