/// @file dpf/shuffle.hpp
/// @brief Honest-majority 3PC shuffle and boolean select / aggregates.
/// @details `shuffle_party` applies one permutation known to every party.
///          `shuffle_hidden_pass` is the hidden order: three passes, each
///          from one pairwise seed, leaving that seed's missing party out.
///          Passes run left-out 2 (`k01`), then 0 (`k12`), then 1 (`k20`).
///          `permute` is `out[i] = in[pi[i]]`, and each pass applies that map,
///          so the clear replay is `permute(permute(permute(v, π01), π12), π20)`.
///          A secret index stays a DPF. This header reorders a column of
///          shares the parties already hold. The manual is
///          [An array of shares, not a secret index](@ref share_shuffle).
#ifndef LIBDPF_INCLUDE_DPF_SHUFFLE_HPP__
#define LIBDPF_INCLUDE_DPF_SHUFFLE_HPP__

#include <algorithm>
#include <cstddef>
#include <cstdint>
#include <numeric>
#include <stdexcept>
#include <utility>
#include <vector>

#include "hedley/hedley.h"

#include "dpf/bit_inject.hpp"
#include "dpf/buffered_prg.hpp"
#include "dpf/prg_aes.hpp"
#include "dpf/random.hpp"
#include "dpf/rss_seed.hpp"
#include "dpf/share_cmp.hpp"

namespace dpf
{
namespace shuffle
{

/// @brief Fisher–Yates permutation from one pairwise seed. Lane 0.
HEDLEY_WARN_UNUSED_RESULT
inline std::vector<std::size_t> permutation_from_seed(
    rss::seed_block seed, std::size_t n, std::uint64_t index)
{
    std::vector<std::size_t> pi(n);
    std::iota(pi.begin(), pi.end(), 0);
    randomness::lane_table<std::uint64_t, prg::aes128> table(seed);
    for (std::size_t i = n; i > 1; --i)
    {
        const auto r = table.value_at(0, index + i);
        const std::size_t j = static_cast<std::size_t>(r % i);
        std::swap(pi[i - 1], pi[j]);
    }
    return pi;
}

/// @brief Derive a shared permutation of `n` from pairwise seeds (dealer view).
HEDLEY_WARN_UNUSED_RESULT
inline std::vector<std::size_t> permutation_from_seeds(
    const rss::seed_bundle & bundle, std::size_t n, std::uint64_t index)
{
    return permutation_from_seed(bundle.k01, n, index);
}

/// @brief Apply permutation to a clear vector.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
std::vector<T> permute(const std::vector<T> & v,
    const std::vector<std::size_t> & pi)
{
    if (v.size() != pi.size())
        throw std::invalid_argument("shuffle permute");
    std::vector<T> out(v.size());
    for (std::size_t i = 0; i < pi.size(); ++i)
        out[i] = v[pi[i]];
    return out;
}

/// @brief Multiset oracle: permute a clear vector (not a party protocol).
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
std::vector<T> shuffle_clear(const std::vector<T> & values,
    const rss::seed_bundle & bundle, std::uint64_t index)
{
    auto pi = permutation_from_seeds(bundle, values.size(), index);
    return permute(values, pi);
}

/// @brief One 3PC party's shuffle of RSS components.
/// @details Party holds replicated vectors `(own[i], next[i])`. Permute both
///          by the shared `pi`, then the caller `send_next`s the permuted `own`
///          and refreshes with the received previous component.
template <typename T>
struct shuffle_party_view
{
    std::vector<T> own;
    std::vector<T> next;
};

template <typename T>
HEDLEY_WARN_UNUSED_RESULT
shuffle_party_view<T> shuffle_party(const shuffle_party_view<T> & in,
    const rss::seed_bundle & bundle, std::uint64_t index)
{
    auto pi = permutation_from_seeds(bundle, in.own.size(), index);
    if (in.next.size() != in.own.size())
        throw std::invalid_argument("shuffle_party size");
    shuffle_party_view<T> out;
    out.own = permute(in.own, pi);
    out.next = permute(in.next, pi);
    return out;
}

/// @brief Party `me` sends to `(me + 1) mod 3` on a hidden pass, or to nobody.
inline unsigned hidden_next(unsigned me) noexcept
{
    return (me + 1u) % 3u;
}

/// @brief The party who forms `u = own + next` on the pass that leaves `left_out` out.
inline unsigned hidden_u_party(unsigned left_out) noexcept
{
    return (left_out + 1u) % 3u;
}

/// @brief The party who permutes its `next` component on that pass.
inline unsigned hidden_side_party(unsigned left_out) noexcept
{
    return (left_out + 2u) % 3u;
}

/// @brief One directed array from `shuffle_hidden_pass`.
template <typename T>
struct hidden_message
{
    std::vector<T> data;
    unsigned to = 0;
    bool sends = false;
};

/// @brief Local view after one pass, plus the array this party sends to `next`.
template <typename T>
struct hidden_pass_result
{
    shuffle_party_view<T> view;
    hidden_message<T> out;
};

namespace detail
{

inline rss::seed_block seed_toward(const rss::party_seeds & seeds, unsigned other)
{
    if (other == hidden_next(seeds.me))
        return seeds.with_next;
    if (other == hidden_next(hidden_next(seeds.me)))
        return seeds.with_prev;
    throw std::invalid_argument("shuffle: seed");
}

template <typename T>
std::vector<T> words_from_seed(rss::seed_block seed, std::uint32_t lane,
    std::size_t n, std::uint64_t index)
{
    randomness::lane_table<std::uint64_t, prg::aes128> table(seed);
    std::vector<T> out(n);
    for (std::size_t i = 0; i < n; ++i)
        out[i] = static_cast<T>(table.value_at(lane, index + i));
    return out;
}

}  // namespace detail

/// @brief One hidden-shuffle pass on replicated `(own, next)`.
/// @details Passes run with `left_out` 2, then 0, then 1. Parties who share
///          the pass seed permute a two-party split; the left-out party
///          receives a fresh component. `inbound` is required for the side
///          party and the left-out party (the array their predecessor sends).
///          The u-party ignores `inbound`. Pad words use lane 1 of the
///          permutation seed. The fresh mask uses lane 2 of the seed shared
///          with the left-out party, so it does not collide with that seed's
///          permutation lane.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
hidden_pass_result<T> shuffle_hidden_pass(unsigned me,
    const rss::party_seeds & seeds, const shuffle_party_view<T> & in,
    std::uint64_t index, unsigned left_out, const std::vector<T> * inbound)
{
    if (me > 2 || left_out > 2)
        throw std::invalid_argument("shuffle: party");
    if (seeds.me != me)
        throw std::invalid_argument("shuffle: party seeds");
    if (in.own.size() != in.next.size())
        throw std::invalid_argument("shuffle: size");
    const std::size_t n = in.own.size();
    const unsigned u = hidden_u_party(left_out);
    const unsigned side = hidden_side_party(left_out);
    hidden_pass_result<T> step;
    if (me == u)
    {
        const auto pi_seed = detail::seed_toward(seeds, side);
        const auto mask_seed = detail::seed_toward(seeds, left_out);
        const auto pi = permutation_from_seed(pi_seed, n, index);
        const auto pad = detail::words_from_seed<T>(pi_seed, 1, n, index);
        const auto mask = detail::words_from_seed<T>(mask_seed, 2, n, index);
        std::vector<T> sum(n);
        for (std::size_t i = 0; i < n; ++i)
            sum[i] = static_cast<T>(in.own[i] + in.next[i] + pad[i]);
        const auto moved = permute(sum, pi);
        step.view.own = mask;
        step.view.next.resize(n);
        step.out.data.resize(n);
        for (std::size_t i = 0; i < n; ++i)
        {
            step.view.next[i] = static_cast<T>(moved[i] - mask[i]);
            step.out.data[i] = step.view.next[i];
        }
        step.out.to = side;
        step.out.sends = true;
        return step;
    }
    if (me == side)
    {
        if (inbound == nullptr || inbound->size() != n)
            throw std::invalid_argument("shuffle: inbound");
        const auto pi_seed = detail::seed_toward(seeds, u);
        const auto pi = permutation_from_seed(pi_seed, n, index);
        const auto pad = detail::words_from_seed<T>(pi_seed, 1, n, index);
        std::vector<T> side_v(n);
        for (std::size_t i = 0; i < n; ++i)
            side_v[i] = static_cast<T>(in.next[i] - pad[i]);
        step.view.next = permute(side_v, pi);
        step.view.own = *inbound;
        step.out.data = step.view.next;
        step.out.to = left_out;
        step.out.sends = true;
        return step;
    }
    if (inbound == nullptr || inbound->size() != n)
        throw std::invalid_argument("shuffle: inbound");
    const auto mask_seed = detail::seed_toward(seeds, u);
    step.view.next = detail::words_from_seed<T>(mask_seed, 2, n, index);
    step.view.own = *inbound;
    step.out.sends = false;
    return step;
}

/// @brief Three hidden passes. Opens by summing the three `own` components.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
std::vector<T> shuffle_hidden_triple(const std::vector<T> & clear,
    const rss::seed_bundle & bundle, std::uint64_t index)
{
    const std::size_t n = clear.size();
    shuffle_party_view<T> held[3];
    for (unsigned p = 0; p < 3; ++p)
    {
        held[p].own.assign(n, T{});
        held[p].next.assign(n, T{});
    }
    for (std::size_t i = 0; i < n; ++i)
    {
        const T a = dpf::uniform_sample<T>();
        const T b = dpf::uniform_sample<T>();
        const T c = static_cast<T>(clear[i] - a - b);
        held[0].own[i] = a;
        held[0].next[i] = b;
        held[1].own[i] = b;
        held[1].next[i] = c;
        held[2].own[i] = c;
        held[2].next[i] = a;
    }
    const unsigned order[3] = {2u, 0u, 1u};
    for (unsigned left : order)
    {
        const unsigned u = hidden_u_party(left);
        const unsigned side = hidden_side_party(left);
        const auto seeds_u = rss::party_seeds::from_bundle(bundle, u);
        const auto seeds_s = rss::party_seeds::from_bundle(bundle, side);
        const auto seeds_l = rss::party_seeds::from_bundle(bundle, left);
        auto u_step = shuffle_hidden_pass<T>(u, seeds_u, held[u], index, left, nullptr);
        auto s_step = shuffle_hidden_pass<T>(
            side, seeds_s, held[side], index, left, &u_step.out.data);
        auto l_step = shuffle_hidden_pass<T>(
            left, seeds_l, held[left], index, left, &s_step.out.data);
        held[u] = std::move(u_step.view);
        held[side] = std::move(s_step.view);
        held[left] = std::move(l_step.view);
    }
    std::vector<T> out(n);
    for (std::size_t i = 0; i < n; ++i)
        out[i] = static_cast<T>(held[0].own[i] + held[1].own[i] + held[2].own[i]);
    return out;
}

/// @brief After ring send of `own`: party stores `(own, recv_from_prev)`.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
shuffle_party_view<T> shuffle_refresh(std::vector<T> own,
    std::vector<T> from_prev)
{
    if (own.size() != from_prev.size())
        throw std::invalid_argument("shuffle_refresh size");
    return shuffle_party_view<T>{std::move(own), std::move(from_prev)};
}

/// @brief Dealer helper: run three parties' permute + ring refresh; open sum.
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
std::vector<T> shuffle_party_triple(const std::vector<T> & clear,
    const rss::seed_bundle & bundle, std::uint64_t index)
{
    // Share as RSS: p0=(v,0), p1=(0,0), p2=(0,v) component-wise is wrong for
    // sum; use additive split into three: sample r0,r1, r2=v-r0-r1.
    const std::size_t n = clear.size();
    shuffle_party_view<T> v0{std::vector<T>(n), std::vector<T>(n)};
    shuffle_party_view<T> v1{std::vector<T>(n), std::vector<T>(n)};
    shuffle_party_view<T> v2{std::vector<T>(n), std::vector<T>(n)};
    for (std::size_t i = 0; i < n; ++i)
    {
        const T a = dpf::uniform_sample<T>();
        const T b = dpf::uniform_sample<T>();
        const T c = static_cast<T>(clear[i] - a - b);
        // p0:(a,b) p1:(b,c) p2:(c,a)
        v0.own[i] = a;
        v0.next[i] = b;
        v1.own[i] = b;
        v1.next[i] = c;
        v2.own[i] = c;
        v2.next[i] = a;
    }
    auto p0 = shuffle_party(v0, bundle, index);
    auto p1 = shuffle_party(v1, bundle, index);
    auto p2 = shuffle_party(v2, bundle, index);
    // Ring send own → next party; refresh: (own, from_prev) with from_prev =
    // previous party's own.
    auto r0 = shuffle_refresh(p0.own, p2.own);
    auto r1 = shuffle_refresh(p1.own, p0.own);
    auto r2 = shuffle_refresh(p2.own, p1.own);
    std::vector<T> out(n);
    for (std::size_t i = 0; i < n; ++i)
        out[i] = static_cast<T>(r0.own[i] + r1.own[i] + r2.own[i]);
    return out;
}

/// @brief Multiset equality (permutation check).
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
bool is_permutation_of(std::vector<T> a, std::vector<T> b)
{
    if (a.size() != b.size())
        return false;
    std::sort(a.begin(), a.end());
    std::sort(b.begin(), b.end());
    return a == b;
}

/// @brief select(pred, row): inject boolean mask onto each column (clear).
template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
std::vector<Ring> select_clear(const std::vector<std::uint8_t> & pred,
    const std::vector<Ring> & row)
{
    if (pred.size() != row.size())
        throw std::invalid_argument("select size");
    std::vector<Ring> out(row.size());
    for (std::size_t i = 0; i < row.size(); ++i)
        out[i] = pred[i] ? row[i] : Ring{};
    return out;
}

template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
Ring sum_if_clear(const std::vector<std::uint8_t> & pred,
    const std::vector<Ring> & row)
{
    Ring s{};
    auto selected = select_clear(pred, row);
    for (auto v : selected)
        s = static_cast<Ring>(s + v);
    return s;
}

template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
Ring count_if_clear(const std::vector<std::uint8_t> & pred)
{
    Ring c{};
    for (auto b : pred)
        c = static_cast<Ring>(c + static_cast<Ring>(b & 1u));
    return c;
}

/// @brief Argmax over a clear vector: value and one-hot index.
template <typename Ring>
struct argmax_result
{
    Ring value{};
    std::vector<std::uint8_t> one_hot;
};

template <typename Ring>
HEDLEY_WARN_UNUSED_RESULT
argmax_result<Ring> argmax_clear(const std::vector<Ring> & v, unsigned nbits)
{
    if (v.empty())
        throw std::invalid_argument("argmax empty");
    argmax_result<Ring> out;
    out.one_hot.assign(v.size(), 0);
    std::size_t best = 0;
    for (std::size_t i = 1; i < v.size(); ++i)
        if (share_cmp::gt_clear(v[i], v[best], nbits))
            best = i;
    out.value = v[best];
    out.one_hot[best] = 1;
    return out;
}

/// @brief Short 2PC Waksman-style shuffle via mux tree (clear oracle).
template <typename T>
HEDLEY_WARN_UNUSED_RESULT
std::vector<T> waksman_clear(const std::vector<T> & v,
    const std::vector<std::uint8_t> & swap_bits)
{
    auto out = v;
    std::size_t bit = 0;
    for (std::size_t len = 2; len <= out.size(); len *= 2)
    {
        for (std::size_t i = 0; i + len / 2 < out.size(); i += len)
        {
            for (std::size_t j = 0; j < len / 2 && i + j + len / 2 < out.size();
                 ++j)
            {
                if (bit < swap_bits.size() && swap_bits[bit])
                    std::swap(out[i + j], out[i + j + len / 2]);
                ++bit;
            }
        }
    }
    return out;
}

}  // namespace shuffle
}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_SHUFFLE_HPP__