#include <array>
#include <cstdint>
#include <iostream>

#include "dpf.hpp"

/// Memoizers are workspaces of interior nodes. Pass a mutable lvalue.
/// A temporary (including the default argument) cannot remember a prefix.
int main()
{
    const std::uint8_t alpha = 42;
    const std::uint64_t beta = 7;
    auto [k0, k1] = dpf::make_dpf(alpha, beta);
    using key_t = dpf::unwrap_party_key_t<std::decay_t<decltype(k0)>>;

    //! [path-memoizer]
    auto path0 = dpf::make_basic_path_memoizer<key_t>();
    auto path1 = dpf::make_basic_path_memoizer<key_t>();
    for (int x = 0; x < 256; ++x)
    {
        auto y0 = *dpf::eval_point(k0, static_cast<std::uint8_t>(x), path0);
        auto y1 = *dpf::eval_point(k1, static_cast<std::uint8_t>(x), path1);
        std::uint64_t got = dpf::reconstruct(y0, y1);
        std::uint64_t expect = (static_cast<std::uint8_t>(x) == alpha) ? beta : 0;
        if (got != expect)
        {
            std::cerr << "basic_path_memoizer\n";
            return 1;
        }
    }
    //! [path-memoizer]

    // One node, no prefix reuse. Correct for a single query.
    auto once0 = dpf::make_nonmemoizing_path_memoizer<key_t>();
    auto once1 = dpf::make_nonmemoizing_path_memoizer<key_t>();
    if (dpf::reconstruct(*dpf::eval_point(k0, alpha, once0),
            *dpf::eval_point(k1, alpha, once1)) != beta)
    {
        std::cerr << "nonmemoizing_path_memoizer\n";
        return 1;
    }

    //! [interval-memoizer]
    const std::uint8_t from = 40;
    const std::uint8_t to = 50;
    // Sized for [from, to]. A wider interval throws std::length_error.
    // `make_full_tree_interval_memoizer` keeps every level instead of two.
    auto memo0 = dpf::make_basic_interval_memoizer<key_t>(from, to);
    auto memo1 = dpf::make_basic_interval_memoizer<key_t>(from, to);
    auto [ibuf0, i0] = dpf::eval_interval(k0, from, to, memo0);
    auto [ibuf1, i1] = dpf::eval_interval(k1, from, to, memo1);
    //! [interval-memoizer]

    auto it0 = std::begin(i0);
    auto it1 = std::begin(i1);
    for (std::uint8_t x = from; x <= to; ++x, ++it0, ++it1)
    {
        std::uint64_t got = dpf::reconstruct(*it0, *it1);
        std::uint64_t expect = (x == alpha) ? beta : 0;
        if (got != expect)
        {
            std::cerr << "basic_interval_memoizer\n";
            return 1;
        }
    }

    // A different key rebuilds into the same memoizer.
    auto [k0b, k1b] = dpf::make_dpf(std::uint8_t{44}, std::uint64_t{9});
    std::tie(ibuf0, i0) = dpf::eval_interval(k0b, from, to, memo0);
    std::tie(ibuf1, i1) = dpf::eval_interval(k1b, from, to, memo1);
    it0 = std::begin(i0);
    it1 = std::begin(i1);
    for (std::uint8_t x = from; x <= to; ++x, ++it0, ++it1)
    {
        std::uint64_t got = dpf::reconstruct(*it0, *it1);
        std::uint64_t expect = (x == 44) ? 9 : 0;
        if (got != expect)
        {
            std::cerr << "interval memoizer reuse\n";
            return 1;
        }
    }

    //! [sequence-memoizer]
    std::array<std::uint8_t, 5> points{1, 7, 42, 100, 200};
    // The memoizer stores a reference to this recipe and checks it by address.
    auto recipe = dpf::make_sequence_recipe<key_t>(points.begin(), points.end());
    auto seq0 = dpf::make_inplace_reversing_sequence_memoizer<key_t>(recipe);
    auto seq1 = dpf::make_inplace_reversing_sequence_memoizer<key_t>(recipe);
    auto [sbuf0, s0] = dpf::eval_sequence(k0, recipe, seq0, dpf::return_output_only_tag_{});
    auto [sbuf1, s1] = dpf::eval_sequence(k1, recipe, seq1, dpf::return_output_only_tag_{});
    //! [sequence-memoizer]

    it0 = std::begin(s0);
    it1 = std::begin(s1);
    for (std::uint8_t x : points)
    {
        std::uint64_t got = dpf::reconstruct(*it0, *it1);
        std::uint64_t expect = (x == alpha) ? beta : 0;
        if (got != expect)
        {
            std::cerr << "sequence memoizer\n";
            return 1;
        }
        ++it0;
        ++it1;
    }

    std::tie(sbuf0, s0) = dpf::eval_sequence(k0b, recipe, seq0, dpf::return_output_only_tag_{});
    std::tie(sbuf1, s1) = dpf::eval_sequence(k1b, recipe, seq1, dpf::return_output_only_tag_{});
    it0 = std::begin(s0);
    it1 = std::begin(s1);
    for (std::uint8_t x : points)
    {
        std::uint64_t got = dpf::reconstruct(*it0, *it1);
        std::uint64_t expect = (x == 44) ? 9 : 0;
        if (got != expect)
        {
            std::cerr << "sequence memoizer reuse\n";
            return 1;
        }
        ++it0;
        ++it1;
    }

    std::cout << beta << "\n";
    return 0;
}