#include <gtest/gtest.h>

#include "dpf.hpp"

#include <cstdint>
#include <vector>

namespace
{

using in_type = std::uint8_t;
using out_type = std::uint64_t;

template <typename T>
out_type raw_of(const T & v)
{
    if constexpr (dpf::is_secret_share_v<std::decay_t<T>>)
        return static_cast<out_type>(v.raw());
    else
        return static_cast<out_type>(v);
}

out_type open(out_type a, out_type b)
{
    return static_cast<out_type>(a - b);
}

template <typename Key>
std::vector<out_type> interval_alone(const Key & key, in_type from, in_type to)
{
    auto buf = dpf::make_output_buffer_for_interval(key, from, to);
    dpf::eval_interval(key, from, to, buf);
    std::vector<out_type> v;
    v.reserve(buf.size());
    for (std::size_t i = 0; i < buf.size(); ++i)
        v.push_back(raw_of(buf[i]));
    return v;
}

}  // namespace

TEST(Cohort, SamePointGenAndPointEval)
{
    const in_type alpha = 7;
    const std::vector<out_type> beta{1, 2, 3, 4, 5};
    auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
    ASSERT_EQ(c0.size(), beta.size());

    std::vector<out_type> y0, y1;
    c0.eval_point(alpha, y0);
    c1.eval_point(alpha, y1);
    ASSERT_EQ(y0.size(), beta.size());
    for (std::size_t k = 0; k < beta.size(); ++k)
    {
        EXPECT_EQ(open(y0[k], y1[k]), beta[k]);
        EXPECT_EQ(y0[k], raw_of(*dpf::eval_point(c0.keys()[k], alpha)));
        EXPECT_EQ(y1[k], raw_of(*dpf::eval_point(c1.keys()[k], alpha)));
    }
    for (std::size_t k = 0; k < beta.size(); ++k)
    {
        const in_type other = static_cast<in_type>(alpha + 1 + k);
        std::vector<out_type> z0, z1;
        c0.eval_point(other, z0);
        c1.eval_point(other, z1);
        EXPECT_EQ(open(z0[k], z1[k]), 0u);
        EXPECT_EQ(z0[k], raw_of(*dpf::eval_point(c0.keys()[k], other)));
    }
}

TEST(Cohort, IntervalInterleavesLeaves)
{
    const in_type alpha = 9;
    const std::vector<out_type> beta{4, 8, 15, 16};
    auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
    const in_type from = 0, to = 15;
    std::vector<out_type> y0, y1;
    c0.eval_interval(from, to, y0);
    c1.eval_interval(from, to, y1);

    using key0 = std::decay_t<decltype(c0.keys()[0])>;
    constexpr std::size_t opl = key0::outputs_per_leaf;
    const std::size_t n = beta.size();
    const auto alone0 = interval_alone(c0.keys()[0], from, to);
    ASSERT_EQ(y0.size(), alone0.size() * n);
    ASSERT_EQ(alone0.size() % opl, 0u);

    for (std::size_t k = 0; k < n; ++k)
    {
        const auto a0 = interval_alone(c0.keys()[k], from, to);
        const auto a1 = interval_alone(c1.keys()[k], from, to);
        ASSERT_EQ(a0.size(), alone0.size());
        for (std::size_t i = 0; i < a0.size(); ++i)
        {
            const std::size_t node = i / opl;
            const std::size_t lane = i % opl;
            const std::size_t slot = (node * n + k) * opl + lane;
            EXPECT_EQ(y0[slot], a0[i]);
            EXPECT_EQ(y1[slot], a1[i]);
            const in_type x = static_cast<in_type>(from + i);
            const out_type opened = open(a0[i], a1[i]);
            if (x == alpha)
                EXPECT_EQ(opened, beta[k]);
            else
                EXPECT_EQ(opened, 0u);
        }
    }
}

TEST(Cohort, SequenceRecipeAndInnerProduct)
{
    const in_type alpha = 5;
    const std::vector<out_type> beta{3, 9, 1};
    auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
    const std::vector<in_type> points{1, 5, 5, 20, 40};
    const std::size_t n = beta.size();

    std::vector<out_type> y0, y1;
    c0.eval_sequence(points.begin(), points.end(), y0);
    c1.eval_sequence(points.begin(), points.end(), y1);
    ASSERT_EQ(y0.size(), points.size() * n);

    using key0 = std::decay_t<decltype(c0.keys()[0])>;
    auto recipe = dpf::make_sequence_recipe<key0>(
        points.begin(), points.end());
    std::vector<out_type> r0, r1;
    c0.eval_sequence(recipe, r0);
    c1.eval_sequence(recipe, r1);

    for (std::size_t q = 0; q < points.size(); ++q)
    {
        for (std::size_t k = 0; k < n; ++k)
        {
            const std::size_t slot = dpf::cohort_index(q, k, n);
            const auto e0 = raw_of(*dpf::eval_point(c0.keys()[k], points[q]));
            const auto e1 = raw_of(*dpf::eval_point(c1.keys()[k], points[q]));
            EXPECT_EQ(y0[slot], e0);
            EXPECT_EQ(y1[slot], e1);
            EXPECT_EQ(r0[slot], e0);
            EXPECT_EQ(r1[slot], e1);
            if (points[q] == alpha)
                EXPECT_EQ(open(e0, e1), beta[k]);
            else
                EXPECT_EQ(open(e0, e1), 0u);
        }
    }

    std::vector<out_type> w(points.size());
    for (std::size_t q = 0; q < w.size(); ++q)
        w[q] = q + 3;
    const auto s0 = c0.eval_sequence_inner_product(recipe, w);
    const auto s1 = c1.eval_sequence_inner_product(points.begin(), points.end(), w);
    ASSERT_EQ(s0.size(), n);
    for (std::size_t k = 0; k < n; ++k)
    {
        out_type acc0 = 0, acc1 = 0;
        for (std::size_t q = 0; q < points.size(); ++q)
        {
            acc0 = static_cast<out_type>(acc0
                + y0[dpf::cohort_index(q, k, n)] * w[q]);
            acc1 = static_cast<out_type>(acc1
                + y1[dpf::cohort_index(q, k, n)] * w[q]);
        }
        EXPECT_EQ(s0[k], acc0);
        EXPECT_EQ(s1[k], acc1);
        EXPECT_EQ(open(s0[k], s1[k]), beta[k] * (w[1] + w[2]));
    }
}

TEST(Cohort, IntervalInnerProduct)
{
    const in_type alpha = 4;
    const std::vector<out_type> beta{6, 7};
    auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
    const in_type from = 0, to = 15;
    const auto alone = interval_alone(c0.keys()[0], from, to);
    std::vector<out_type> w(alone.size());
    for (std::size_t i = 0; i < w.size(); ++i)
        w[i] = (i * 5u) + 1u;

    const auto a0 = c0.eval_interval_inner_product(from, to, w);
    const auto a1 = c1.eval_interval_inner_product(from, to, w);
    for (std::size_t k = 0; k < beta.size(); ++k)
    {
        const auto b0 = interval_alone(c0.keys()[k], from, to);
        const auto b1 = interval_alone(c1.keys()[k], from, to);
        out_type e0 = 0;
        out_type e1 = 0;
        for (std::size_t i = 0; i < w.size(); ++i)
        {
            e0 = static_cast<out_type>(e0 + b0[i] * w[i]);
            e1 = static_cast<out_type>(e1 + b1[i] * w[i]);
        }
        EXPECT_EQ(a0[k], e0);
        EXPECT_EQ(a1[k], e1);

        // Same weights against one-key batched leaf inner product.
        auto m0 = dpf::make_basic_interval_memoizer(c0.keys()[k], from, to);
        auto m1 = dpf::make_basic_interval_memoizer(c1.keys()[k], from, to);
        const auto ip0 = dpf::eval_inner_product(c0.keys()[k], from, to, w, m0);
        const auto ip1 = dpf::eval_inner_product(c1.keys()[k], from, to, w, m1);
        EXPECT_EQ(raw_of(ip0), a0[k]);
        EXPECT_EQ(raw_of(ip1), a1[k]);
        EXPECT_EQ(open(a0[k], a1[k]), beta[k] * w[alpha - from]);
    }
}

TEST(Cohort, InterleaveLeavesMatchesSequenceLayout)
{
    // Sequence cohort layout is `out[q * n + k]`, the same order
    // `interleave_leaves` writes for whole leaf values.
    const in_type alpha = 6;
    const std::vector<out_type> beta{2, 3, 5};
    auto [c0, c1] = dpf::make_dpf_cohort(alpha, beta.begin(), beta.end());
    (void)c1;
    const std::vector<in_type> pts{1, 6, 10, 20};
    const std::size_t n = beta.size();
    std::vector<out_type> cohort;
    c0.eval_sequence(pts.begin(), pts.end(), cohort);

    std::vector<std::vector<out_type>> per_key(n);
    std::vector<const out_type *> ptrs(n);
    for (std::size_t k = 0; k < n; ++k)
    {
        per_key[k].resize(pts.size());
        for (std::size_t q = 0; q < pts.size(); ++q)
            per_key[k][q] = raw_of(*dpf::eval_point(c0.keys()[k], pts[q]));
        ptrs[k] = per_key[k].data();
    }
    std::vector<out_type> interleaved(pts.size() * n);
    dpf::interleave_leaves<out_type>(
        interleaved.data(), ptrs.data(), n, pts.size());
    ASSERT_EQ(interleaved.size(), cohort.size());
    for (std::size_t i = 0; i < interleaved.size(); ++i)
        EXPECT_EQ(interleaved[i], cohort[i]) << "slot " << i;
}

TEST(Cohort, HalfTreeSamePoint)
{
    using ht = dpf::prg::aes128_ccr;
    const in_type alpha = 3;
    const std::vector<out_type> beta{11, 13, 17, 19};
    auto [c0, c1] = dpf::make_dpf_cohort<ht, dpf::prg::aes128>(
        alpha, beta.begin(), beta.end());
    std::vector<out_type> y0, y1;
    c0.eval_point(alpha, y0);
    c1.eval_point(alpha, y1);
    for (std::size_t k = 0; k < beta.size(); ++k)
    {
        EXPECT_EQ(open(y0[k], y1[k]), beta[k]);
        EXPECT_EQ(y0[k], raw_of(*dpf::eval_point(c0.keys()[k], alpha)));
    }
}