/// @file dpf/placement.hpp
/// @brief Prefix placement (`at<N>`), phantom cmp tag, and slot-meta machinery.
/// @details Shared by `dpf_key` (unified key) and `incremental.hpp` (gen/eval).
///          Holds the pure type-level pieces so `dpf_key.hpp` can build a
///          `slot_meta` table for multi-level / cmp keys without depending on
///          the generation / evaluation code that lives in `incremental.hpp`.
/// @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_PLACEMENT_HPP__
#define LIBDPF_INCLUDE_DPF_PLACEMENT_HPP__

#include <cstddef>
#include <array>
#include <tuple>
#include <type_traits>
#include <utility>
#include <algorithm>

#include "dpf/utils.hpp"
#include "dpf/leaf_node.hpp"
#include "dpf/wildcard.hpp"

namespace dpf
{

// ---------------------------------------------------------------------------
// Public placement sugar: `at<N>(y, ys...)`
// ---------------------------------------------------------------------------

template <std::size_t N, typename OutputT, typename ...OutputTs>
struct at_pack
{
    static constexpr std::size_t prefix = N;
    using outputs_tuple = std::tuple<OutputT, OutputTs...>;
    outputs_tuple values;

    at_pack() = delete;
    explicit at_pack(OutputT y, OutputTs ...ys)
      : values{std::move(y), std::move(ys)...} { }
};

template <std::size_t N>
struct at_fn
{
    template <typename OutputT, typename ...OutputTs>
    constexpr auto operator()(OutputT y, OutputTs ...ys) const
    {
        return at_pack<N, OutputT, OutputTs...>(std::move(y), std::move(ys)...);
    }
};

template <std::size_t N>
inline constexpr at_fn<N> at{};

template <typename T> struct is_at : std::false_type {};
template <std::size_t N, typename O, typename ...Os>
struct is_at<at_pack<N, O, Os...>> : std::true_type {};
template <typename T> inline constexpr bool is_at_v = is_at<T>::value;

/// Phantom pack element for a key's comparison (DCF) channel. Not a leaf: it
/// only records the cmp prefix depth in the key's type. `Depth` is the number
/// of tree levels the comparison walks (0 is reserved for "no cmp").
/// `OutBits` is the comparison output group width (bits of the β payload),
/// so the value CWs / addend can be stored at group width instead of a full
/// padded `uint64_t` per level.
template <std::size_t Depth, std::size_t OutBits = 0, bool Wild = false>
struct cmp_channel_tag
{
    static constexpr std::size_t depth = Depth;
    static constexpr std::size_t out_bits = OutBits;
    /// True when the comparison payload (β) is a wildcard to be assigned
    /// after keygen. Concrete (non-wildcard) cmp keys keep `Wild == false`
    /// so their layout / type name is unchanged.
    static constexpr bool wild = Wild;
};

template <typename T> struct is_cmp_channel_tag : std::false_type {};
template <std::size_t Depth, std::size_t OutBits, bool Wild>
struct is_cmp_channel_tag<cmp_channel_tag<Depth, OutBits, Wild>>
    : std::true_type {};
template <typename T>
inline constexpr bool is_cmp_channel_tag_v =
    is_cmp_channel_tag<std::decay_t<T>>::value;

namespace detail
{
namespace incr
{

// ---------------------------------------------------------------------------
// A concrete placed output: an output type `OutputT` planted at prefix `N`.
// ---------------------------------------------------------------------------

template <std::size_t N, typename OutputT>
struct placed
{
    static constexpr std::size_t prefix = N;
    using output_type = OutputT;
    OutputT value;
    OutputT addend{};  // public if_false for eq(...); party 0 absorbs at eval
};

template <typename T> struct is_placed : std::false_type {};
template <std::size_t N, typename O>
struct is_placed<placed<N, O>> : std::true_type {};
template <typename T>
inline constexpr bool is_placed_v = is_placed<std::decay_t<T>>::value;

template <typename NodeT, typename OutputT>
inline constexpr std::size_t out_bits_v =
    utils::bitlength_of_output_v<concrete_type_t<OutputT>, NodeT>;

template <typename NodeT, typename OutputT>
inline constexpr std::size_t lg_opl_v =
    dpf::lg_outputs_per_leaf_v<concrete_type_t<OutputT>, NodeT>;

template <typename NodeT, std::size_t N, typename OutputT>
inline constexpr std::size_t level_of_v = N - lg_opl_v<NodeT, OutputT>;

template <typename NodeT, std::size_t N, typename OutputT>
inline constexpr bool prefix_ok_v = (N >= lg_opl_v<NodeT, OutputT>);

// ---------------------------------------------------------------------------
// Per-output slot metadata (packing / tree-level table).
// ---------------------------------------------------------------------------

struct slot_meta
{
    std::size_t prefix;
    std::size_t tree_level;
    std::size_t pos_base;
    std::size_t group_id;
    std::size_t index_in_group;
    std::size_t out_bits;
    std::size_t lg_opl;
    std::size_t block_len;
};

template <typename NodeT, typename PlacedTuple, std::size_t ...Is>
constexpr std::size_t max_tree_level_impl(std::index_sequence<Is...>)
{
    std::size_t m = 0;
    ((m = std::max(m, level_of_v<NodeT,
        std::tuple_element_t<Is, PlacedTuple>::prefix,
        typename std::tuple_element_t<Is, PlacedTuple>::output_type>)), ...);
    return m;
}

template <typename NodeT, typename PlacedTuple>
inline constexpr std::size_t max_tree_level_v =
    max_tree_level_impl<NodeT, PlacedTuple>(
        std::make_index_sequence<std::tuple_size_v<PlacedTuple>>{});

template <typename NodeT, typename PlacedTuple, std::size_t ...Is>
constexpr bool all_prefixes_ok_impl(std::index_sequence<Is...>)
{
    return (prefix_ok_v<NodeT,
        std::tuple_element_t<Is, PlacedTuple>::prefix,
        typename std::tuple_element_t<Is, PlacedTuple>::output_type> && ...);
}

template <typename NodeT, typename PlacedTuple>
inline constexpr bool all_prefixes_ok_v =
    all_prefixes_ok_impl<NodeT, PlacedTuple>(
        std::make_index_sequence<std::tuple_size_v<PlacedTuple>>{});

template <typename NodeT, typename PlacedTuple, std::size_t I>
constexpr void fill_slot_basics(
    std::array<slot_meta, std::tuple_size_v<PlacedTuple>> & meta)
{
    using P = std::tuple_element_t<I, PlacedTuple>;
    using O = typename P::output_type;
    meta[I].prefix = P::prefix;
    meta[I].lg_opl = lg_opl_v<NodeT, O>;
    meta[I].out_bits = out_bits_v<NodeT, O>;
    meta[I].tree_level = P::prefix - meta[I].lg_opl;
    meta[I].block_len =
        dpf::block_length_of_leaf_v<concrete_type_t<O>, NodeT>;
    meta[I].pos_base = 0;
    meta[I].group_id = 0;
    meta[I].index_in_group = 0;
}

template <typename NodeT, typename PlacedTuple, std::size_t ...Is>
constexpr void fill_all_basics(
    std::array<slot_meta, std::tuple_size_v<PlacedTuple>> & meta,
    std::index_sequence<Is...>)
{
    (fill_slot_basics<NodeT, PlacedTuple, Is>(meta), ...);
}

template <typename NodeT, typename PlacedTuple>
constexpr auto build_meta()
{
    constexpr std::size_t n = std::tuple_size_v<PlacedTuple>;
    std::array<slot_meta, n> meta{};
    fill_all_basics<NodeT, PlacedTuple>(meta, std::make_index_sequence<n>{});

    std::size_t next_gid = 0;
    for (std::size_t i = 0; i < n; ++i)
    {
        std::size_t gid = static_cast<std::size_t>(-1);
        for (std::size_t j = 0; j < i; ++j)
        {
            if (meta[j].prefix == meta[i].prefix
                && meta[j].out_bits == meta[i].out_bits)
            {
                gid = meta[j].group_id;
                break;
            }
        }
        if (gid == static_cast<std::size_t>(-1))
            gid = next_gid++;
        meta[i].group_id = gid;
    }

    const std::size_t ngroups = next_gid;
    std::array<std::size_t, n> group_count{};
    for (std::size_t i = 0; i < n; ++i)
        group_count[i] = 0;
    for (std::size_t i = 0; i < n; ++i)
        meta[i].index_in_group = group_count[meta[i].group_id]++;

    std::array<std::size_t, n> g_level{};
    std::array<std::size_t, n> g_blocks{};
    std::array<std::size_t, n> g_first{};
    for (std::size_t g = 0; g < ngroups; ++g)
    {
        g_blocks[g] = 0;
        g_first[g] = n;
        g_level[g] = 0;
    }
    for (std::size_t i = 0; i < n; ++i)
    {
        const auto g = meta[i].group_id;
        g_level[g] = meta[i].tree_level;
        g_blocks[g] += meta[i].block_len;
        if (i < g_first[g])
            g_first[g] = i;
    }

    constexpr std::size_t depth = max_tree_level_v<NodeT, PlacedTuple>;
    std::array<bool, n> g_done{};
    for (std::size_t g = 0; g < ngroups; ++g)
        g_done[g] = false;

    for (std::size_t level = 0; level <= depth; ++level)
    {
        std::size_t cursor = (level == depth) ? 0 : 2;
        for (;;)
        {
            std::size_t best_g = n;
            std::size_t best_first = n;
            for (std::size_t g = 0; g < ngroups; ++g)
            {
                if (g_done[g] || g_level[g] != level)
                    continue;
                if (g_first[g] < best_first)
                {
                    best_first = g_first[g];
                    best_g = g;
                }
            }
            if (best_g == n)
                break;
            for (std::size_t i = 0; i < n; ++i)
            {
                if (meta[i].group_id == best_g)
                    meta[i].pos_base = cursor;
            }
            cursor += g_blocks[best_g];
            g_done[best_g] = true;
        }
    }
    return meta;
}

template <typename MetaArray, std::size_t NGroups>
constexpr auto build_group_order(const MetaArray & meta, std::size_t n)
{
    std::array<std::size_t, NGroups == 0 ? 1 : NGroups> order{};
    std::array<bool, NGroups == 0 ? 1 : NGroups> used{};
    for (std::size_t g = 0; g < NGroups; ++g)
        used[g] = false;
    for (std::size_t k = 0; k < NGroups; ++k)
    {
        std::size_t best = NGroups;
        std::size_t best_lvl = static_cast<std::size_t>(-1);
        std::size_t best_first = n;
        for (std::size_t g = 0; g < NGroups; ++g)
        {
            if (used[g])
                continue;
            std::size_t lvl = 0, first = n;
            for (std::size_t i = 0; i < n; ++i)
            {
                if (meta[i].group_id == g)
                {
                    lvl = meta[i].tree_level;
                    if (i < first)
                        first = i;
                }
            }
            if (lvl < best_lvl || (lvl == best_lvl && first < best_first))
            {
                best_lvl = lvl;
                best_first = first;
                best = g;
            }
        }
        order[k] = best;
        used[best] = true;
    }
    return order;
}

template <typename InputT>
constexpr InputT lane_input(InputT x, std::size_t prefix, std::size_t bitlen)
{
    if (prefix >= bitlen)
        return x;
    // Shift on the integral representation so this works for `modint`,
    // `keyword` (whose `>>` yields a parent `modint`, not the keyword), and
    // signed/bitstring inputs. Reconstruct the input type from the shifted
    // integral value via `make_from_integral_value` (a friend of `keyword`).
    constexpr auto to_int = utils::to_integral_type<InputT>{};
    using FromI = typename utils::make_from_integral_value<InputT>::integral_type;
    const auto shifted = static_cast<FromI>(to_int(x) >> (bitlen - prefix));
    return utils::make_from_integral_value<InputT>{}(shifted);
}

// Concatenate index sequences
template <typename... Seqs> struct cat_seq;
template <> struct cat_seq<> { using type = std::index_sequence<>; };
template <std::size_t... Is>
struct cat_seq<std::index_sequence<Is...>>
{
    using type = std::index_sequence<Is...>;
};
template <std::size_t... Is, std::size_t... Js, typename... Rest>
struct cat_seq<std::index_sequence<Is...>, std::index_sequence<Js...>, Rest...>
{
    using type = typename cat_seq<std::index_sequence<Is..., Js...>, Rest...>::type;
};
template <typename... Seqs>
using cat_seq_t = typename cat_seq<Seqs...>::type;

template <std::size_t G, typename MetaHolder, std::size_t I>
using keep_if_group = std::conditional_t<
    MetaHolder::value[I].group_id == G,
    std::index_sequence<I>,
    std::index_sequence<>>;

template <std::size_t G, typename MetaHolder, typename Seq>
struct filter_group;
template <std::size_t G, typename MetaHolder, std::size_t... Is>
struct filter_group<G, MetaHolder, std::index_sequence<Is...>>
{
    using type = cat_seq_t<keep_if_group<G, MetaHolder, Is>...>;
};
template <std::size_t G, typename MetaHolder, std::size_t N>
using filter_group_t = typename filter_group<G, MetaHolder,
    std::make_index_sequence<N>>::type;

template <typename KeyT>
struct meta_holder
{
    static constexpr auto value = KeyT::meta;
};

// ---------------------------------------------------------------------------
// Normalize a `dpf_key` output pack into (PlacedTuple, CmpDepth).
//
// Each pack element is one of:
//   - a bare output `T`            -> placed<BitLen, T>
//   - a `placed<N, T>`             -> placed<N, T>  (from `at<N>`)
//   - a `cmp_channel_tag<Depth>`   -> not a leaf; contributes Depth to CmpDepth
// ---------------------------------------------------------------------------

template <std::size_t BitLen, typename Elem>
struct normalize_one
{
    using placed_tuple = std::tuple<placed<BitLen, Elem>>;
    static constexpr std::size_t cmp_depth = 0;
    static constexpr std::size_t cmp_out_bits = 0;
    static constexpr bool cmp_wild = false;
};
template <std::size_t BitLen, std::size_t N, typename T>
struct normalize_one<BitLen, placed<N, T>>
{
    using placed_tuple = std::tuple<placed<N, T>>;
    static constexpr std::size_t cmp_depth = 0;
    static constexpr std::size_t cmp_out_bits = 0;
    static constexpr bool cmp_wild = false;
};
template <std::size_t BitLen, std::size_t Depth, std::size_t OutBits, bool Wild>
struct normalize_one<BitLen, cmp_channel_tag<Depth, OutBits, Wild>>
{
    using placed_tuple = std::tuple<>;
    static constexpr std::size_t cmp_depth = Depth;
    static constexpr std::size_t cmp_out_bits = OutBits;
    static constexpr bool cmp_wild = Wild;
};

template <std::size_t BitLen, typename ...Elems>
struct normalize_pack
{
    using placed_tuple =
        decltype(std::tuple_cat(
            std::declval<typename normalize_one<BitLen, Elems>::placed_tuple>()...));
    static constexpr std::size_t cmp_depth =
        (std::size_t{0} + ... + normalize_one<BitLen, Elems>::cmp_depth);
    // At most one comparison channel per key, so the sum is that channel's
    // output width (0 when there is no cmp channel).
    static constexpr std::size_t cmp_out_bits =
        (std::size_t{0} + ... + normalize_one<BitLen, Elems>::cmp_out_bits);
    // At most one comparison channel per key, so the OR is that channel's
    // wildcard flag (false when there is no cmp channel).
    static constexpr bool cmp_wild =
        (false || ... || normalize_one<BitLen, Elems>::cmp_wild);
};

template <std::size_t BitLen>
struct normalize_pack<BitLen>
{
    using placed_tuple = std::tuple<>;
    static constexpr std::size_t cmp_depth = 0;
    static constexpr std::size_t cmp_out_bits = 0;
    static constexpr bool cmp_wild = false;
};

/// True iff the pack is "classic-shaped": every element is a bare output (no
/// `placed<>` from `at<>` and no `cmp_channel_tag<>`).
template <typename ...Elems>
inline constexpr bool is_classic_pack_v =
    !((is_placed_v<Elems> || is_cmp_channel_tag_v<Elems>) || ...);

}  // namespace incr
}  // namespace detail

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_PLACEMENT_HPP__