// Standalone verification tool; not an Aner runtime change or an RSS profiler.
// Build against the matching Release Aner C++ libraries, C++20, strict FP.
// The charged-byte formula below targets the recorded macOS arm64 libc++ ABI.
// It deliberately fails if a changed ABI or vector capacity changes admission.
#include "aner/ml.hpp"
#include "aner/parallel_jobs.hpp"

#include <algorithm>
#include <bit>
#include <cstdint>
#include <iostream>
#include <optional>
#include <stdexcept>
#include <string>
#include <vector>

namespace ml = aner::ml;
namespace nn = aner::nn;
namespace par = aner::parallel;

static void exact(const std::vector<double>& actual, const std::vector<double>& expected) {
    if (actual.size() != expected.size()) throw std::runtime_error("output shape mismatch");
    for (std::size_t i = 0; i < actual.size(); ++i)
        if (std::bit_cast<std::uint64_t>(actual[i]) != std::bit_cast<std::uint64_t>(expected[i]))
            throw std::runtime_error("serial/context output bits differ");
}
static std::uint64_t next(std::uint64_t& seed) {
    seed += UINT64_C(0x9e3779b97f4a7c15);
    auto value = seed;
    value = (value ^ (value >> 30)) * UINT64_C(0xbf58476d1ce4e5b9);
    value = (value ^ (value >> 27)) * UINT64_C(0x94d049bb133111eb);
    return value ^ (value >> 31);
}

int main() {
    try {
        if (sizeof(double) != 8 || sizeof(std::size_t) != 8)
            throw std::runtime_error("probe requires the recorded 64-bit ABI");
        if (par::discover_cpu_budget().hard_limit < 4)
            throw std::runtime_error("probe requires admission of four CPU workers");
        struct Workload { const char* name; std::size_t n, q, d; };
        bool first = true;
        std::cout << "[\n";
        for (const auto& workload : {Workload{"batch", 2048, 256, 16}, Workload{"wide", 4096, 64, 64}}) {
            const auto [name, n, q, d] = workload;
            constexpr std::size_t c = 5, k = 5;
            std::uint64_t seed = 20261030;
            std::vector<double> train(n * d), labels(n), queries(q * d);
            for (auto& value : train) value = static_cast<double>(next(seed) % 4097) / 256.0 - 8;
            for (auto& value : queries) value = static_cast<double>(next(seed) % 4097) / 256.0 - 8;
            for (std::size_t row = 0; row < n; ++row) labels[row] = static_cast<double>(row % c);
            // Exercise selected zero-distance duplicates as well as ordinary queries.
            std::copy_n(train.begin(), d, train.begin() + static_cast<std::ptrdiff_t>(d));
            std::copy_n(train.begin(), d, queries.begin());
            const auto x = nn::tensor(n, d, train), y = nn::tensor(n, 1, labels), input = nn::tensor(q, d, queries);
            for (const std::string weights : {"uniform", "distance"}) {
                const auto model = ml::fit_knn(x, y, k, weights);
                const auto serial_labels = ml::predict(model, input).values();
                const auto serial_probabilities = ml::predict_proba(model, input).values();
                const auto serial_indices = ml::neighbor_indices(model, input).values();
                const auto serial_distances = ml::neighbor_distances(model, input).values();
                for (const std::size_t workers : {0, 1, 2, 4}) {
                    std::optional<par::Context> context;
                    if (workers) context.emplace(par::local_context(workers));
                    const auto lanes = workers ? std::min(workers, q) : 1;
                    const auto bytes = 344 + 8 * q + lanes * (136 + 16 * n + 8 * c + (weights == "distance" ? 24 * c : 0));
                    ml::KNNQueryOptions options;
                    if (context) options.execution = &*context;
                    options.max_workspace_bytes = bytes - 1;
                    bool rejected = false;
                    try { (void)ml::predict(model, input, options); }
                    catch (const ml::Error& error) {
                        rejected = std::string(error.what()).find("workspace") != std::string::npos;
                    }
                    if (!rejected) throw std::runtime_error("one-byte-below workspace was not rejected");
                    options.max_workspace_bytes = bytes;
                    exact(ml::predict(model, input, options).values(), serial_labels);
                    // Probability and neighbor results have different output sizes;
                    // their equality checks use the default budget, not predict's cap.
                    options.max_workspace_bytes = 64 * 1024 * 1024;
                    exact(ml::predict_proba(model, input, options).values(), serial_probabilities);
                    exact(ml::neighbor_indices(model, input, options).values(), serial_indices);
                    exact(ml::neighbor_distances(model, input, options).values(), serial_distances);
                    if (!first) std::cout << ",\n";
                    first = false;
                    std::cout << "{\"workload\":\"" << name << "\",\"n\":" << n << ",\"q\":" << q
                              << ",\"d\":" << d << ",\"k\":" << k << ",\"c\":" << c
                              << ",\"weights\":\"" << weights << "\",\"operation\":\"predict\",\"execution\":\""
                              << (workers ? "context" : "serial") << "\",\"workers_requested\":" << workers
                              << ",\"active_lanes\":" << lanes << ",\"charged_workspace_bytes\":" << bytes
                              << ",\"one_byte_below_rejected\":true,\"at_budget_accepted\":true,\"bitwise_serial_equal\":true"
                              << ",\"bitwise_readouts\":[\"labels\",\"probabilities\",\"neighbor_indices\",\"neighbor_distances\"]}";
                }
            }
        }
        std::cout << "\n]\n";
        return 0;
    } catch (const std::exception& error) {
        std::cerr << error.what() << '\n';
        return 1;
    }
}