
#pragma once

#include <disc/disc/Encoding.hxx>
#include <relaxent/IntersectingFactorization.hxx>
#include <relaxent/ModelCost.hxx>

namespace sd::disc
{

template <typename Callback, typename Factors>
void erase_or_keep_factors(const std::vector<size_t>& u,
                                     Factors&                   facs,
                                     Callback&&                 call)
{
    for (size_t i = 0; i < u.size(); ++i)
    {
        if (u[i] != 0)
        {
            const auto& fac = facs[i].factor.itemsets.set;

            for (const auto& y : fac)
            {
                call(y);
            }
        }
        else if (!is_singleton(facs[i].range))
        {
            facs[i].factor.itemsets.set.clear();
            facs[i].factor.singletons.set.clear();
            facs[i].range.clear();
        }
    }

    facs.erase(
        std::remove_if(facs.begin(), facs.end(), [](const auto& f) { return f.range.empty(); }),
        facs.end());
}

template <typename float_type, typename pattern_type, typename Factors>
void erase_unused_factors(const std::vector<size_t>&                u,
                            LabeledDataset<float_type, pattern_type>& s,
                            Factors&                                  facs)
{
    erase_or_keep_factors(u, facs, [&](const auto& y) {
        if (!is_singleton(y.point) && !disc::contains_any(y.point, s))
        {
            s.insert(y.frequency, y.point);
        }
    });
}

template <typename float_type, typename pattern_type, typename Factors>
void erase_unused_factors(const std::vector<size_t>&                u,
                            LabeledDataset<float_type, pattern_type>& s,
                            sparse_dynamic_bitset<size_t>&            a,
                            Factors&                                  facs)
{
    erase_or_keep_factors(u, facs, [&](const auto& y) {
        if (!is_singleton(y.point) && !disc::contains_any(y.point, s))
        {
            a.insert(s.size());
            s.insert(y.frequency, y.point);
        }
    });
}

template <typename Result, typename Model, typename Data, typename Buffer>
void calc_factor_usage(
    const Result& s, std::vector<size_t>& u, Model& m, const Data& data, Buffer& x_buf)
{
    u.resize(m.model.factors.size());
    std::fill(begin(u), end(u), 0);

    for (const auto& x : data)
    {
        x_buf.assign(point(x));
        viva::factorize_model(m.model, s.data.dim, x_buf, [&](size_t index) { u[index]++; });
    }
}

template <typename Trait>
void prune_set(PatternsetResult<Trait>& s, const MiningSettings&)
{
    using pattern_type = typename Trait::pattern_type;

    disc::itemset<pattern_type> x_buf;
    std::vector<size_t>         u(s.model.model.factors.size());

    while (s.summary.size() > s.data.dim)
    {
        s.summary.pop_back();
    }

    // for (size_t i = 0; i < s.summary.size() && s.summary.size() > s.data.dim;) 
    // {
    //     if (is_singleton(s.summary.point(i))
    //     {
    //         std::swap(s.summary[i], s.summary.back());
    //         s.summary.pop_back();
    //     }
    //     else {
    //         ++i;
    //     }
    // }

    calc_factor_usage(s, u, s.model, s.data, x_buf);
    erase_unused_factors(u, s.summary, s.model.model.factors);
}

template <typename Trait>
void prune_set(Composition<Trait>& s, const MiningSettings&)
{
    using pattern_type = typename Trait::pattern_type;

    disc::itemset<pattern_type> x_buf;
    std::vector<size_t>         u;

    while (s.summary.size() > s.data.dim)
    {
        s.summary.pop_back();
    }

    for (size_t k = 0; k < s.data.num_components(); ++k)
    {
        s.assignment[k].clear();
        for (size_t i = 0; i < s.summary.size(); ++i)
            s.assignment[k].insert(i);
    }

    for (size_t k = 0; k < s.data.num_components(); ++k)
    {
        calc_factor_usage(s, u, s.models[k], s.data.subset(k), x_buf);
        erase_unused_factors(u, s.summary, s.assignment[k], s.models[k].model.factors);
    }

    compute_frequency_matrix(s.data, s.summary, s.frequency);
}

} // namespace sd::disc