/// @file dpf/blocked_dcf.hpp
/// @brief Blocked-checkpoint comparison: one ring word per block of levels.
/// @details The seed spine stays dense. Ring words are published only at a
///          public checkpoint schedule. Point eval expands parked siblings
///          up to the next checkpoint; a full-domain memoizer already holds
///          those nodes. `q` tail bits, when the comparison sets the key
///          depth, are a residual table on the node at height `h`.
/// @copyright Copyright (c) 2019-2026 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license.

#ifndef LIBDPF_INCLUDE_DPF_BLOCKED_DCF_HPP__
#define LIBDPF_INCLUDE_DPF_BLOCKED_DCF_HPP__

#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <utility>
#include <vector>

#include "hedley/hedley.h"

#include "dpf/dcf.hpp"
#include "dpf/path_memoizer.hpp"
#include "dpf/twiddle.hpp"
#include "dpf/utils.hpp"

namespace dpf
{
namespace detail
{
namespace blocked
{

template <std::size_t H, std::size_t B>
struct schedule
{
    static constexpr std::size_t count =
        (B == 0 || H == 0) ? 0 : (H + B - 1) / B;

    static constexpr auto depths = [] {
        std::array<std::size_t, count == 0 ? 1 : count> cs{};
        if constexpr (count == 0)
            return cs;
        const std::size_t base = H / count;
        const std::size_t extra = H % count;
        std::size_t acc = 0;
        for (std::size_t i = 0; i < count; ++i)
        {
            acc += base + (i < extra ? 1 : 0);
            cs[i] = acc;
        }
        return cs;
    }();

    HEDLEY_CONST
    HEDLEY_NO_THROW
    static constexpr bool contains(std::size_t depth) noexcept
    {
        for (std::size_t i = 0; i < count; ++i)
        {
            if (depths[i] == depth)
                return true;
        }
        return false;
    }

    HEDLEY_CONST
    HEDLEY_NO_THROW
    static constexpr std::size_t index(std::size_t depth) noexcept
    {
        for (std::size_t i = 0; i < count; ++i)
        {
            if (depths[i] == depth)
                return i;
        }
        return static_cast<std::size_t>(-1);
    }
};

template <typename PRG, typename Node>
HEDLEY_ALWAYS_INLINE
HEDLEY_NO_THROW
uint64_t rho_of(const Node & node, uint64_t mask) noexcept
{
    auto kids = PRG::eval01(dpf::unset_lo_2bits(node));
    return dcf_impl::convert_node(kids[0], mask);
}

HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr int control_sign(uint8_t t0, uint8_t t1) noexcept
{
    return static_cast<int>(t0) - static_cast<int>(t1);
}

HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr uint64_t mul_sgn(int sgn, uint64_t v, uint64_t mask) noexcept
{
    if (sgn > 0)
        return v & mask;
    if (sgn < 0)
        return dcf_impl::neg_m(v, mask);
    return 0;
}

/// Group element `sgn` (`+1`, `-1`, or `0`) used as an `assign_cmp` coefficient.
HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr uint64_t sgn_coeff(int sgn, uint64_t mask) noexcept
{
    if (sgn > 0)
        return 1ULL & mask;
    if (sgn < 0)
        return dcf_impl::neg_m(1ULL, mask);
    return 0;
}

template <typename PRG, typename Node>
HEDLEY_NO_THROW
uint64_t checkpoint_word(const Node & n0, const Node & n1, uint64_t beta,
    uint64_t mask) noexcept
{
    const int sgn = control_sign(
        static_cast<uint8_t>(dpf::get_lo_bit(n0)),
        static_cast<uint8_t>(dpf::get_lo_bit(n1)));
    const uint64_t r0 = rho_of<PRG>(n0, mask);
    const uint64_t r1 = rho_of<PRG>(n1, mask);
    const uint64_t inner =
        (beta + dcf_impl::neg_m(r0, mask) + r1) & mask;
    return mul_sgn(sgn, inner, mask);
}

template <typename Node>
HEDLEY_NO_THROW
uint64_t checkpoint_coeff(const Node & n0, const Node & n1,
    uint64_t mask) noexcept
{
    return sgn_coeff(control_sign(
        static_cast<uint8_t>(dpf::get_lo_bit(n0)),
        static_cast<uint8_t>(dpf::get_lo_bit(n1))), mask);
}

template <typename PRG, typename Node>
HEDLEY_NON_NULL(4)
HEDLEY_NO_THROW
void suffix_masks(const Node & seed, std::size_t q, uint64_t mask,
    uint64_t * out) noexcept
{
    Node cur[4]{};
    Node nxt[8]{};
    cur[0] = seed;
    std::size_t n = 1;
    for (std::size_t lvl = 0; lvl < q; ++lvl)
    {
        std::size_t m = 0;
        for (std::size_t i = 0; i < n; ++i)
        {
            auto kids = PRG::eval01(dpf::unset_lo_2bits(cur[i]));
            nxt[m++] = kids[0];
            nxt[m++] = kids[1];
        }
        for (std::size_t i = 0; i < m; ++i)
            cur[i] = nxt[i];
        n = m;
    }
    for (std::size_t i = 0; i < n; ++i)
        out[i] = dcf_impl::convert_node(cur[i], mask);
}

template <typename PRG, typename Node>
HEDLEY_NON_NULL(8)
HEDLEY_NO_THROW
void tail_words(const Node & n0, const Node & n1, uint64_t beta, uint64_t mask,
    bool include_eq, uint64_t suffix, std::size_t q, uint64_t * words,
    uint64_t * coeffs) noexcept
{
    uint64_t u0[4]{};
    uint64_t u1[4]{};
    suffix_masks<PRG>(n0, q, mask, u0);
    suffix_masks<PRG>(n1, q, mask, u1);
    const int sgn = control_sign(
        static_cast<uint8_t>(dpf::get_lo_bit(n0)),
        static_cast<uint8_t>(dpf::get_lo_bit(n1)));
    const uint64_t coeff = sgn_coeff(sgn, mask);
    const std::size_t n = std::size_t{1} << q;
    for (std::size_t z = 0; z < n; ++z)
    {
        const bool pred = include_eq
            ? (z <= suffix)
            : (z < suffix);
        const uint64_t inner =
            ((pred ? beta : 0ULL) + dcf_impl::neg_m(u0[z], mask) + u1[z]) & mask;
        words[z] = mul_sgn(sgn, inner, mask);
        if (coeffs != nullptr)
            coeffs[z] = pred ? coeff : 0ULL;
    }
}

template <typename KeyT>
HEDLEY_NO_THROW
uint64_t add_membership(uint64_t acc, const typename KeyT::interior_node & node,
    uint64_t word, uint64_t mask, int party) noexcept
{
    using prg = typename KeyT::interior_prg;
    const uint8_t t = static_cast<uint8_t>(dpf::get_lo_bit(node));
    const uint64_t y =
        (rho_of<prg>(node, mask) + (t ? word : 0ULL)) & mask;
    return (acc + (party ? dcf_impl::neg_m(y, mask) : y)) & mask;
}

template <typename KeyT>
uint64_t add_frontier(uint64_t acc, const typename KeyT::interior_node & seed,
    std::size_t from_depth, std::size_t to_depth, const KeyT & dpf, uint64_t word,
    uint64_t mask, int party)
{
    using node = typename KeyT::interior_node;
    std::vector<node> cur;
    std::vector<node> nxt;
    cur.push_back(seed);
    for (std::size_t lvl = from_depth; lvl < to_depth; ++lvl)
    {
        nxt.clear();
        nxt.reserve(cur.size() * 2);
        const node cw0 = dpf.correction_word(lvl, false);
        const node cw1 = dpf.correction_word(lvl, true);
        for (const node & fs : cur)
        {
            auto kids = KeyT::traverse_interior01(fs, cw0, cw1);
            nxt.push_back(kids[0]);
            nxt.push_back(kids[1]);
        }
        cur.swap(nxt);
    }
    for (const node & fs : cur)
        acc = add_membership<KeyT>(acc, fs, word, mask, party);
    return acc;
}

template <typename KeyT, typename InputT>
HEDLEY_CONST
HEDLEY_NO_THROW
HEDLEY_ALWAYS_INLINE
constexpr uint64_t query_suffix(InputT tx, std::size_t nbits, std::size_t q) noexcept
{
    if (q == 0)
        return 0;
    uint64_t z = 0;
    auto bit_mask = KeyT::msb_mask >> (nbits - q);
    for (std::size_t j = 0; j < q; ++j, bit_mask >>= 1)
        z = (z << 1) | static_cast<uint64_t>(!!(bit_mask & tx));
    return z;
}

template <typename KeyT>
HEDLEY_NO_THROW
uint64_t finish_share(const KeyT & dpf, uint64_t suffix, uint64_t acc,
    const typename KeyT::interior_node & at_h, int party) noexcept
{
    using namespace dcf_impl;
    using prg = typename KeyT::interior_prg;
    const auto & ch = dpf.cmp();
    const uint64_t mask = ch.mask;
    constexpr std::size_t q = KeyT::cmp_q;
    constexpr std::size_t h = KeyT::cmp_h;
    if constexpr (q == 0)
    {
        if (ch.include_eq)
        {
            constexpr auto wi = schedule<h, KeyT::cmp_block>::index(h);
            acc = add_membership<KeyT>(acc, at_h, dpf.value_cw(wi), mask, party);
        }
        (void)suffix;
    }
    else
    {
        const uint64_t z = suffix;
        uint64_t u[4]{};
        suffix_masks<prg>(at_h, q, mask, u);
        const uint8_t t = static_cast<uint8_t>(dpf::get_lo_bit(at_h));
        const uint64_t y = (u[z] + (t ? dpf.tail_cw(z) : 0ULL)) & mask;
        acc = (acc + (party ? neg_m(y, mask) : y)) & mask;
    }
    if (ch.eval_as_ge)
        acc = neg_m(acc, mask);
    const uint64_t add = [&]() -> uint64_t {
        if constexpr (is_party_key_v<KeyT>)
            return dpf.cmp_addend().raw();
        else
            return dpf.cmp_addend();
    }();
    return (acc + add) & mask;
}

template <typename KeyT, typename InputT, typename PathMemoizer>
uint64_t eval_share(const KeyT & dpf, InputT tx, PathMemoizer & path)
{
    using node = typename KeyT::interior_node;
    const auto & ch = dpf.cmp();
    const uint64_t mask = ch.mask;
    const uint64_t add = [&]() -> uint64_t {
        if constexpr (is_party_key_v<KeyT>)
            return dpf.cmp_addend().raw();
        else
            return dpf.cmp_addend();
    }();
    if (ch.trivial == cmp_trivial::always_true
        || ch.trivial == cmp_trivial::always_false)
        return add & mask;

    constexpr std::size_t h = KeyT::cmp_h;
    using sched = schedule<h, KeyT::cmp_block>;
    const std::size_t nbits = static_cast<std::size_t>(ch.nbits);
    const int party = dpf::get_lo_bit(dpf.root()) ? 1 : 0;

    dpf::detail::ensure_level(dpf, tx, path, h);

    struct parked
    {
        node seed;
        std::size_t depth;
    };
    parked pend[128];
    std::size_t npend = 0;

    uint64_t acc = 0;
    auto bit_mask = KeyT::msb_mask;
    for (std::size_t level = 0; level < h; ++level, bit_mask >>= 1)
    {
        const bool xi = !!(bit_mask & tx);
        const node & parent = path[level];
        const node right = KeyT::traverse_interior(parent,
            dpf.correction_word(level, true), true);
        if (!xi)
        {
            pend[npend].seed = right;
            pend[npend].depth = level + 1;
            ++npend;
        }
        const std::size_t c = level + 1;
        if (sched::contains(c))
        {
            const uint64_t word = dpf.value_cw(sched::index(c));
            for (std::size_t p = 0; p < npend; ++p)
            {
                acc = add_frontier<KeyT>(acc, pend[p].seed, pend[p].depth, c,
                    dpf, word, mask, party);
            }
            npend = 0;
        }
    }
    const uint64_t suffix = query_suffix<KeyT>(tx, nbits, KeyT::cmp_q);
    return finish_share(dpf, suffix, acc, path[h], party);
}

template <typename KeyT, typename Integral, typename Memo>
HEDLEY_NO_THROW
bool memo_has(const Memo & memo, Integral prefix, std::size_t depth,
    Integral from_lane, Integral to_excl) noexcept
{
    const auto shift = KeyT::cmp_depth - depth;
    const auto from_p = from_lane >> shift;
    if (prefix < from_p)
        return false;
    const auto idx = static_cast<std::size_t>(prefix - from_p);
    const auto count = memo.get_nodes_at_level(depth, from_lane, to_excl);
    return idx < count;
}

template <typename KeyT, typename Integral, typename Memo>
const typename KeyT::interior_node & memo_node(const Memo & memo, Integral prefix,
    std::size_t depth, Integral from_lane)
{
    const auto shift = KeyT::cmp_depth - depth;
    const auto from_p = from_lane >> shift;
    const auto idx = static_cast<std::size_t>(prefix - from_p);
    return memo[depth][idx];
}

template <typename KeyT, typename Integral, typename Memo>
uint64_t eval_share_memo(const KeyT & dpf, Integral lane,
    Integral from_lane, Integral to_excl, const Memo & memo)
{
    using node = typename KeyT::interior_node;
    const auto & ch = dpf.cmp();
    const uint64_t mask = ch.mask;
    const uint64_t add = [&]() -> uint64_t {
        if constexpr (is_party_key_v<KeyT>)
            return dpf.cmp_addend().raw();
        else
            return dpf.cmp_addend();
    }();
    if (ch.trivial == cmp_trivial::always_true
        || ch.trivial == cmp_trivial::always_false)
        return add & mask;

    constexpr std::size_t h = KeyT::cmp_h;
    using sched = schedule<h, KeyT::cmp_block>;
    const int party = dpf::get_lo_bit(dpf.root()) ? 1 : 0;

    struct parked
    {
        node seed;
        Integral prefix;
        std::size_t depth;
    };
    parked pend[128];
    std::size_t npend = 0;

    uint64_t acc = 0;
    Integral path_pref = 0;
    const std::size_t nbits = static_cast<std::size_t>(ch.nbits);
    for (std::size_t level = 0; level < h; ++level)
    {
        const bool xi = ((lane >> (nbits - 1 - level)) & Integral{1}) != 0;
        const node & parent = memo_node<KeyT>(memo, path_pref, level, from_lane);
        const Integral sib = static_cast<Integral>((path_pref << 1) | Integral{1});
        if (!xi)
        {
            const node right = KeyT::traverse_interior(parent,
                dpf.correction_word(level, true), true);
            pend[npend].seed = right;
            pend[npend].prefix = sib;
            pend[npend].depth = level + 1;
            ++npend;
        }
        path_pref = static_cast<Integral>((path_pref << 1) | Integral{xi ? 1 : 0});
        const std::size_t c = level + 1;
        if (!sched::contains(c))
            continue;
        const uint64_t word = dpf.value_cw(sched::index(c));
        for (std::size_t p = 0; p < npend; ++p)
        {
            const std::size_t extra = c - pend[p].depth;
            const Integral leftmost =
                static_cast<Integral>(pend[p].prefix << extra);
            const Integral rightmost = static_cast<Integral>(
                leftmost + static_cast<Integral>((Integral{1} << extra) - 1));
            const bool covered =
                memo_has<KeyT>(memo, leftmost, c, from_lane, to_excl)
                && memo_has<KeyT>(memo, rightmost, c, from_lane, to_excl);
            if (covered && extra < 16)
            {
                const Integral nleaf = static_cast<Integral>(Integral{1} << extra);
                for (Integral k = 0; k < nleaf; ++k)
                {
                    const auto pref = static_cast<Integral>(leftmost + k);
                    acc = add_membership<KeyT>(acc,
                        memo_node<KeyT>(memo, pref, c, from_lane), word, mask,
                        party);
                }
            }
            else
            {
                acc = add_frontier<KeyT>(acc, pend[p].seed, pend[p].depth, c,
                    dpf, word, mask, party);
            }
        }
        npend = 0;
    }
    const node & at_h = memo_node<KeyT>(memo, path_pref, h, from_lane);
    const uint64_t suffix = static_cast<uint64_t>(lane)
        & (KeyT::cmp_q == 0 ? 0ULL : ((1ULL << KeyT::cmp_q) - 1ULL));
    return finish_share(dpf, suffix, acc, at_h, party);
}

}  // namespace blocked
}  // namespace detail
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_BLOCKED_DCF_HPP__