#pragma once

#include <desc/Desc.hxx>

namespace sd::disc
{

constexpr size_t iterated_log2(size_t n)
{
    // iterated log
    if (n <= 1) return 0;
    if (n <= 2) return 1;
    if (n <= 4) return 2;
    if (n <= 16) return 3;
    if (n <= 65536) return 4;
    return 5; // if (n <= 2^65536)
    // never 6 for our case.
}

constexpr double universal_code(size_t n) { return iterated_log2(n) + 1.5186; }
// 1.5186 = log2(2.865064)

template <typename T, typename X>
auto constant_mdl_cost(Component<T> const& c, const X& x)
{
    using std::log2;
    typename T::float_type acc = universal_code(count(x));
    for_each_item(x, [&](size_t item) {
        auto s = c.frequency[item] * c.data.size();
        acc -= log2((s + (s == 0)) / (c.data.size() + (s == 0)));
    });
    return acc;
}

template <typename T, typename X>
auto constant_mdl_cost(const Composition<T>& c, const X& x)
{
    using std::log2;
    typename T::float_type acc = universal_code(count(x)) + universal_code(c.data.num_components());
    for (size_t j = 0; j < c.data.num_components(); ++j) {
        typename T::float_type l_j = 0;
        for_each_item(x, [&](size_t i) {
            auto s = c.frequency(i, j) * c.data.subset(j).size();
            l_j -= log2((s + (s == 0)) / (c.data.subset(j).size() + (s == 0)));
        });
        acc = j == 0 ? l_j : std::min(acc, l_j);
    }
    return acc;
}
auto additional_cost_mdl(size_t support)
{
    return universal_code(support); // encode-per-component-support
}

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

    float_type acc = -constant_mdl_cost(c, x.pattern) - universal_code(c.data.num_components());

    for (size_t i = 0; i < c.data.num_components(); ++i) {
        auto p = c.models[i].expectation(x.pattern);
        auto s = size_of_intersection(x.row_ids, c.masks[i]);
        auto q = static_cast<float_type>(s) / c.data.subset(i).size();
        using std::log2;
        auto h = s == 0 ? 0 : log2(q / p);

        assert(0 <= p && p <= 1);

        acc += h - additional_cost_mdl(s);
    }

    return acc;
}

template <typename C, typename Distribution, typename Candidate>
auto desc_heuristic_mdl_1(const C& c, const Distribution& pr, const Candidate& x)
{
    using float_type = typename C::float_type;

    // const auto s = static_cast<float_type>(x.support);
    const auto s = x.support;
    const auto q = static_cast<float_type>(s) / c.data.size();
    const auto p = pr.expectation(x.pattern);

    assert(0 <= p && p <= 1);
    using std::log2;
    // float_type cnt = count(x.pattern);

    // std::cout << (q * log2(q / p) + (1 - q) * log2((1 - q) / (1 - p))) << " "
    //           << (cnt * log(c.data.dim / cnt)) << std::endl;
     return c.data.size() * q * log2(q / p) - additional_cost_mdl(s) - constant_mdl_cost(c, x.pattern);

    // return c.data.size() * (q * log2(q / p) + (1 - q) * log2((1 - q) / (1 - p))) -
    //        (cnt * log(c.data.dim / cnt));

    // log2(c.data.dim) - count(x.pattern) * (log2(c.data.size()) + 2.44);
    // return s * log2(q / p) - constant_mdl_cost(c, x.pattern) - additional_cost_mdl(s);
}

struct IDescMDL : DefaultPatternsetMinerInterface
{
    template <typename T, typename Candidate, typename Config>
    static auto heuristic(Component<T>& c, Candidate& x, const Config&)
    {
        return desc_heuristic_mdl_1(c, c.model, x);
    }

    template <typename T, typename Candidate, typename Config>
    static auto heuristic(Composition<T>& c, Candidate& x, const Config&)
    {
        if (c.data.num_components() == 1) {
            return desc_heuristic_mdl_1(c, c.models.front(), x);
        }
        else
        {
            return desc_heuristic_mdl_multi(c, x);
        }
    }

    template <typename C, typename Config>
    static auto finish(C& c, const Config& cfg)
    {
        IDesc::finish(c, cfg);
    }
};

} // namespace sd::disc