/// @file test/profile/grotto_profile.cpp
/// @brief Workload for grotto evaluation: tables, offset walks, prefix parity.
///
/// Key generation is a separate case from the eval that consumes those keys.
/// Table cases probe the domain once, then time only inputs that evaluate.
/// `items` is the number of table evaluations inside one sample, or 1 for a
/// single offset / prefix walk (both parties).
///
///   profile_grotto --list
///   profile_grotto --case window_gelu_p16 --repeat 100
///
/// Profile-guided build, separate from COVERAGE:
///   cmake -S test -B build-pgo -DLIBDPF_PGO=generate -DCMAKE_BUILD_TYPE=Release
///   cmake --build build-pgo --target profile_grotto
///   build-pgo/bin/profile_grotto --repeat 40 --warmup 2
///   cmake -S test -B build-pgo -DLIBDPF_PGO=use -DCMAKE_BUILD_TYPE=Release
///   cmake --build build-pgo --target profile_grotto

#include "harness.hpp"

#include "dpf.hpp"
#include "grotto/closed_form.hpp"
#include "grotto/exact_steps.hpp"
#include "grotto/offset_horner.hpp"
#include "grotto/offset_jet.hpp"
#include "grotto/offset_poly.hpp"
#include "grotto/offset_repr.hpp"
#include "grotto/offset_twist.hpp"
#include "grotto/prefix_parity.hpp"
#include "grotto/range_lut.hpp"
#include "grotto/window_lut.hpp"

#include <array>
#include <cstdint>
#include <functional>
#include <initializer_list>
#include <memory>
#include <string>
#include <vector>

namespace
{

using profile::sample;
using profile::touch_arr;
using profile::touch_vec;
using profile::touch_word;
using profile::work;

const char kUsage[] =
    "profile_grotto [--list] [--family grotto] [--slice S] [--case NAME]\n"
    "               [--repeat N=30] [--warmup W=2]\n"
    "Slices: window, principal, closed, reduced, exact, horner, poly, jet,\n"
    "        twist, repr, prefix, keygen.\n"
    "Times grotto table evals and offset / prefix walks.\n";

inline work cell(const char * slice, std::string name, std::uint64_t items,
    std::function<sample()> fn)
{
    return profile::make_work("grotto", slice, std::move(name), items, std::move(fn));
}

template <typename Enum, typename Fn>
std::vector<std::int64_t> probe(Enum which, unsigned bits, Fn fn,
    const std::vector<std::int64_t> & cand)
{
    std::vector<std::int64_t> ok;
    ok.reserve(cand.size());
    for (const std::int64_t raw : cand)
    {
        try
        {
            (void)fn(which, bits, raw);
            ok.push_back(raw);
        }
        catch (const std::exception &)
        {
        }
    }
    return ok;
}

std::vector<std::int64_t> around_one(unsigned bits)
{
    const std::int64_t one = std::int64_t{1} << bits;
    return {
        -4 * one, -2 * one, -one, -one / 2, -1, 0, 1,
        one / 20, one / 5, one / 3, one / 2, one, (3 * one) / 2, 2 * one, 4 * one
    };
}

template <typename Fn>
void add_table(std::vector<work> & works, const char * slice, std::string name, Fn fn,
    std::vector<std::int64_t> samples)
{
    if (samples.empty())
    {
        std::cerr << name << " skipped: no in-domain sample\n";
        return;
    }
    auto held = std::make_shared<std::vector<std::int64_t>>(std::move(samples));
    const std::uint64_t items = static_cast<std::uint64_t>(held->size()) * 32u;
    std::function<sample()> body = [held, fn]() -> sample {
        std::uint64_t h = 0x84222325cbf29ce4ull;
        for (int rep = 0; rep < 32; ++rep)
        {
            for (const std::int64_t raw : *held)
            {
                // Multiply-add, not a plain xor: an even number of identical
                // xors cancels, and the compiler will delete the evals.
                h = (h * 0x100000001b3ull) ^ static_cast<std::uint64_t>(fn(raw))
                    ^ static_cast<std::uint64_t>(rep);
            }
        }
        return touch_word(h, held->size() * sizeof(std::int64_t));
    };
    works.push_back(cell(slice, std::move(name), items, std::move(body)));
}

template <typename T>
std::vector<T> knots_of(std::size_t n, T gap)
{
    std::vector<T> knots(n);
    for (std::size_t i = 0; i < n; ++i)
        knots[i] = static_cast<T>(1 + gap * static_cast<T>(i));
    return knots;
}

template <std::size_t Degree>
std::vector<std::array<std::uint64_t, Degree + 1>> horner_rows(std::size_t n)
{
    std::vector<std::array<std::uint64_t, Degree + 1>> rows(n);
    for (std::size_t i = 0; i < n; ++i)
        for (std::size_t m = 0; m <= Degree; ++m)
            rows[i][m] = (i + 1) * (m + 3);
    return rows;
}

std::vector<std::vector<std::uint64_t>> poly_rows(std::size_t n, std::size_t degree)
{
    std::vector<std::vector<std::uint64_t>> rows(n, std::vector<std::uint64_t>(degree + 1));
    for (std::size_t i = 0; i < n; ++i)
        for (std::size_t m = 0; m <= degree; ++m)
            rows[i][m] = (i + 1) * (m + 3);
    return rows;
}

std::vector<std::uint64_t> coeff_col(std::size_t degree)
{
    std::vector<std::uint64_t> c(degree + 1);
    for (std::size_t m = 0; m <= degree; ++m)
        c[m] = m + 5;
    return c;
}

template <std::size_t N, typename T>
std::array<T, N> ends_of(T start, T step)
{
    std::array<T, N> ends{};
    for (std::size_t i = 0; i < N; ++i)
        ends[i] = static_cast<T>(start + step * static_cast<T>(i));
    return ends;
}

}  // namespace

template <std::size_t Degree>
void add_horner_degree(std::vector<work> & works)
{
    using in_t = std::uint32_t;
    const in_t center = 50;
    const in_t eta = 7;
    works.push_back(cell("keygen", "horner_keygen_d" + std::to_string(Degree), 1, [] {
        auto built = grotto::make_offset_horner_keys<in_t, Degree>(in_t{50});
        return touch_word(built.wrap_share[0][0] ^ built.wrap_share[Degree][1],
            sizeof(built.wrap_share));
    }));
    using mat_t = decltype(grotto::make_offset_horner_keys<in_t, Degree>(center));
    auto mat = std::make_shared<mat_t>(grotto::make_offset_horner_keys<in_t, Degree>(center));
    for (const std::size_t pieces : std::initializer_list<std::size_t>{1, 4, 16, 64})
    {
        auto knots = std::make_shared<std::vector<in_t>>(knots_of<in_t>(pieces, 1000));
        auto rows = std::make_shared<std::vector<std::array<std::uint64_t, Degree + 1>>>(
            horner_rows<Degree>(pieces));
        const auto name = "horner_eval_d" + std::to_string(Degree)
            + "_p" + std::to_string(pieces);
        works.push_back(cell("horner", name, 1, [mat, knots, rows, eta] {
            const auto a = grotto::offset_horner_eval<0, Degree>(*mat, *knots, *rows, eta);
            const auto b = grotto::offset_horner_eval<1, Degree>(*mat, *knots, *rows, eta);
            return touch_word(a ^ b, 16);
        }));
    }
}

template <typename Input, typename BitPtr, typename GtPtr, std::size_t N>
void add_prefix_n(std::vector<work> & works, const char * width, Input step,
    const BitPtr & bits, const GtPtr & gts)
{
    const auto ends = ends_of<N>(Input{1}, step);
    const auto ntag = std::to_string(N);
    works.push_back(cell("prefix", std::string("prefix_") + width + "_e" + ntag, N,
        [bits, ends] {
            const auto a = grotto::prefix_parities(std::get<0>(*bits), ends);
            const auto b = grotto::prefix_parities(std::get<1>(*bits), ends);
            return touch_arr(std::get<0>(a)) + touch_arr(std::get<0>(b));
        }));
    works.push_back(cell("prefix", std::string("signed_prefix_") + width + "_e" + ntag, N,
        [gts, ends] {
            const auto a = grotto::signed_prefix_parities(std::get<0>(*gts), ends);
            const auto b = grotto::signed_prefix_parities(std::get<1>(*gts), ends);
            return touch_arr(a) + touch_arr(b);
        }));
}

template <typename Input>
void add_prefix_width(std::vector<work> & works, const char * width, Input step)
{
    const auto bit_keys = dpf::make_dpf(Input{1000}, dpf::bit{1});
    const auto gt_keys = dpf::make_dpf(Input{1000}, dpf::gt(std::uint64_t{1}));
    using bit_ptr = std::shared_ptr<std::decay_t<decltype(bit_keys)>>;
    using gt_ptr = std::shared_ptr<std::decay_t<decltype(gt_keys)>>;
    auto bits = std::make_shared<std::decay_t<decltype(bit_keys)>>(bit_keys);
    auto gts = std::make_shared<std::decay_t<decltype(gt_keys)>>(gt_keys);
    add_prefix_n<Input, bit_ptr, gt_ptr, 1>(works, width, step, bits, gts);
    add_prefix_n<Input, bit_ptr, gt_ptr, 8>(works, width, step, bits, gts);
    add_prefix_n<Input, bit_ptr, gt_ptr, 32>(works, width, step, bits, gts);
    add_prefix_n<Input, bit_ptr, gt_ptr, 128>(works, width, step, bits, gts);
}

int main(int argc, char ** argv)
{
    const auto opt = profile::parse_args(argc, argv, 30, 2, kUsage);
    std::vector<work> works;

    const unsigned precisions[] = {8u, 16u, 24u, 32u};
    const auto add_window = [&](const char * stem, grotto::window which, unsigned bits) {
        auto samples = probe(which, bits,
            [](grotto::window w, unsigned b, std::int64_t raw) {
                return grotto::eval_window(w, b, raw);
            }, around_one(bits));
        add_table(works, "window",
            std::string("window_") + stem + "_p" + std::to_string(bits),
            [which, bits](std::int64_t raw) {
                return grotto::eval_window(which, bits, raw);
            }, std::move(samples));
    };
    const auto add_principal = [&](const char * stem, grotto::principal which, unsigned bits) {
        auto samples = probe(which, bits,
            [](grotto::principal w, unsigned b, std::int64_t raw) {
                return grotto::eval_principal(w, b, raw);
            }, around_one(bits));
        add_table(works, "principal",
            std::string("principal_") + stem + "_p" + std::to_string(bits),
            [which, bits](std::int64_t raw) {
                return grotto::eval_principal(which, bits, raw);
            }, std::move(samples));
    };
    const auto add_closed = [&](const char * stem, grotto::closed which, unsigned bits) {
        auto samples = probe(which, bits,
            [](grotto::closed w, unsigned b, std::int64_t raw) {
                return grotto::eval_closed(w, b, raw);
            }, around_one(bits));
        add_table(works, "closed",
            std::string("closed_") + stem + "_p" + std::to_string(bits),
            [which, bits](std::int64_t raw) {
                return grotto::eval_closed(which, bits, raw);
            }, std::move(samples));
    };
    const auto add_reduced = [&](const char * stem, grotto::reduced which, unsigned bits) {
        auto samples = probe(which, bits,
            [](grotto::reduced w, unsigned b, std::int64_t raw) {
                return grotto::eval_reduced(w, b, raw);
            }, around_one(bits));
        add_table(works, "reduced",
            std::string("reduced_") + stem + "_p" + std::to_string(bits),
            [which, bits](std::int64_t raw) {
                return grotto::eval_reduced(which, bits, raw);
            }, std::move(samples));
    };

    const struct { const char * stem; grotto::window which; } windows[] = {
        {"gelu", grotto::window::gelu},
        {"sigmoid", grotto::window::sigmoid},
        {"tanh", grotto::window::tanh},
        {"erf", grotto::window::erf},
        {"erfc", grotto::window::erfc},
        {"silu", grotto::window::silu},
        {"mish", grotto::window::mish},
        {"softplus", grotto::window::softplus},
        {"probit", grotto::window::probit},
        {"smoothstep", grotto::window::smoothstep},
        {"asin", grotto::window::asin},
        {"acos", grotto::window::acos},
        {"hardelish", grotto::window::hardelish},
        {"lecun_tanh", grotto::window::lecun_tanh},
    };
    for (const auto & spec : windows)
        for (const unsigned bits : precisions)
            add_window(spec.stem, spec.which, bits);

    const struct { const char * stem; grotto::principal which; } principals[] = {
        {"ln", grotto::principal::ln},
        {"exp", grotto::principal::exp},
        {"sin", grotto::principal::sin},
        {"sqrt", grotto::principal::sqrt},
        {"sinh", grotto::principal::sinh},
        {"inv", grotto::principal::inv},
    };
    for (const auto & spec : principals)
        for (const unsigned bits : precisions)
            add_principal(spec.stem, spec.which, bits);

    const struct { const char * stem; grotto::closed which; } closeds[] = {
        {"atan", grotto::closed::atan},
        {"cbrt", grotto::closed::cbrt},
        {"sinc", grotto::closed::sinc},
        {"softsign", grotto::closed::softsign},
        {"logistic", grotto::closed::logistic},
    };
    for (const auto & spec : closeds)
        for (const unsigned bits : precisions)
            add_closed(spec.stem, spec.which, bits);

    const struct { const char * stem; grotto::reduced which; } reduceds[] = {
        {"ln", grotto::reduced::ln},
        {"exp", grotto::reduced::exp},
        {"sin", grotto::reduced::sin},
        {"expm1", grotto::reduced::expm1},
        {"log1p", grotto::reduced::log1p},
        {"sqrt", grotto::reduced::sqrt},
    };
    for (const auto & spec : reduceds)
        for (const unsigned bits : precisions)
            add_reduced(spec.stem, spec.which, bits);

    const auto add_exact = [&](const std::string & name, unsigned bits, auto fn) {
        std::vector<std::int64_t> samples;
        for (const std::int64_t raw : around_one(bits))
        {
            try
            {
                (void)fn(raw);
                samples.push_back(raw);
            }
            catch (const std::exception &)
            {
            }
        }
        add_table(works, "exact", name, fn, std::move(samples));
    };
    for (const unsigned bits : {8u, 16u, 32u})
    {
        const auto tag = "_p" + std::to_string(bits);
        add_exact("exact_dec_floor" + tag, bits, [bits](std::int64_t raw) {
            return grotto::eval_dec_floor(raw, bits);
        });
        add_exact("exact_bit_width" + tag, bits, [bits](std::int64_t raw) {
            return grotto::eval_bit_width<std::int64_t>(raw, bits);
        });
        add_exact("exact_dec_width" + tag, bits, [bits](std::int64_t raw) {
            return grotto::eval_dec_width(raw, bits);
        });
        add_exact("exact_ilog256" + tag, bits, [bits](std::int64_t raw) {
            return grotto::eval_ilog256<std::int64_t>(raw, bits);
        });
    }

    add_horner_degree<1>(works);
    add_horner_degree<3>(works);

    using in_t = std::uint32_t;
    const in_t center = 50;
    const in_t eta = 7;
    for (const std::size_t degree : std::initializer_list<std::size_t>{1, 4, 8, 16})
    {
        works.push_back(cell("keygen", "poly_keygen_d" + std::to_string(degree), 1,
            [degree] {
                auto built = grotto::make_offset_poly_keys(in_t{50}, degree);
                return touch_word(built.wrap_share[0][0],
                    built.wrap_share.size() * sizeof(built.wrap_share[0]));
            }));
        using mat_t = decltype(grotto::make_offset_poly_keys(center, degree));
        auto mat = std::make_shared<mat_t>(grotto::make_offset_poly_keys(center, degree));
        for (const std::size_t pieces : std::initializer_list<std::size_t>{4, 8, 16})
        {
            if (degree == 16 && pieces == 16)
                continue;
            auto knots = std::make_shared<std::vector<in_t>>(knots_of<in_t>(pieces, 800));
            auto rows = std::make_shared<std::vector<std::vector<std::uint64_t>>>(
                poly_rows(pieces, degree));
            const auto name = "poly_eval_d" + std::to_string(degree)
                + "_p" + std::to_string(pieces);
            works.push_back(cell("poly", name, 1, [mat, knots, rows] {
                const auto a = grotto::offset_poly_eval<0>(*mat, *knots, *rows, eta, nullptr);
                const auto b = grotto::offset_poly_eval<1>(*mat, *knots, *rows, eta, nullptr);
                return touch_word(a ^ b, 16);
            }));
        }
    }

    for (const std::size_t degree : std::initializer_list<std::size_t>{1, 4, 8, 16})
    {
        works.push_back(cell("keygen", "jet_keygen_d" + std::to_string(degree), 1,
            [degree] {
                auto built = grotto::make_offset_jet_keys(in_t{50}, degree);
                return touch_word(built.wrap_share[0][0],
                    built.wrap_share.size() * sizeof(built.wrap_share[0]));
            }));
        using mat_t = decltype(grotto::make_offset_jet_keys(center, degree));
        auto mat = std::make_shared<mat_t>(grotto::make_offset_jet_keys(center, degree));
        for (const std::size_t pieces : std::initializer_list<std::size_t>{4, 16})
        {
            auto knots = std::make_shared<std::vector<in_t>>(knots_of<in_t>(pieces, 800));
            auto col = std::make_shared<std::vector<std::uint64_t>>(coeff_col(degree));
            const auto name = "jet_eval_d" + std::to_string(degree)
                + "_p" + std::to_string(pieces);
            works.push_back(cell("jet", name, 1, [mat, knots, col] {
                const auto a = grotto::offset_jet_eval<0>(*mat, *knots, *col, eta);
                const auto b = grotto::offset_jet_eval<1>(*mat, *knots, *col, eta);
                return touch_word(a ^ b, 16);
            }));
        }
    }

    for (const std::size_t degree : std::initializer_list<std::size_t>{1, 4, 8})
    {
        auto knots = std::make_shared<std::vector<in_t>>(knots_of<in_t>(8, 900));
        auto col = std::make_shared<std::vector<std::uint64_t>>(coeff_col(degree));
        using odd_t = decltype(grotto::make_offset_twist_keys(center, degree, std::uint64_t{3}));
        using half_t = decltype(grotto::make_offset_twist_keys(center, degree, grotto::twist_half));
        auto odd = std::make_shared<odd_t>(
            grotto::make_offset_twist_keys(center, degree, std::uint64_t{3}));
        auto half = std::make_shared<half_t>(
            grotto::make_offset_twist_keys(center, degree, grotto::twist_half));
        const auto dtag = std::to_string(degree);
        works.push_back(cell("keygen", "twist_keygen_odd_d" + dtag, 1, [degree] {
            auto built = grotto::make_offset_twist_keys(in_t{50}, degree, std::uint64_t{3});
            return touch_word(built.wrap_share[0][0],
                built.wrap_share.size() * sizeof(built.wrap_share[0]));
        }));
        works.push_back(cell("twist", "twist_eval_odd_d" + dtag + "_p8", 1,
            [odd, knots, col] {
                const auto a = grotto::offset_twist_eval<0>(*odd, *knots, *col, eta);
                const auto b = grotto::offset_twist_eval<1>(*odd, *knots, *col, eta);
                return touch_word(a ^ b, 16);
            }));
        works.push_back(cell("twist", "twist_eval_half_d" + dtag + "_p8", 1,
            [half, knots, col] {
                const auto a = grotto::offset_twist_eval<0>(*half, *knots, *col, eta);
                const auto b = grotto::offset_twist_eval<1>(*half, *knots, *col, eta);
                return touch_word(a ^ b, 16);
            }));
    }

    {
        const std::vector<std::uint64_t> fib_state{1, 0};
        const std::vector<std::uint64_t> wide_state{1, 2, 3, 4};
        using fib_t = decltype(grotto::make_offset_repr_keys(center, fib_state));
        using wide_t = decltype(grotto::make_offset_repr_keys(center, wide_state));
        auto fib = std::make_shared<fib_t>(grotto::make_offset_repr_keys(center, fib_state));
        auto wide = std::make_shared<wide_t>(grotto::make_offset_repr_keys(center, wide_state));
        auto fib_m = std::make_shared<std::vector<std::vector<std::uint64_t>>>(
            grotto::offset_repr_fibonacci_matrix());
        auto tri_m = std::make_shared<std::vector<std::vector<std::uint64_t>>>(
            std::vector<std::vector<std::uint64_t>>{
                {1, 1, 0, 0},
                {0, 1, 1, 0},
                {0, 0, 1, 1},
                {0, 0, 0, 1},
            });
        works.push_back(cell("keygen", "repr_keygen_dim2", 1, [] {
            auto built = grotto::make_offset_repr_keys(in_t{50}, std::vector<std::uint64_t>{1, 0});
            return touch_word(built.wrap_share[0][0],
                built.wrap_share.size() * sizeof(built.wrap_share[0]));
        }));
        works.push_back(cell("keygen", "repr_keygen_dim4", 1, [] {
            auto built = grotto::make_offset_repr_keys(in_t{50},
                std::vector<std::uint64_t>{1, 2, 3, 4});
            return touch_word(built.wrap_share[0][0],
                built.wrap_share.size() * sizeof(built.wrap_share[0]));
        }));
        for (const std::size_t pieces : std::initializer_list<std::size_t>{4, 8, 16})
        {
            auto knots = std::make_shared<std::vector<in_t>>(knots_of<in_t>(pieces, 700));
            const auto tag = "_p" + std::to_string(pieces);
            works.push_back(cell("repr", "repr_eval_fib" + tag, 1, [fib, fib_m, knots] {
                const auto a = grotto::offset_repr_eval<0>(*fib, *fib_m, *knots, eta);
                const auto b = grotto::offset_repr_eval<1>(*fib, *fib_m, *knots, eta);
                return touch_vec(a) + touch_vec(b);
            }));
            works.push_back(cell("repr", "repr_eval_dim4" + tag, 1, [wide, tri_m, knots] {
                const auto a = grotto::offset_repr_eval<0>(*wide, *tri_m, *knots, eta);
                const auto b = grotto::offset_repr_eval<1>(*wide, *tri_m, *knots, eta);
                return touch_vec(a) + touch_vec(b);
            }));
        }
        works.push_back(cell("repr", "repr_crc32_jump", 64, [] {
            std::uint32_t state = 0x12345678u;
            for (unsigned i = 0; i < 64; ++i)
                state = grotto::offset_repr_crc32_jump(state, 1ull << (i % 17));
            return touch_word(state, 64u * 32u * sizeof(std::uint32_t));
        }));
    }

    add_prefix_width<std::uint16_t>(works, "u16", 400);
    add_prefix_width<std::uint32_t>(works, "u32", 100000);

    return profile::run_works(opt, works);
}