#pragma once

#include <relaxent/IntersectingFactorization.hxx>
#include <relaxent/UsageHeuristic.hxx>

#include <disc/disc/AssignPatternToComponents.hxx>
#include <disc/disc/Composition.hxx>
#include <disc/disc/PatternsetResult.hxx>

#include <cctype>
#include <type_traits>

namespace sd::disc
{

template <typename Trait, typename Candidate, typename Config, typename Objective>
bool checked_insert_copy(PatternsetResult<Trait>& c,
                         Candidate&               x,
                         const Config&            cfg,
                         Objective&&              objective_fn)
{
    using float_type = typename Trait::float_type;

    if (x.score < 0)
        return false;

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

    c.model.insert(fr, x.pattern, true);
    auto enc = objective_fn(std::as_const(c), cfg);

    if (enc.objective() < c.encoding.objective())
    {
        c.encoding = enc;
        c.summary.insert(fr, x.pattern);
        return true;
    }
    else
    {
        c.model = std::move(old);
        return false;
    }
}

template <typename Trait, typename Candidate, typename Config, typename Objective>
bool checked_insert_append(PatternsetResult<Trait>& c,
                           Candidate&               x,
                           const Config&            cfg,
                           Objective&&              objective_fn)
{
    using float_type = typename Trait::float_type;

    if (x.score <= 0)
        return false;

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

    auto enc = c.encoding;
    {
        c.model.insert(fr, x.pattern, true);
        c.summary.insert(fr, x.pattern);
        enc = objective_fn(c, cfg);
    }

    if (enc.objective() < c.encoding.objective())
    {
        c.encoding = enc;
        return true;
    }
    else
    {
        c.model.model.factors.pop_back();
        c.summary.pop_back();
        return false;
    }
}

template <typename Trait, typename Candidate, typename Config, typename Objective>
bool checked_insert_copy(Composition<Trait>& c,
                         Candidate&          x,
                         const Config&       cfg,
                         Objective&&         objective_fn)
{
    size_t assignment_count = 0;

    using float_type = typename Trait::float_type;

    auto old = c.models; // expensive

    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)
        {
            c.models[i].insert(fr, x.pattern, true);
            ++assignment_count;
        }
    }

    if (assignment_count > 0)
    {
        auto enc = objective_fn(c, cfg);

        if (enc.objective() < c.encoding.objective())
        {
            c.encoding = enc;
            insert_pattern_to_summary(c, x, c.masks);
            return true;
        }
        else
        {
            c.models = std::move(old);
            for (auto& a : c.assignment)
                a.erase(c.summary.size());
            return false;
        }
    }
    return false;
}

template <typename Trait, typename Candidate, typename Config, typename Objective>
bool checked_insert_append(Composition<Trait>& c,
                           Candidate&          x,
                           const Config&       cfg,
                           Objective&&         objective_fn)
{
    size_t assignment_count = 0;

    using float_type = typename Trait::float_type;

    float_type additional_cost = 0;

    if (!cfg.use_bic)
    {
        additional_cost += constant_mdl_cost(c, x.pattern);
    }

    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)
        {
            c.models[i].insert(fr, x.pattern, true);
            ++assignment_count;
            if (cfg.use_bic)
            {
                additional_cost += additional_cost_bic(d.size());
            }
            else
            {
                additional_cost += additional_cost_mdl(x.support);
            }
        }
    }

    if (assignment_count > 0)
    {
        auto enc = objective_fn(c, cfg);
        enc.of_model += additional_cost;

        if (enc.objective() < c.encoding.objective())
        {
            insert_pattern_to_summary(c, x, c.masks);
            c.encoding          = enc;
            c.encoding.of_model = encode_model(c, cfg);
            return true;
        }
        else
        {
            for (size_t i = 0; i < c.data.num_components(); ++i)
            {
                if (c.assignment[i].contains(c.summary.size()))
                {
                    c.assignment[i].contains(c.summary.size());
                    c.models[i].model.factors.pop_back();
                }
            }
            return false;
        }
    }
    return false;
}

template <typename C, typename Candidate, typename Config, typename Objective>
decltype(auto) checked_insert(C& c, Candidate& x, const Config& cfg, Objective objective_fn)
{

    if constexpr (is_append_only_dist<typename std::decay_t<C>::distribution_type>())
    {
        // debug_print(c, x, cfg, objective_fn);
        return checked_insert_append(c, x, cfg, objective_fn);
    }
    else
    {
        // debug_print_short(c, x, cfg, objective_fn);
        return checked_insert_copy(c, x, cfg, objective_fn);
    }
}

} // namespace sd::disc