/// @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);

    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`.
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 and fold a VDPF proof token for the walked path.
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");
    assert_not_wildcard_output<I>(dpf);
    using output_type = typename DpfKey::concrete_output_type<I>;

    detail::vdpf::init_proof(pr.token, dpf);
    auto tx = dpf.offset_x(x);
    return make_eval_dpf_output<DpfKey, output_type>(
        internal::eval_point<I>(dpf, tx, path, &pr.token), tx);
}

/// Evaluate several outputs at `x`.
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)...);
}

/// Fold every point in `[from, to]` into `pi` (caller must `init_proof` first,
/// or pass a fresh token via `prove_interval` below).
template <typename KeyT, typename InputT>
void prove_fold_interval(const KeyT & key, InputT from, InputT to,
    proof_token & pi)
{
    static_assert(KeyT::is_verifiable,
        "prove_fold_interval: key must carry dpf::verifiable");
    using input_type = typename KeyT::input_type;
    auto cur = static_cast<input_type>(from);
    const auto last = static_cast<input_type>(to);
    for (;;)
    {
        nonmemoizing_path_memoizer<KeyT> path{};
        auto tx = key.offset_x(cur);
        utils::flip_msb_if_signed_integral(tx);
        internal::eval_point_interior(key, tx, path, &pi);
        if (cur == last)
            break;
        ++cur;
    }
}

/// Initialise `pr.token` and fold `[from, to]`.
template <typename KeyT, typename InputT>
void prove_interval(const KeyT & key, InputT from, InputT to, prove_ref pr)
{
    detail::vdpf::init_proof(pr.token, key);
    prove_fold_interval(key, from, to, pr.token);
}

}  // namespace dpf

#endif  // LIBDPF_INCLUDE_DPF_EVAL_POINT_HPP__