#pragma once

#include <disc/distribution/Distribution.hxx>
#include <disc/distribution/FactorizedModel.hxx>
#include <disc/distribution/IterativeScaling.hxx>
#include <disc/storage/Itemset.hxx>

#include <ndarray/ndarray.hxx>
#include <setcover/GreedySetcover.hxx>

namespace sd::viva
{

struct SuperSetcoverWeightFunction
{
    template <typename Set, typename To_Cover>
    float operator()(size_t, const Set& f, const To_Cover& y) const noexcept
    {
        const auto s = size_of_intersection(y, f.range);

        if (s > 1 && f.factor.itemsets.get_precomputed_expectation(y))
        {
            return std::numeric_limits<float>::max();
        }

        return static_cast<float>(s) - 1;
    }
};

struct SetMinusFactor
{
    template <typename To_Cover, typename Set>
    void operator()(To_Cover& s, const Set& f) const noexcept
    {
        sd::setminus(s, f.range);
    }
};

template <typename U, typename V, typename T, typename F>
void factorize_model(
    FactorizedModel<U, V> const& m, size_t, T& t, F&& f, size_t limit = 1'000'000)
{
    for (size_t i = 0; i < m.factors.size() && limit-- > 0; ++i)
    {
        if (intersects(t, m.factors[i].range))
        {
            f(i);
        }
    }
}

template <typename U, typename V, typename Mode>
struct IntersectingFactorsModel;

template <typename U, typename V, typename S, typename F, typename Weight>
void factorize_model_itemsets(IntersectingFactorsModel<U, V, Weight> const& m,
                              size_t                                        dim,
                              S&                                            x,
                              F&&                                           f,
                              size_t                                        limit = 1'000'000)
{
    auto phi    = sd::make_slice(m.factors).cut(dim);
    auto report = [&](const auto& index) { f(index + dim); };
    thread_local disc::itemset<disc::tag_dense> ignore(1024);

    ignore.resize(phi.size());
    ignore.reset();

    setcover::greedy_set_cover_no_minus(
        ignore, phi, x, report, {limit, -1}, Weight{}, SetMinusFactor{});

}

template <typename U, typename V, typename S, typename F, typename Mode>
void factorize_model_singletons(IntersectingFactorsModel<U, V, Mode> const&,
                                size_t,
                                S&  x,
                                F&& f)
{
    foreach(x, f);
}

template <typename U, typename V, typename S, typename F, typename Mode>
void factorize_model(IntersectingFactorsModel<U, V, Mode> const& m,
                     size_t                                      dim,
                     S&                                          x,
                     F&&                                         f,
                     size_t                                      limit = 1'000'000)
{
    factorize_model_itemsets(m, dim, x, f, limit);
    factorize_model_singletons(m, dim, x, f);
}

template <typename U, typename V, typename T>
std::pair<size_t, size_t>
size_and_width_of_new_factor_after_insert(FactorizedModel<U, V> const& p, const T& t)
{
    size_t total_size  = 0;
    size_t total_width = 0;
    for (auto& f : p.factors)
    {
        if (is_subset(t, f.range))
        {
            return {f.factor.itemsets.set.size() + 1, f.factor.singletons.set.size()};
        }
        if (intersects(t, f.range))
        {
            // this works, because all factors are covering disjoint sets
            total_size += f.factor.itemsets.set.size();
            total_width += count(f.range);
        }
    }
    return {total_size, total_width};
}

template <typename U, typename V, typename T, typename Mode>
std::pair<size_t, size_t>
size_and_width_of_new_factor_after_insert(IntersectingFactorsModel<U, V, Mode> const& p,
                                          const T&                                    t)
{
    size_t total_size  = 0;
    size_t total_width = 0;

    thread_local disc::itemset<U> x_buf;
    x_buf.assign(t);

    factorize_model_itemsets(p, p.dim, x_buf, [&](const auto& index) {
        total_size += p.factors[index].factor.itemsets.set.size();
        total_width += size_of_intersection(p.factors[index].range, x_buf);
    });

    return {total_size + 1, total_width + count(x_buf)};
}

template <typename U, typename V, typename Weight = SuperSetcoverWeightFunction>
struct IntersectingFactorsModel : public FactorizedModel<U, V>
{
    using float_type  = typename FactorizedModel<U, V>::float_type;
    using factor_type = typename FactorizedModel<U, V>::factor_type;

private:
    using FactorizedModel<U, V>::insert_pattern;
    using FactorizedModel<U, V>::is_pattern_feasible;
    using FactorizedModel<U, V>::insert;

public:
    using FactorizedModel<U, V>::FactorizedModel;
    using FactorizedModel<U, V>::join_factors;

    template <typename T>
    void insert_pattern(float_type fr,
                        const T&   t,
                        size_t     max_factor_size,
                        size_t     max_factor_width,
                        bool       estimate = false)
    {

        if (max_factor_size <= 1)
        {
            insert_pattern_join_only_singletons(fr, t, estimate);
            return;
        }

        disc::itemset<U> x_buf(t);
        factor_type      next(this->dim);

        factorize_model(*this, this->dim, x_buf, [&](size_t index) {
            join_factors(next, this->factors[index]);
        });

        const auto next_elems = next.factor.itemsets.set.size();
        const auto next_items = next.factor.singletons.set.size();

        if (next_elems >= (max_factor_size * 0.8) || next_items > max_factor_width)
        {
            prune_greedily(next, fr, t, max_factor_size, max_factor_width);
        }
        else
        {
            next.factor.insert(fr, t, estimate);
        }

        next.range.insert(t);

        this->factors.emplace_back(std::move(next));
    }

    template <typename T>
    bool is_pattern_feasible(const T& t, size_t max_factor_size, size_t max_factor_width) const
    {
        if (count(t) > max_factor_width)
            return false;

        auto [total_size, total_width] = size_and_width_of_new_factor_after_insert(*this, t);

        if (total_size >= max_factor_size || total_width >= max_factor_width)
            return false;
        else
            return true;
    }

    template <typename T>
    bool is_pattern_feasible(const T& t) const
    {
        return is_pattern_feasible(t, this->max_factor_size, this->max_factor_width);
    }

    template <typename T>
    void insert(float_type frequency, const T& t, bool estimate = false)
    {
        using this_type = IntersectingFactorsModel<U, V, Weight>;
        using base_type = FactorizedModel<U, V>;

        if (is_singleton(t))
            base_type::insert_singleton(frequency, t, estimate);
        else
            this_type::insert_pattern(
                frequency, t, this->max_factor_size, this->max_factor_width, estimate);
    }

private:
    template <typename T>
    void prune_greedily(factor_type& next,
                        float_type   frequency,
                        const T&     t,
                        size_t       max_factor_size,
                        size_t       max_factor_width,
                        bool         estimate = true) const
    {

        auto&       is = next.factor.itemsets.set;
        factor_type replacement(this->dim);

        // singletons
        for (const auto& i : next.factor.singletons.set)
        {
            replacement.factor.insert_singleton(i.frequency, i.element);
        }
        replacement.range.insert(t);

        // most informative next pattern
        replacement.factor.insert(frequency, t, estimate);

        // greedy algorithm to select the sets from next that happen to increase the likelihood
        hybrid_bitset<size_t, 1> in_use;

        for (size_t count = 1, n = is.size(); count < max_factor_size;)
        {
            std::pair best{n, float_type(0)};
            for (size_t j = 0; j < n; ++j)
            {
                if (in_use.test(j))
                    continue;

                const auto sou = size_of_union(is[j].point, replacement.range);
                if (sou > max_factor_width)
                {
                    in_use.insert(j);
                    continue;
                }

                auto p = expected_frequency(replacement.factor, is[j].point);
                auto q = is[j].frequency;
                auto g = std::log2(q / p);

                if (g > best.second)
                {
                    best = {j, g};
                }
            }

            if (best.first >= n)
                break;

            const auto& best_y = is[best.first].point;
            const auto& best_q = is[best.first].frequency;

            if (best.second >= 0)
            {
                ++count;
                in_use.insert(best.first);
                for (const auto& i : next.factor.singletons.set)
                {
                    if (is_subset(i.element, best_y) &&
                        !is_subset(i.element, replacement.range))
                    {
                        replacement.factor.insert_singleton(i.frequency, i.element);
                    }
                }
                replacement.range.insert(best_y);
                replacement.factor.insert(best_q, best_y, estimate);
            }
            else
            {
                break;
            }
        }

        auto& sis = replacement.factor.singletons.set;
        sis.erase(std::remove_if(
                      sis.begin(),
                      sis.end(),
                      [&](const auto& i) { return !is_subset(i.element, replacement.range); }),
                  sis.end());
        next = std::move(replacement);
    }

    template <typename T>
    void insert_pattern_join_only_singletons(float_type frequency, const T& t, bool estimate)
    {
        factor_type next(this->dim);
        for (auto& f : this->factors)
        {
            if (is_singleton(f.range) && intersects(f.range, t))
            {
                join_factors(next, f);
            }
        }
        next.range.insert(t);
        next.factor.insert(frequency, t, estimate);
        this->factors.emplace_back(std::move(next));
    }
};

template <typename S, typename T, typename U, typename Mode>
auto expected_frequency(IntersectingFactorsModel<S, T, Mode> const& m, U const& x)
{
    thread_local disc::itemset<S> part, x_buf;
    x_buf.clear();
    x_buf.insert(x);
    typename IntersectingFactorsModel<S, T>::float_type fr = 1;

    factorize_model_itemsets(m, m.dim, x_buf, [&](size_t index) {
        part.clear();
        intersection(x, m.factors[index].range, part);
        fr *= expected_frequency(m.factors[index].factor, part);
    });
    factorize_model_singletons(m, m.dim, x_buf, [&](size_t index) {
        fr *= m.factors[index].factor.singletons.set.front().probability;
    });

    return fr;
}

template <typename query_type, typename pattern_type, typename float_type, typename Mode>
auto expected_generalized_frequency(
    const IntersectingFactorsModel<pattern_type, float_type, Mode>& m, const query_type& x)
{
    float_type p = expected_frequency(m, x);
    for (size_t i = 0; i < m.dimension(); ++i)
    {
        const auto& f = m.factors[i];
        const auto& e = f.factor.singletons.set;
        if (e.size() == 1 && !is_subset(e.front().element, x))
        {
            p *= float_type(1) - e.front().probability;
            // p *= probability_of_absent_items(f.factor, x);
        }
    }
    return p;
}

} // namespace sd::viva

namespace sd::disc
{

template <typename S, typename T, typename W = viva::SuperSetcoverWeightFunction>
struct IntersectingFactorsDistribution
    : MEDistributionImpl<viva::IntersectingFactorsModel<S, T, W>>
{
    using base = MEDistributionImpl<viva::IntersectingFactorsModel<S, T, W>>;

    IntersectingFactorsDistribution() = default;

    IntersectingFactorsDistribution(size_t dimension,
                                    size_t length,
                                    size_t max_factor_size,
                                    size_t max_factor_width)
        : base(dimension, length)
    {
        this->model.max_factor_size = max_factor_size;
        this->model.max_factor_width   = max_factor_width;
    }

    IntersectingFactorsDistribution(size_t                      dimension,
                                    size_t                      length,
                                    const disc::MiningSettings& cfg = {})
        : IntersectingFactorsDistribution(
              dimension, length, cfg.max_factor_size, cfg.max_factor_width)
    {
    }

    template <typename pattern_t>
    auto expected_frequency(const pattern_t& t) const
    {
        return viva::expected_frequency(this->model, t);
    }

    template <typename pattern_t>
    auto expected_generalized_frequency(const pattern_t& t) const
    {
        return viva::expected_generalized_frequency(this->model, t);
    }
};

template <typename U, typename V, typename F = double, typename Mode>
void estimate_model(IntersectingFactorsDistribution<U, V, Mode>& m,
                    viva::IterativeScalingSettings<F> const&     opts = {})
{
    estimate_model(static_cast<viva::FactorizedModel<U, V>&>(m.model), opts);
}

template <typename D>
constexpr bool is_append_only_dist()
{
    using U  = std::decay_t<D>;
    using S  = typename U::pattern_type;
    using T  = typename U::float_type;
    // using O1 = IntersectingFactorsDistribution<S, T, viva::SetcoverWeightFunction>;
    using O2 = IntersectingFactorsDistribution<S, T, viva::SuperSetcoverWeightFunction>;
    // using O3 = IntersectingFactorsDistribution<S, T, viva::WeightedSuperSetcoverWeightFunction>;

    return std::is_same_v<U, O2>;
    // return std::is_same_v<U, O1> || std::is_same_v<U, O2> || std::is_same_v<U, O3>;
}

} // namespace sd::disc
