#include "aner/nn.hpp"
#include "aner/gpu.hpp"
#include <algorithm>
#include <chrono>
#include <cmath>
#include <cstdint>
#include <iomanip>
#include <iostream>
#include <numeric>
#include <limits>
#include <stdexcept>
#include <string>
#include <vector>

namespace nn = aner::nn;

// Isolate selection from forward work. Fresh resident logits precede every
// sample, so the old path cannot reuse a host mirror from a previous sample.
int main(int argc, char **argv) {
    try {
        std::string device = "cpu", dtype = "Float64";
        std::int64_t rows = 8, vocabulary = 4096;
        int repetitions = 12;
        for (int i = 1; i < argc; i += 2) {
            if (i + 1 >= argc) throw std::runtime_error("option needs value");
            const std::string flag = argv[i], value = argv[i + 1];
            if (flag == "--device") device = value;
            else if (flag == "--dtype") dtype = value;
            else if (flag == "--rows") rows = std::stoll(value);
            else if (flag == "--vocabulary") vocabulary = std::stoll(value);
            else if (flag == "--repetitions") repetitions = std::stoi(value);
            else throw std::runtime_error("unknown option");
        }
        if (rows < 1 || vocabulary < 4 || rows > 1000000 / vocabulary || repetitions < 2 || repetitions > 1000)
            throw std::runtime_error("invalid benchmark dimensions/repetitions");
        if (!nn::gpu::supports_dtype(device, dtype) && device != "cpu")
            throw std::runtime_error("requested device/precision unavailable; no fallback");
        const auto bos = vocabulary - 3, eos = vocabulary - 2, pad = vocabulary - 1;
        std::vector<double> values(static_cast<std::size_t>(rows * vocabulary));
        std::vector<std::int64_t> expected;
        for (std::int64_t row = 0; row < rows; ++row) {
            const auto winner = (row * 997 + 31) % bos;
            expected.push_back(winner);
            for (std::int64_t col = 0; col < vocabulary; ++col)
                values[static_cast<std::size_t>(row * vocabulary + col)] =
                    static_cast<double>((col * 179 + row * 173) % 8192) / 1024 - 4;
            values[static_cast<std::size_t>(row * vocabulary + winner)] = 20;
            values[static_cast<std::size_t>(row * vocabulary + bos)] = 100;
            values[static_cast<std::size_t>(row * vocabulary + pad)] = 101;
        }
        const auto input = nn::to(nn::astype(nn::tensor(rows, vocabulary, values), dtype), device);
        using Clock = std::chrono::steady_clock;
        for (int iteration = -2; iteration < repetitions; ++iteration) {
            for (int slot = 0; slot < 3; ++slot) {
                const auto mode = (iteration + 2 + slot) % 3;
                const bool resident = mode == 0;
                const bool linear = mode == 2;
                const auto scores = nn::reshape(nn::scale(input, 1), {rows, 1, vocabulary});
                nn::synchronize(device);
                const auto before = nn::memory_stats();
                nn::reset_memory_peaks();
                const auto begin = Clock::now();
                std::vector<std::int64_t> ids;
                if (resident) ids = nn::detail::generation_greedy(scores, bos, pad, eos);
                else {
                    const auto &host = scores.values();
                    for (std::int64_t row = 0; row < rows; ++row) {
                        if (linear) {
                            std::int64_t chosen = -1;
                            double maximum = -std::numeric_limits<double>::infinity();
                            for (std::int64_t id = 0; id < vocabulary; ++id) {
                                const auto score = host[static_cast<std::size_t>(row * vocabulary + id)];
                                if (!std::isfinite(score)) throw std::runtime_error("non-finite selection logits");
                                if (id == bos || id == pad) continue;
                                if (chosen < 0 || score > maximum) { chosen = id; maximum = score; }
                            }
                            ids.push_back(chosen);
                            continue;
                        }
                        std::vector<std::int64_t> candidates;
                        for (std::int64_t id = 0; id < vocabulary; ++id)
                            if (id != bos && id != pad) candidates.push_back(id);
                        std::sort(candidates.begin(), candidates.end(), [&](auto a, auto b) {
                            const auto x = host[static_cast<std::size_t>(row * vocabulary + a)];
                            const auto y = host[static_cast<std::size_t>(row * vocabulary + b)];
                            return x == y ? a < b : x > y;
                        });
                        ids.push_back(candidates.front());
                    }
                }
                nn::synchronize(device);
                const auto seconds = std::chrono::duration<double>(Clock::now() - begin).count();
                const auto after = nn::memory_stats();
                if (ids != expected) throw std::runtime_error("selection IDs disagree with independent fixture");
                if (iteration >= 0) {
                    std::cout << std::setprecision(17) << "{\"schema\":\"aner.gen-greedy-selection/1\",\"device\":\"" << device
                              << "\",\"dtype\":\"" << dtype << "\",\"rows\":" << rows << ",\"vocabulary\":" << vocabulary
                              << ",\"iteration\":" << iteration << ",\"order\":" << slot << ",\"mode\":\""
                              << (resident ? "native_greedy" : linear ? "host_linear" : "host_sort") << "\",\"seconds\":" << seconds
                              << ",\"selection_backend\":\"" << (resident ? (device == "cpu" ? "cpu_linear" :
                                  after.host_readback_cache.current_bytes > before.host_readback_cache.current_bytes ? "host_linear" : "gpu_argmax") :
                                  linear ? "host_linear" : "host_sort") << "\""
                              << ",\"logit_host_mirror_bytes\":" << after.host_readback_cache.current_bytes - before.host_readback_cache.current_bytes
                              << ",\"ids\":[";
                    for (std::size_t i = 0; i < ids.size(); ++i) std::cout << (i ? "," : "") << ids[i];
                    std::cout << "]}\n";
                }
            }
        }
        return 0;
    } catch (const std::exception &error) {
        std::cerr << error.what() << '\n';
        return 1;
    }
}