#include "aner/ml.hpp"
#include "aner/parallel_jobs.hpp"

#include <algorithm>
#include <bit>
#include <chrono>
#include <ctime>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <numeric>
#include <stdexcept>
#include <string>
#include <vector>

namespace ml = aner::ml;
namespace nn = aner::nn;
namespace par = aner::parallel;
namespace {
using Clock = std::chrono::steady_clock;
struct Arguments {
    std::size_t n = 2048, q = 256, d = 16, k = 5, workers = 1, repeats = 5, warmups = 1;
    std::uint64_t seed = 20261030;
    std::string mode = "serial", weights = "uniform";
};
Arguments arguments(int argc, char** argv) {
    Arguments a;
    for (int i = 1; i < argc; i += 2) {
        if (i + 1 == argc) throw std::runtime_error("every benchmark flag requires a value");
        const std::string flag = argv[i], value = argv[i + 1];
        if (flag == "--mode") a.mode = value;
        else if (flag == "--weights") a.weights = value;
        else {
            std::size_t consumed = 0;
            const auto number = std::stoull(value, &consumed);
            if (consumed != value.size()) throw std::runtime_error("invalid integer argument");
            if (flag == "--n") a.n = number;
            else if (flag == "--q") a.q = number;
            else if (flag == "--d") a.d = number;
            else if (flag == "--k") a.k = number;
            else if (flag == "--workers") a.workers = number;
            else if (flag == "--repeats") a.repeats = number;
            else if (flag == "--warmups") a.warmups = number;
            else if (flag == "--seed") a.seed = number;
            else throw std::runtime_error("unknown benchmark flag: " + flag);
        }
    }
    if (!a.n || !a.q || !a.d || !a.k || a.k > a.n || !a.repeats || a.repeats > 100 || a.warmups > 100 || !a.workers)
        throw std::runtime_error("invalid benchmark dimensions or repetition count");
    if (a.n > 1'000'000 / a.d || a.q > 1'000'000 / a.d || a.q > 50'000'000 / a.n / a.d)
        throw std::runtime_error("benchmark exceeds Aner input or whole-call work admission");
    if (a.mode != "serial" && a.mode != "context" && a.mode != "cold") throw std::runtime_error("mode must be serial, context or cold");
    if (a.weights != "uniform" && a.weights != "distance") throw std::runtime_error("weights must be uniform or distance");
#ifdef ANER_KNN_BASELINE
    if (a.mode != "serial" || a.weights != "uniform") throw std::runtime_error("original baseline supports serial uniform KNN only");
#endif
    if (a.workers > par::discover_cpu_budget().hard_limit) throw std::runtime_error("workers exceed discovered CPU allocation");
    return a;
}
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);
}
struct Neighbor { double distance; std::size_t index; };
bool less(const Neighbor& a, const Neighbor& b) {
    return a.distance < b.distance || (a.distance == b.distance && a.index < b.index);
}
struct Expected { std::vector<double> prediction, probability, indices, distances; };
Expected reference(const Arguments& a, const std::vector<double>& train, const std::vector<double>& labels,
                   const std::vector<double>& query) {
    Expected result;
    result.prediction.resize(a.q); result.probability.resize(a.q * 5);
    result.indices.resize(a.q * a.k); result.distances.resize(a.q * a.k);
    for (std::size_t row = 0; row < a.q; ++row) {
        std::vector<Neighbor> neighbors(a.n);
        for (std::size_t training = 0; training < a.n; ++training) {
            double norm = 0;
            for (std::size_t col = 0; col < a.d; ++col)
                norm = std::hypot(norm, query[row * a.d + col] - train[training * a.d + col]);
            neighbors[training] = {norm, training};
        }
        // Correctness oracle uses a complete sort and original per-coordinate
        // hypot; the measured original baseline calls the original ml.cpp API.
        std::sort(neighbors.begin(), neighbors.end(), less);
        std::vector<long double> votes(5, 0);
        const bool zero_only = a.weights == "distance" && neighbors.front().distance == 0;
        const long double scale = zero_only ? 1 : neighbors.front().distance;
        for (std::size_t i = 0; i < a.k; ++i) {
            const auto& neighbor = neighbors[i];
            result.indices[row * a.k + i] = static_cast<double>(neighbor.index);
            result.distances[row * a.k + i] = neighbor.distance;
            const auto c = static_cast<std::size_t>(labels[neighbor.index]);
            if (a.weights == "uniform") votes[c] += 1;
            else if (zero_only) votes[c] += neighbor.distance == 0 ? 1 : 0;
            else votes[c] += scale / static_cast<long double>(neighbor.distance);
        }
        result.prediction[row] = static_cast<double>(std::max_element(votes.begin(), votes.end()) - votes.begin());
        const auto denominator = std::accumulate(votes.begin(), votes.end(), static_cast<long double>(0));
        for (std::size_t c = 0; c < 5; ++c) result.probability[row * 5 + c] = static_cast<double>(votes[c] / denominator);
    }
    return result;
}
void exact(const std::vector<double>& actual, const std::vector<double>& expected, const char* name) {
    if (actual.size() != expected.size()) throw std::runtime_error(std::string(name) + " 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(std::string(name) + " exact check failed at cell " + std::to_string(i));
}
double median(std::vector<double> values) {
    std::sort(values.begin(), values.end());
    const auto middle = values.size() / 2;
    return values.size() % 2 ? values[middle] : (values[middle - 1] + values[middle]) / 2;
}
std::string quoted(const std::string& input) {
    std::string output = "\"";
    for (const auto c : input) { if (c == '\\' || c == '"') output += '\\'; output += c; }
    return output + '"';
}
}
int main(int argc, char** argv) {
    try {
        const auto a = arguments(argc, argv);
        if (a.n < 5 || a.q > 1'000'000 / a.k || a.q > 1'000'000 / 5)
            throw std::runtime_error("benchmark requires five fitted classes and admissible verification outputs");
        std::vector<double> train(a.n * a.d), labels(a.n), query(a.q * a.d);
        auto seed = a.seed;
        for (auto& value : train) value = static_cast<double>(next(seed) % 4097) / 256.0 - 8;
        for (auto& value : query) value = static_cast<double>(next(seed) % 4097) / 256.0 - 8;
        for (std::size_t row = 0; row < a.n; ++row) labels[row] = static_cast<double>(row % 5);
        const auto x = nn::tensor(a.n, a.d, train), y = nn::tensor(a.n, 1, labels), input = nn::tensor(a.q, a.d, query);
#ifdef ANER_KNN_BASELINE
        const auto model = ml::fit_knn(x, y, a.k);
        const std::string flavor = "original";
#else
        const auto model = ml::fit_knn(x, y, a.k, a.weights);
        const std::string flavor = "optimized";
        par::Context context;
        if (a.mode == "context") context = par::local_context(a.workers);
        ml::KNNQueryOptions options;
        if (a.mode == "context") options.execution = &context;
#endif
        // Independent reference, all readouts and model fitting stay outside
        // timed regions. Every measured mode must pass this verification.
        const auto expected = reference(a, train, labels, query);
        exact(ml::predict(model, input).values(), expected.prediction, "full-sort labels");
        exact(ml::neighbor_indices(model, input).values(), expected.indices, "full-sort indices");
        exact(ml::neighbor_distances(model, input).values(), expected.distances, "original hypot distance bits");
        const auto serial_probability = ml::predict_proba(model, input).values();
        if (a.weights == "uniform") exact(serial_probability, expected.probability, "uniform probabilities");
        else for (std::size_t i = 0; i < serial_probability.size(); ++i)
            if (std::abs(serial_probability[i] - expected.probability[i]) > 2e-14)
                throw std::runtime_error("weighted oracle probability check failed");
#ifndef ANER_KNN_BASELINE
        if (a.mode == "context") {
            exact(ml::predict(model, input, options).values(), expected.prediction, "context labels");
            exact(ml::predict_proba(model, input, options).values(), serial_probability, "context probability bits");
            exact(ml::neighbor_indices(model, input, options).values(), expected.indices, "context indices");
            exact(ml::neighbor_distances(model, input, options).values(), expected.distances, "context distance bits");
        }
#endif
        const auto invoke = [&]() -> nn::Tensor {
#ifdef ANER_KNN_BASELINE
            return ml::predict(model, input);
#else
            if (a.mode == "cold") {
                auto cold = par::local_context(a.workers);
                ml::KNNQueryOptions cold_options; cold_options.execution = &cold;
                return ml::predict(model, input, cold_options);
            }
            return ml::predict(model, input, options);
#endif
        };
        for (std::size_t i = 0; i < a.warmups; ++i) exact(invoke().values(), expected.prediction, "warmup output");
        std::vector<double> samples, cpu_samples, cpu_wall_ratios;
        for (std::size_t i = 0; i < a.repeats; ++i) {
            const auto cpu_begin = std::clock();
            const auto begin = Clock::now();
            const auto result = invoke();
            samples.push_back(std::chrono::duration<double>(Clock::now() - begin).count());
            cpu_samples.push_back(static_cast<double>(std::clock() - cpu_begin) / CLOCKS_PER_SEC);
            cpu_wall_ratios.push_back(cpu_samples.back() / samples.back());
            exact(result.values(), expected.prediction, "timed output");
        }
        const auto center = median(samples);
        std::vector<double> deviations;
        for (const auto value : samples) deviations.push_back(std::abs(value - center));
        const auto budget = par::discover_cpu_budget();
        std::cout << std::setprecision(17) << "{\"flavor\":" << quoted(flavor)
                  << ",\"mode\":" << quoted(a.mode) << ",\"weights\":" << quoted(a.weights)
                  << ",\"n\":" << a.n << ",\"q\":" << a.q << ",\"d\":" << a.d << ",\"k\":" << a.k
                  << ",\"workers_requested\":" << (a.mode == "serial" ? 0 : a.workers)
                  << ",\"cpu_hard_limit\":" << budget.hard_limit << ",\"cpu_auto_workers\":" << budget.automatic_workers
                  << ",\"seed\":" << a.seed << ",\"warmups\":" << a.warmups
                  << ",\"median_seconds\":" << center << ",\"mad_seconds\":" << median(deviations)
                  << ",\"median_process_cpu_seconds\":" << median(cpu_samples)
                  << ",\"median_process_cpu_wall_ratio\":" << median(cpu_wall_ratios)
                  << ",\"cpu_timing_scope\":\"process CPU seconds over predict; ratio estimates average CPU equivalents, not instantaneous utilization\""
                  << ",\"rows_per_second\":" << static_cast<double>(a.q) / center
                  << ",\"compiler\":" << quoted(__VERSION__)
                  << ",\"cplusplus\":" << __cplusplus
#ifdef NDEBUG
                  << ",\"ndebug\":true"
#else
                  << ",\"ndebug\":false"
#endif
#ifdef __FAST_MATH__
                  << ",\"fast_math\":true"
#else
                  << ",\"fast_math\":false"
#endif
                  << ",\"verification\":\"exact labels,indices,distances; exact uniform probabilities; weighted oracle tolerance2e-14 and exact context/serial probabilities\""
                  << ",\"timing_scope\":\"predict native call; excludes fit,input creation,verification,process startup; cold includes context creation/destruction\""
                  << ",\"raw_seconds\":[";
        for (std::size_t i = 0; i < samples.size(); ++i) { if (i) std::cout << ','; std::cout << samples[i]; }
        std::cout << "],\"raw_process_cpu_seconds\":[";
        for (std::size_t i = 0; i < cpu_samples.size(); ++i) { if (i) std::cout << ','; std::cout << cpu_samples[i]; }
        std::cout << "]}\n";
        return 0;
    } catch (const std::exception& error) {
        std::cerr << "KNN benchmark failed: " << error.what() << '\n';
        return 1;
    }
}