#include <cstdint>
#include <iostream>
#include <iterator>
#include <vector>

#include "dpf.hpp"
#include "dpf/app_flow.hpp"
#include "dpf/app_plans.hpp"

// MPC SUBLEQ, the DPF steps of one instruction (Jiang and Henry).
// Offline: expand wildcard unit keys with `defer_eval_full` before the
// addresses are known. Online: assign each address into `offset_x`, read
// by rotating the prepaid buffer (no second AES pass), write by scaling
// the same `e_B` view, and branch with a path evaluation of `x ≤ 0`.
//
// Instruction fetch is the same prepaid unit dotted against three sliding
// windows of D; this listing starts after (A, B, C) are already shares.
//
//   c++ -std=c++17 -march=native -I include -I thirdparty \
//       examples/applications/subleq.cpp

namespace
{

using addr_t = std::uint8_t;
using word_t = std::uint32_t;
constexpr std::size_t n = 256;

template <typename Key0, typename Key1, typename T>
void assign_input(Key0 & k0, Key1 & k1, T alpha)
{
    const T a0 = static_cast<T>(0x12);
    const T a1 = static_cast<T>(alpha - a0);
    const auto s0 = k0.offset_x.compute_and_get_share(a0);
    const auto s1 = k1.offset_x.compute_and_get_share(a1);
    k0.offset_x.reconstruct(s1);
    k1.offset_x.reconstruct(s0);
}

/// Opened unit · public memory over `[0, n)`.
template <typename View0, typename View1>
word_t dot_prefix(View0 && v0, View1 && v1, const std::vector<word_t> & mem)
{
    word_t acc = 0;
    auto it0 = std::begin(v0);
    auto it1 = std::begin(v1);
    for (std::size_t i = 0; i < n; ++i, ++it0, ++it1)
        acc += dpf::reconstruct(*it0, *it1) * mem[i];
    return acc;
}

template <typename View0, typename View1>
void add_scaled_prefix(std::vector<word_t> & mem, View0 && v0, View1 && v1,
    word_t scale)
{
    auto it0 = std::begin(v0);
    auto it1 = std::begin(v1);
    for (std::size_t i = 0; i < n; ++i, ++it0, ++it1)
        mem[i] += scale * dpf::reconstruct(*it0, *it1);
}

}  // namespace

int main()
{
    // Two's-complement words in a uint32_t container (same bits as int32_t).
    constexpr addr_t A = 3;
    constexpr addr_t B = 7;
    constexpr addr_t C = 2;
    constexpr addr_t pc = 0;
    std::vector<word_t> D(n);
    D[A] = 5;
    D[B] = 3;   // after SUBLEQ: D[B] = 3 - 5 = -2 ≤ 0 → pc' = C

    // --- Offline: wildcard unit keys, full-domain expand at identity -----
    auto [kA0, kA1] = dpf::make_dpf(dpf::wildcard_value<addr_t>{}, word_t{1});
    auto [kB0, kB1] = dpf::make_dpf(dpf::wildcard_value<addr_t>{}, word_t{1});
    auto bufA0 = dpf::make_output_buffer_for_full(kA0);
    auto bufA1 = dpf::make_output_buffer_for_full(kA1);
    auto bufB0 = dpf::make_output_buffer_for_full(kB0);
    auto bufB1 = dpf::make_output_buffer_for_full(kB1);
    auto defA0 = dpf::defer_eval_full(kA0, bufA0);
    auto defA1 = dpf::defer_eval_full(kA1, bufA1);
    auto defB0 = dpf::defer_eval_full(kB0, bufB0);
    auto defB1 = dpf::defer_eval_full(kB1, bufB1);

    // --- Online: open addresses, rotate prepaid unit vectors -------------
    assign_input(kA0, kA1, A);
    assign_input(kB0, kB1, B);

    const word_t DA = dot_prefix(defA0.get(), defA1.get(), D);
    const word_t DB = dot_prefix(defB0.get(), defB1.get(), D);
    if (DA != D[A] || DB != D[B])
    {
        std::cerr << "subleq read\n";
        return 1;
    }

    const word_t x = static_cast<word_t>(DB - DA);  // wraps to -2 as uint32_t

    // Write D[B] ← D[B] - D[A] by adding (-DA) · e_B. The protocol Beavers
    // the scale; the opened -DA stands in here.
    add_scaled_prefix(D, defB0.get(), defB1.get(), static_cast<word_t>(-DA));
    if (D[B] != static_cast<word_t>(3 - 5) || D[A] != 5)
    {
        std::cerr << "subleq write\n";
        return 1;
    }

    // Branch: path eval only — never expand the word-domain key.
    // `leq` at knot 0, evaluated at x: 1 iff x ≤ 0 in signed order.
    auto [kZ0, kZ1] = dpf::make_dpf(std::int32_t{0}, dpf::leq(std::uint64_t{1}));
    const auto b = dpf::reconstruct(
        dpf::eval_point(dpf::cmp, kZ0, static_cast<std::int32_t>(x)),
        dpf::eval_point(dpf::cmp, kZ1, static_cast<std::int32_t>(x)));
    const addr_t pc_next = b ? C : static_cast<addr_t>(pc + 3);
    if (b != 1 || pc_next != C)
    {
        std::cerr << "subleq branch\n";
        return 1;
    }

    {
        if (int rc = dpf::app::run_measured("subleq",
                dpf::protocol::subleq_instruction_plan(0), 8))
            return rc;
    }

    std::cout << static_cast<std::int32_t>(D[B]) << " "
              << static_cast<unsigned>(pc_next) << "\n";
    return 0;
}