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

namespace gen = aner::gen;
namespace nn = aner::nn;
namespace token = aner::token;
using IDs = std::vector<std::vector<std::int64_t>>;
using Clock = std::chrono::steady_clock;

namespace {
void require(bool value, const std::string &message) {
    if (!value) throw std::runtime_error(message);
}
template<class T> void array(const std::vector<T> &values) {
    std::cout << '[';
    for (std::size_t i = 0; i < values.size(); ++i) std::cout << (i ? "," : "") << values[i];
    std::cout << ']';
}
void nested(const IDs &rows) {
    std::cout << '[';
    for (std::size_t i = 0; i < rows.size(); ++i) { if (i) std::cout << ','; array(rows[i]); }
    std::cout << ']';
}
struct Options {
    std::string device = "mps", dtype = "Float32", phase = "natural", name = "tiny";
    std::int64_t batch = 8, width = 32, layers = 2;
    int repetitions = 3, warmups = 2;
    bool native_plan = true;
};
struct Work {
    IDs prompts;
    std::vector<std::int64_t> lengths, budgets, capacities, seeds;
};
Work work(const Options &options) {
    Work result;
    if (options.name == "ragged") {
        require(options.batch == 8, "ragged workload requires eight logical rows");
        result.lengths = {1, 3, 7, 31, 2, 5, 13, 47};
        result.budgets = {32, 1, 16, 2, 24, 3, 8, 4};
    } else {
        require(options.name == "tiny" || options.name == "long", "unknown workload");
        result.lengths.assign(static_cast<std::size_t>(options.batch), 8);
        result.budgets.assign(static_cast<std::size_t>(options.batch), options.name == "tiny" ? 8 : 64);
    }
    for (std::size_t row = 0; row < result.lengths.size(); ++row) {
        std::vector<std::int64_t> prompt{token::bos_id};
        for (std::int64_t step = 1; step < result.lengths[row]; ++step)
            prompt.push_back(97 + (37 * step + 11 * static_cast<std::int64_t>(row) + 5) % 26);
        result.prompts.push_back(std::move(prompt));
        result.capacities.push_back(result.lengths[row] + result.budgets[row] - 1);
        result.seeds.push_back(7301 + 101 * static_cast<std::int64_t>(row));
    }
    return result;
}
IDs forced_inputs(const Work &workload, std::int64_t step) {
    IDs result(workload.prompts.size());
    for (std::size_t row = 0; row < result.size(); ++row)
        if (step < workload.budgets[row] - 1)
            result[row] = {97 + (7 * step + 13 * static_cast<std::int64_t>(row) + 3) % 26};
    return result;
}
template<class Cache = gen::detail::GenerationCache>
Cache single_cache(const gen::Model &model, std::int64_t capacity, bool native) {
    if constexpr (std::is_constructible_v<Cache, const gen::Model &, bool, std::int64_t, bool>)
        return Cache(model, true, capacity, native);
    else {
        require(native, "baseline does not expose the primitive-plan testing selector");
        return Cache(model, true, capacity);
    }
}
template<class Cache = gen::detail::GenerationBatchCache>
Cache batch_cache(const gen::Model &model, const Work &workload, bool native) {
    if constexpr (std::is_constructible_v<Cache, const gen::Model &, const std::vector<std::int64_t> &, bool, bool, bool>)
        return Cache(model, workload.capacities, true, false, native);
    else {
        require(native, "baseline does not expose the primitive-plan testing selector");
        return Cache(model, workload.capacities, true, false);
    }
}
struct Statistics {
    bool available = false;
    std::uint64_t forwards = 0, buffers = 0, reads = 0, gemms = 0, encoders = 0, packed_rows = 0;
};
template<class Cache> Statistics statistics(const Cache &cache) {
    if constexpr (requires { cache.statistics(); }) {
        const auto source = cache.statistics();
        return {source.available, source.forward_calls, source.command_buffers, source.status_reads,
                source.gemm_calls, source.compute_encoders, source.packed_output_rows};
    } else return {};
}
void statistics_json(const Statistics &stats) {
    std::cout << "{\"available\":" << (stats.available ? "true" : "false")
              << ",\"forward_calls\":" << stats.forwards << ",\"command_buffers\":" << stats.buffers
              << ",\"status_reads\":" << stats.reads << ",\"gemm_calls\":" << stats.gemms
              << ",\"compute_encoders\":" << stats.encoders << ",\"packed_output_rows\":" << stats.packed_rows << '}';
}
void prefix(const Options &options, const Work &workload) {
    std::cout << "{\"schema\":\"aner.metal-decoder-benchmark/1\",\"phase\":\"" << options.phase
              << "\",\"device\":\"" << options.device << "\",\"dtype\":\"" << options.dtype
              << "\",\"case\":\"" << options.name << "\",\"batch_size\":" << options.batch
              << ",\"width\":" << options.width << ",\"layers\":" << options.layers
              << ",\"heads\":4,\"context\":96,\"vocabulary\":259,\"model_seed\":20261015,\"dropout\":0"
              << ",\"prompt_lengths\":";
    array(workload.lengths); std::cout << ",\"budgets\":"; array(workload.budgets);
    std::cout << ",\"row_seeds\":"; array(workload.seeds);
}
void memory_json(const nn::MemoryStats &before, const nn::MemoryStats &after) {
    require(after.all_buffers.current_bytes == before.all_buffers.current_bytes,
            "generation retained tracked numeric buffers after release");
    std::cout << ",\"memory\":{\"retained_increment_bytes\":0"
              << ",\"all_peak_increment_bytes\":" << after.all_buffers.peak_bytes - before.all_buffers.current_bytes
              << ",\"device_peak_increment_bytes\":" << after.device_buffers.peak_bytes - before.device_buffers.current_bytes
              << ",\"host_peak_increment_bytes\":" << after.host_buffers.peak_bytes - before.host_buffers.current_bytes
              << ",\"readback_peak_increment_bytes\":" << after.host_readback_cache.peak_bytes - before.host_readback_cache.current_bytes
              << '}';
}
void natural(const gen::Model &model, const Work &workload, const Options &options) {
    IDs expected;
    for (int iteration = -options.warmups; iteration < options.repetitions; ++iteration) {
        nn::synchronize(options.device);
        const auto before = nn::memory_stats(); nn::reset_memory_peaks();
        const auto started = Clock::now();
        IDs output;
        if (options.batch == 1)
            output = {gen::generate(model, workload.prompts[0], workload.budgets[0], 0., 0, 1., workload.seeds[0], true)};
        else
            output = gen::generate_batch(model, workload.prompts, workload.budgets, 0., 0, 1., workload.seeds, true, false);
        nn::synchronize(options.device);
        const auto seconds = std::chrono::duration<double>(Clock::now() - started).count();
        const auto after = nn::memory_stats();
        require(output.size() == workload.prompts.size(), "natural generation changed row inventory");
        for (std::size_t row = 0; row < output.size(); ++row)
            require(!output[row].empty() && output[row].size() <= static_cast<std::size_t>(workload.budgets[row]),
                    "natural generation violated an output budget");
        if (expected.empty()) expected = output;
        else require(output == expected, "one runtime produced inconsistent repeated greedy IDs");
        if (iteration < 0) continue;
        prefix(options, workload);
        std::cout << ",\"iteration\":" << iteration << ",\"seconds\":" << seconds << ",\"output_ids\":";
        nested(output);
        std::vector<std::int64_t> counts, ended;
        for (const auto &row : output) { counts.push_back(static_cast<std::int64_t>(row.size())); ended.push_back(row.back() == token::eos_id); }
        std::cout << ",\"actual_new_tokens_per_row\":"; array(counts);
        std::cout << ",\"ended_eos_per_row\":"; array(ended);
        std::cout << ",\"actual_total_new_tokens\":" << std::accumulate(counts.begin(), counts.end(), std::int64_t(0));
        memory_json(before, after);
        std::cout << ",\"statistics\":{\"available\":false}}\n";
    }
}
void forced(const gen::Model &model, const Work &workload, const Options &options) {
    const auto steps = *std::max_element(workload.budgets.begin(), workload.budgets.end()) - 1;
    std::vector<IDs> inputs;
    for (std::int64_t step = 0; step < steps; ++step) inputs.push_back(forced_inputs(workload, step));
    for (int iteration = -options.warmups; iteration < options.repetitions; ++iteration) {
        nn::synchronize(options.device);
        const auto before = nn::memory_stats(); nn::reset_memory_peaks();
        const auto started = Clock::now();
        Statistics stats;
        std::uint64_t bytes = 0, elements = 0, allocations = 0;
        if (options.batch == 1) {
            auto cache = single_cache(model, workload.capacities[0], options.native_plan);
            auto output = cache.append(workload.prompts[0]);
            for (const auto &input : inputs) { output = nn::Tensor(); output = cache.append(input[0]); }
            require(cache.position() == workload.capacities[0], "forced single final position differs");
            stats = statistics(cache); bytes = cache.key_value_bytes(); elements = cache.key_value_elements(); allocations = cache.key_value_allocations();
        } else {
            auto cache = batch_cache(model, workload, options.native_plan);
            auto output = cache.append(workload.prompts);
            if (!inputs.empty()) cache.prepare_decode(static_cast<std::int64_t>(std::count_if(inputs[0].begin(), inputs[0].end(), [](const auto &x) { return !x.empty(); })));
            for (const auto &input : inputs) { output = nn::Tensor(); output = cache.append(input); }
            require(cache.positions() == workload.capacities, "forced batch final positions differ");
            stats = statistics(cache); bytes = cache.key_value_bytes(); elements = cache.key_value_elements(); allocations = cache.key_value_allocations();
        }
        nn::synchronize(options.device);
        const auto seconds = std::chrono::duration<double>(Clock::now() - started).count();
        const auto after = nn::memory_stats();
        if (iteration < 0) continue;
        prefix(options, workload);
        std::cout << ",\"iteration\":" << iteration << ",\"seconds\":" << seconds
                  << ",\"timed_logit_readback\":false,\"forced_next_logit_rows\":"
                  << std::accumulate(workload.budgets.begin(), workload.budgets.end(), std::int64_t(0))
                  << ",\"kv_bytes\":" << bytes << ",\"kv_elements\":" << elements
                  << ",\"kv_allocations\":" << allocations;
        memory_json(before, after);
        std::cout << ",\"statistics\":"; statistics_json(stats); std::cout << "}\n";
    }
}
void oracle(const gen::Model &model, const Work &workload, const Options &options) {
    const auto reference_model = gen::decoder(token::byte_tokenizer(), {96, options.width, 4, options.layers, 20261015, 0});
    std::vector<std::unique_ptr<gen::detail::GenerationCache>> references;
    for (auto capacity : workload.capacities)
        references.push_back(std::make_unique<gen::detail::GenerationCache>(reference_model, true, capacity));
    std::unique_ptr<gen::detail::GenerationCache> single;
    std::unique_ptr<gen::detail::GenerationBatchCache> packed;
    if (options.batch == 1)
        single = std::make_unique<gen::detail::GenerationCache>(single_cache(model, workload.capacities[0], options.native_plan));
    else
        packed = std::make_unique<gen::detail::GenerationBatchCache>(batch_cache(model, workload, options.native_plan));
    const auto steps = *std::max_element(workload.budgets.begin(), workload.budgets.end()) - 1;
    const double tolerance = options.dtype == "Float64" ? 1e-10 : 5e-4;
    for (std::int64_t step = -1; step < steps; ++step) {
        const auto input = step < 0 ? workload.prompts : forced_inputs(workload, step);
        const auto output = options.batch == 1 ? single->append(input[0]) : packed->append(input);
        const auto actual = output.values();
        std::vector<double> expected; std::vector<std::int64_t> active;
        double maximum = 0;
        for (std::size_t row = 0; row < input.size(); ++row) if (!input[row].empty()) {
            const auto next = references[row]->append(input[row]);
            const auto &values = next.values();
            expected.insert(expected.end(), values.begin(), values.end()); active.push_back(static_cast<std::int64_t>(row));
        }
        require(actual.size() == expected.size(), "oracle shape differs");
        for (std::size_t i = 0; i < actual.size(); ++i) {
            require(std::isfinite(actual[i]) && std::isfinite(expected[i]), "oracle observed nonfinite logit");
            maximum = std::max(maximum, std::abs(actual[i] - expected[i]) / (1 + std::abs(expected[i])));
        }
        require(maximum <= tolerance, "forced logits exceed the existing CPU Float64 reference tolerance");
        prefix(options, workload); std::cout << ",\"step\":" << step << ",\"active_rows\":"; array(active);
        std::cout << ",\"max_scaled_logit_error\":" << maximum << ",\"scaled_tolerance\":" << tolerance << ",\"actual_logits\":"; array(actual);
        std::cout << ",\"reference_logits\":"; array(expected); std::cout << "}\n";
    }
}
void weights(const gen::Model &model) {
    auto names = gen::parameter_names(model); std::sort(names.begin(), names.end());
    auto write64 = [](std::uint64_t value) { std::cout.write(reinterpret_cast<const char *>(&value), sizeof(value)); };
    std::cout << "aner.parameters.widened-f64/1\n" << gen::dtype(model) << '\n'; write64(names.size());
    for (const auto &name : names) {
        const auto parameter = gen::parameter(model, name); const auto &values = parameter.values();
        write64(name.size()); std::cout.write(name.data(), static_cast<std::streamsize>(name.size()));
        write64(parameter.rows()); write64(parameter.cols()); write64(values.size());
        std::cout.write(reinterpret_cast<const char *>(values.data()), static_cast<std::streamsize>(values.size() * sizeof(double)));
    }
}
} // namespace

int main(int argc, char **argv) {
    try {
        Options options;
        for (int i = 1; i < argc; i += 2) {
            require(i + 1 < argc, "option needs value");
            const std::string flag = argv[i], value = argv[i + 1];
            if (flag == "--device") options.device = value;
            else if (flag == "--dtype") options.dtype = value;
            else if (flag == "--phase") options.phase = value;
            else if (flag == "--case") options.name = value;
            else if (flag == "--batch-size") options.batch = std::stoll(value);
            else if (flag == "--width") options.width = std::stoll(value);
            else if (flag == "--layers") options.layers = std::stoll(value);
            else if (flag == "--repetitions") options.repetitions = std::stoi(value);
            else if (flag == "--warmups") options.warmups = std::stoi(value);
            else if (flag == "--native-plan") { require(value == "true" || value == "false", "native-plan requires true/false"); options.native_plan = value == "true"; }
            else throw std::runtime_error("unknown option");
        }
        require(options.batch == 1 || options.batch == 8, "batch must be one or eight");
        require((options.width == 32 && options.layers == 2) || (options.width == 128 && options.layers == 3), "unsupported bounded model");
        require(options.repetitions >= 2 && options.repetitions <= 100 && options.warmups >= 1 && options.warmups <= 4, "invalid repeat protocol");
        require(gen::device_supported(options.device, options.dtype), "explicit device/precision unavailable; no fallback");
        const auto workload = work(options);
        const auto model = gen::to(gen::decoder(token::byte_tokenizer(), {96, options.width, 4, options.layers, 20261015, 0}), options.device, options.dtype);
        std::cout << std::setprecision(17);
        if (options.phase == "natural") natural(model, workload, options);
        else if (options.phase == "forced") forced(model, workload, options);
        else if (options.phase == "oracle") oracle(model, workload, options);
        else if (options.phase == "weights") weights(model);
        else throw std::runtime_error("unknown phase");
        return 0;
    } catch (const std::exception &error) {
        std::cerr << error.what() << '\n'; return 1;
    }
}