/// @file dpf/eval_point.hpp
/// @brief Evaluate one DPF input.
/// @details `eval_point(key, x)` returns a handle; `*handle` is that party's
///          share of output 0. `eval_point<I>` selects another output.
///          `eval_point<I0, I1, ...>` returns a tuple of shares.
///          Pass a `basic_path_memoizer` lvalue to resume a previous path.
///          An unassigned wildcard output throws `std::runtime_error`.
///          `eval_point(key, x, dpf::prove(π))` folds a VDPF proof token.
///          @snippet evaluation/eval_point.cpp eval-point
/// @author Ryan Henry <ryan.henry@ucalgary.ca>
/// @author Christopher Jiang <christopher.jiang@ucalgary.ca>
/// @copyright Copyright (c) 2019-2024 Ryan Henry and [others](@ref authors)
/// @license Released under a GNU General Public v2.0 (GPLv2) license;
///          see [LICENSE.md](@ref license) for details.

#ifndef LIBDPF_INCLUDE_DPF_EVAL_POINT_HPP__
#define LIBDPF_INCLUDE_DPF_EVAL_POINT_HPP__

#include <portable-snippets/builtin/builtin.h>
#include "hedley/hedley.h"

#include <cstddef>
#include <tuple>

#include "dpf/dpf_key.hpp"
#include "dpf/eval_common.hpp"
#include "dpf/eval_target.hpp"
#include "dpf/path_memoizer.hpp"
#include "dpf/verifiable.hpp"

namespace dpf
{

namespace internal
{

template <typename DpfKey,
          typename InputT,
          typename PathMemoizer>
inline auto eval_point_interior(const DpfKey & dpf, InputT && x, PathMemoizer && path,
    proof_token * pi = nullptr)
{
    using dpf_type = DpfKey;

    auto level_index = detail::path_resume_for_level(path, dpf, x, dpf.depth);

    // Same guard as `ensure_level`: skip the `msb_mask` shift when already done.
    if (level_index <= dpf.depth)
    {
        DPF_UNROLL_LOOP
        for (auto mask = dpf.msb_mask>>(level_index-1);
            level_index <= dpf.depth; ++level_index, mask>>=1)
        {
            bool bit = !!(mask & x);
            auto cw = dpf.correction_word(level_index-1, bit);
            const bool is_last = dpf_type::tree::is_last_level(level_index - 1,
                dpf.depth);
            path[level_index] = dpf_type::traverse_interior(path[level_index-1],
                cw, bit, is_last);
            if constexpr (dpf_type::is_verifiable)
            {
                if (pi != nullptr)
                {
                    const auto x_bits = static_cast<psnip_uint64_t>(
                        utils::to_integral_type<std::decay_t<InputT>>{}(x)
                        >> (utils::bitlength_of_v<std::decay_t<InputT>>
                            - level_index));
                    detail::vdpf::fold_node(*pi, level_index - 1, x_bits,
                        path[level_index],
                        dpf.correction_seeds()[level_index - 1]);
                }
            }
        }
    }
    detail::path_note_filled_to(path, dpf.depth);
}

template <std::size_t I,
          typename DpfKey,
          typename PathMemoizer>
inline auto eval_point_exterior(const DpfKey & dpf, PathMemoizer && path)
{
    assert_not_wildcard_output<I>(dpf);

    auto interior = path[dpf.depth];
    return dpf.template traverse_exterior<I>(interior);
}

template <std::size_t I,
          typename DpfKey,
          typename InputT,
          typename PathMemoizer>
HEDLEY_ALWAYS_INLINE
auto eval_point(const DpfKey & dpf, InputT && x, PathMemoizer && path,
    proof_token * pi = nullptr)
{
    utils::flip_msb_if_signed_integral(x);
    internal::eval_point_interior(dpf, x, path, pi);
    return internal::eval_point_exterior<I>(dpf, path);
}

}  // namespace internal

/// Evaluate output `I` at `x`.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <std::size_t I = 0,
          typename DpfKey,
          typename InputT,
          typename PathMemoizer = dpf::nonmemoizing_path_memoizer<DpfKey>,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey> && !is_multilevel_key_v<DpfKey>, bool> = true>
HEDLEY_ALWAYS_INLINE
auto eval_point(const DpfKey & dpf, InputT && x, PathMemoizer && path = PathMemoizer{})
{
    assert_not_wildcard_output<I>(dpf);
    using output_type = typename DpfKey::concrete_output_type<I>;

    auto tx = dpf.offset_x(x);
    return make_eval_dpf_output<DpfKey, output_type>(
        internal::eval_point<I>(dpf, tx, path), tx);
}

/// Evaluate at `x`, folding newly walked nodes into an existing proof token.
/// @details Does not `init_proof`; used by bulk sequence evals that share a path.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <std::size_t I = 0,
          typename DpfKey,
          typename InputT,
          typename PathMemoizer,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey> && !is_multilevel_key_v<DpfKey>, bool> = true>
HEDLEY_ALWAYS_INLINE
auto eval_point(const DpfKey & dpf, InputT && x, PathMemoizer && path,
    proof_token * pi)
{
    assert_not_wildcard_output<I>(dpf);
    using output_type = typename DpfKey::concrete_output_type<I>;

    auto tx = dpf.offset_x(x);
    return make_eval_dpf_output<DpfKey, output_type>(
        internal::eval_point<I>(dpf, tx, path, pi), tx);
}

/// Evaluate and fold a VDPF proof token for the walked path.
/// @details A fresh token always folds every node on the path, including
///          nodes already stored in a warm path memoizer (hash the cached
///          seeds; do not re-expand the PRG). Callers that continue one
///          accumulator across many points use the `proof_token *` overload
///          without `init_proof`.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <std::size_t I = 0,
          typename DpfKey,
          typename InputT,
          typename PathMemoizer = dpf::nonmemoizing_path_memoizer<DpfKey>,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey>, bool> = true>
HEDLEY_ALWAYS_INLINE
auto eval_point(const DpfKey & dpf, InputT && x, prove_ref pr,
    PathMemoizer && path = PathMemoizer{})
{
    static_assert(DpfKey::is_verifiable,
        "eval_point(..., prove(π)): key must carry dpf::verifiable");
    detail::vdpf::init_proof(pr.token, dpf);
    auto walk_x = dpf.offset_x(x);
    utils::flip_msb_if_signed_integral(walk_x);
    if constexpr (DpfKey::is_verifiable)
    {
        const auto resume = detail::path_resume_for_level(
            path, dpf, walk_x, dpf.depth);
        for (std::size_t level_index = 1; level_index < resume; ++level_index)
        {
            const auto x_bits = static_cast<psnip_uint64_t>(
                utils::to_integral_type<std::decay_t<decltype(walk_x)>>{}(
                    walk_x)
                >> (utils::bitlength_of_v<std::decay_t<decltype(walk_x)>>
                    - level_index));
            detail::vdpf::fold_node(pr.token, level_index - 1, x_bits,
                path[level_index],
                dpf.correction_seeds()[level_index - 1]);
        }
    }
    auto out = eval_point<I>(dpf, std::forward<InputT>(x),
        std::forward<PathMemoizer>(path), &pr.token);
    detail::vdpf::fold_output_binding(pr.token, dpf);
    return out;
}

/// Evaluate and fold the output into a weight-1 sketch.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <std::size_t I = 0,
          typename DpfKey,
          typename InputT,
          typename PathMemoizer = dpf::nonmemoizing_path_memoizer<DpfKey>,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey>, bool> = true>
HEDLEY_ALWAYS_INLINE
auto eval_point(const DpfKey & dpf, InputT && x, sketch_ref & sk,
    PathMemoizer && path = PathMemoizer{})
{
    static_assert(DpfKey::is_extractable,
        "eval_point(..., sketch(σ)): key must carry dpf::extractable");
    auto out = eval_point<I>(dpf, std::forward<InputT>(x),
        std::forward<PathMemoizer>(path));
    if constexpr (DpfKey::is_extractable)
        sk.absorb((*out).raw());
    return out;
}

/// @brief Rvalue convenience for a one-shot `sketch(local, rs)` temporary.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <std::size_t I = 0,
          typename DpfKey,
          typename InputT,
          typename PathMemoizer = dpf::nonmemoizing_path_memoizer<DpfKey>,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey>, bool> = true>
HEDLEY_ALWAYS_INLINE
auto eval_point(const DpfKey & dpf, InputT && x, sketch_ref && sk,
    PathMemoizer && path = PathMemoizer{})
{
    return eval_point<I>(dpf, std::forward<InputT>(x), sk,
        std::forward<PathMemoizer>(path));
}

/// Evaluate several outputs at `x`.
/// \complexity O(n) time. n is `depth`. One interior traversal per level from the memoizer resume index through the leaf. Extra space is the path memoizer (O(n) nodes, or one node if it does not memoize). A proof token adds one fold per level walked.
template <std::size_t I0,
          std::size_t I1,
          std::size_t ...Is,
          typename DpfKey,
          typename InputT,
          typename PathMemoizer = dpf::basic_path_memoizer<DpfKey>,
          std::enable_if_t<looks_like_dpf_key_v<DpfKey> && !is_multilevel_key_v<DpfKey>, bool> = true>
HEDLEY_ALWAYS_INLINE
auto eval_point(const DpfKey & dpf, InputT && x, PathMemoizer && path = PathMemoizer{})
{
    return std::make_tuple(
        *eval_point<I0>(dpf, x, path),
        *eval_point<I1>(dpf, x, path),
        *eval_point<Is>(dpf, x, path)...);
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_EVAL_POINT_HPP__