/// @file dpf/idpf_agg.hpp
/// @brief Max and k-th order statistic on a list of incremental DPF keys.
/// @details Each secret value is one `idpf` with a unit payload on every
///          prefix length. Servers add prefix shares along the live spine
///          with `eval_until`. Communication tracks the bit length of the
///          domain, not how many secret inputs were summed
///          (Cheng–Mitrokotsa–Zhang–Hartmann, ePrint 2024/1190, on the
///          S&P 2021 incremental DPF / Poplar walk).
/// @see dpf/eval_until.hpp, dpf/placement.hpp (`idpf`), examples/applications/idpf_agg.cpp
/// @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_IDPF_AGG_HPP__
#define LIBDPF_INCLUDE_DPF_IDPF_AGG_HPP__

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

#include "hedley/hedley.h"

#include "dpf/eval_until.hpp"
#include "dpf/placement.hpp"
#include "dpf/secret_share.hpp"
#include "dpf/utils.hpp"

namespace dpf
{

/// @brief Unit payload on every prefix length `1 .. N`.
template <std::size_t N, typename Beta = std::uint64_t, std::size_t... Is>
HEDLEY_ALWAYS_INLINE
constexpr auto idpf_ones_impl(std::index_sequence<Is...>)
{
    return idpf(((void)Is, Beta{1})...);
}

/// @brief `idpf(1, 1, …, 1)` of length `N` (one unit per prefix length).
template <std::size_t N, typename Beta = std::uint64_t>
HEDLEY_ALWAYS_INLINE
constexpr auto idpf_ones()
{
    static_assert(N > 0, "idpf_ones: N must be positive");
    return idpf_ones_impl<N, Beta>(std::make_index_sequence<N>{});
}

namespace detail
{
namespace idpf_agg_detail
{

template <typename Share, typename = void>
struct has_raw_member : std::false_type {};
template <typename Share>
struct has_raw_member<Share,
    std::void_t<decltype(std::declval<const Share &>().raw())>>
    : std::true_type {};

template <typename Share>
HEDLEY_ALWAYS_INLINE
auto share_raw(const Share & s)
{
    if constexpr (has_raw_member<Share>::value)
        return s.raw();
    else
        return s;
}

template <typename Key0T, typename Key1T>
auto open_prefix_counts(std::vector<idpf_eval_ctx<Key0T>> & ctx0,
    std::vector<idpf_eval_ctx<Key1T>> & ctx1, std::size_t level,
    const std::vector<typename unwrap_party_key_t<Key0T>::input_type> & prefs)
{
    std::vector<std::uint64_t> counts(prefs.size(), 0);
    for (std::size_t i = 0; i < ctx0.size(); ++i)
    {
        auto s0 = eval_until(ctx0[i], level, prefs);
        auto s1 = eval_until(ctx1[i], level, prefs);
        for (std::size_t j = 0; j < prefs.size(); ++j)
        {
            const auto opened = reconstruct(s0[j], s1[j]);
            counts[j] += static_cast<std::uint64_t>(share_raw(opened));
        }
    }
    return counts;
}

template <typename Key0T, typename Key1T, typename InputT>
void retain_all(std::vector<idpf_eval_ctx<Key0T>> & ctx0,
    std::vector<idpf_eval_ctx<Key1T>> & ctx1, InputT prefix)
{
    for (std::size_t i = 0; i < ctx0.size(); ++i)
    {
        ctx0[i].retain(prefix);
        ctx1[i].retain(prefix);
    }
}

}  // namespace idpf_agg_detail
}  // namespace detail

/// @brief Walk both parties' idpf contexts to the numeric maximum.
/// @details At each bit, keep the `1`-child when it holds any mass; otherwise
///          keep the `0`-child. That is the MSB-first spine of the max value.
/// \complexity O(n · N · depth) party work for `n` keys on an `N`-bit domain:
///          two `eval_until` calls per key per level (both children), then
///          `retain`. Communication of opened counts is O(depth), independent
///          of `n` (ePrint 2024/1190).
template <typename Key0T, typename Key1T>
HEDLEY_WARN_UNUSED_RESULT
auto idpf_agg_max(const std::vector<Key0T> & keys0,
    const std::vector<Key1T> & keys1)
{
    using key_type = unwrap_party_key_t<Key0T>;
    using input_type = typename key_type::input_type;
    constexpr auto bits = key_type::input_bits;
    static_assert(std::is_same_v<key_type, unwrap_party_key_t<Key1T>>,
        "idpf_agg_max: party keys must unwrap to the same DPF type");
    if (keys0.size() != keys1.size() || keys0.empty())
        throw std::invalid_argument("idpf_agg_max: nonempty equal key lists");

    std::vector<idpf_eval_ctx<Key0T>> ctx0;
    std::vector<idpf_eval_ctx<Key1T>> ctx1;
    ctx0.reserve(keys0.size());
    ctx1.reserve(keys1.size());
    for (std::size_t i = 0; i < keys0.size(); ++i)
    {
        ctx0.emplace_back(keys0[i]);
        ctx1.emplace_back(keys1[i]);
    }

    input_type prefix{0};
    for (std::size_t level = 1; level <= bits; ++level)
    {
        const input_type left = static_cast<input_type>(prefix << 1);
        const input_type right = static_cast<input_type>((prefix << 1) | 1);
        std::vector<input_type> prefs{left, right};
        auto counts = detail::idpf_agg_detail::open_prefix_counts(ctx0, ctx1,
            level, prefs);
        prefix = (counts[1] > 0) ? right : left;
        detail::idpf_agg_detail::retain_all(ctx0, ctx1, prefix);
    }
    return prefix;
}

/// @brief Walk both parties' idpf contexts to the k-th largest (1-based).
/// @details `k == 1` is the maximum; `k == n` is the minimum. At each bit,
///          take the `1`-child when its count is at least `k`; otherwise
///          subtract that count and take the `0`-child.
/// \complexity Same as `idpf_agg_max`: O(n · N · depth) local work,
///          O(depth) opened counts.
template <typename Key0T, typename Key1T>
HEDLEY_WARN_UNUSED_RESULT
auto idpf_agg_kth(const std::vector<Key0T> & keys0,
    const std::vector<Key1T> & keys1, std::size_t k)
{
    using key_type = unwrap_party_key_t<Key0T>;
    using input_type = typename key_type::input_type;
    constexpr auto bits = key_type::input_bits;
    static_assert(std::is_same_v<key_type, unwrap_party_key_t<Key1T>>,
        "idpf_agg_kth: party keys must unwrap to the same DPF type");
    if (keys0.size() != keys1.size() || keys0.empty())
        throw std::invalid_argument("idpf_agg_kth: nonempty equal key lists");
    if (k == 0 || k > keys0.size())
        throw std::invalid_argument("idpf_agg_kth: k out of range");

    std::vector<idpf_eval_ctx<Key0T>> ctx0;
    std::vector<idpf_eval_ctx<Key1T>> ctx1;
    ctx0.reserve(keys0.size());
    ctx1.reserve(keys1.size());
    for (std::size_t i = 0; i < keys0.size(); ++i)
    {
        ctx0.emplace_back(keys0[i]);
        ctx1.emplace_back(keys1[i]);
    }

    input_type prefix{0};
    for (std::size_t level = 1; level <= bits; ++level)
    {
        const input_type left = static_cast<input_type>(prefix << 1);
        const input_type right = static_cast<input_type>((prefix << 1) | 1);
        std::vector<input_type> prefs{left, right};
        auto counts = detail::idpf_agg_detail::open_prefix_counts(ctx0, ctx1,
            level, prefs);
        if (counts[1] >= k)
            prefix = right;
        else
        {
            k -= counts[1];
            prefix = left;
        }
        detail::idpf_agg_detail::retain_all(ctx0, ctx1, prefix);
    }
    return prefix;
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_IDPF_AGG_HPP__