// test-fusion: verify the backend fusion logic against a per-device baseline.
//
// for every dummy model generated by test-llama-archs, the tool runs the model on a single
// device with fusion enabled and disabled, and reports:
//   - the per-fusion-type counters for each mode (prefill / decode, merged into "any" when the
//     per-graph counts match)
//   - the NMSE between the fused and unfused logits
//   - the NMSE between the device and a CPU reference
//
// the per-fusion-type counters are compared against a per-device baseline file (CSV) so a
// fusion pattern that silently stops matching (or fires when it should not) is caught as a
// regression.
//
// usage:
//   test-fusion --models DIR --device MTL0 --record baseline.csv   # generate a baseline
//   test-fusion --models DIR --device MTL0 --check  baseline.csv   # validate against it
//   test-fusion --model FILE --device MTL0 --check  baseline.csv   # validate a single model

#include "common.h"
#include "log.h"
#include "llama-cpp.h"

#include "ggml.h"
#include "gguf.h"

#include <algorithm>
#include <array>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <iomanip>
#include <iostream>
#include <map>
#include <random>
#include <string>
#include <vector>

// generic fusion debugging API, resolved through the ad-hoc get_proc_address mechanism
// (not part of the official ggml backend interface yet). a backend that adopts fusion debugging
// exports these exact names.
typedef void * ggml_backend_fusion_t;

typedef ggml_backend_fusion_t ( * fusion_get_t)       (ggml_backend_dev_t);
typedef void ( * fusion_stats_init_t)  (ggml_backend_fusion_t);
typedef void ( * fusion_stats_reset_t) (ggml_backend_fusion_t);
typedef int  ( * fusion_stats_get_t)   (ggml_backend_fusion_t, const char **, uint64_t *, int);
typedef void ( * fusion_set_enabled_t) (ggml_backend_fusion_t, bool);

static bool silent_model_load_progress(float, void *) {
    return true;
}

struct gguf_context_ptr {
    gguf_context * ctx;
    gguf_context_ptr(gguf_context * c) : ctx(c) {}
    ~gguf_context_ptr() { if (ctx) { gguf_free(ctx); } }
    gguf_context * get() const { return ctx; }
    gguf_context_ptr(const gguf_context_ptr &) = delete;
    gguf_context_ptr & operator=(const gguf_context_ptr &) = delete;
};

// NMSE between two vectors (same as tests/test-llama-archs.cpp)
static double nmse(const std::vector<float> & a, const std::vector<float> & b) {
    GGML_ASSERT(a.size() == b.size());
    double mse_a_b = 0.0;
    double mse_a_0 = 0.0;

    for (size_t i = 0; i < a.size(); i++) {
        const float a_i = a[i];
        const float b_i = b[i];

        mse_a_b += (a_i - b_i) * (a_i - b_i);
        mse_a_0 += a_i * a_i;
    }

    return mse_a_b / mse_a_0;
}

// deterministic token sequence
static std::vector<llama_token> get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed) {
    std::mt19937 gen(seed);
    std::uniform_int_distribution<> dis(0, n_vocab - 1);
    std::vector<llama_token> ret;
    ret.reserve(n_tokens);
    for (uint32_t i = 0; i < n_tokens; i++) {
        ret.push_back(dis(gen));
    }
    return ret;
}

// trim leading/trailing whitespace (used when parsing padded CSV columns)
static std::string trim(const std::string & s) {
    const size_t b = s.find_first_not_of(" \t\r\n");
    if (b == std::string::npos) {
        return "";
    }
    const size_t e = s.find_last_not_of(" \t\r\n");
    return s.substr(b, e - b + 1);
}

static std::string get_arch(const std::string & path) {
    gguf_init_params params = { /*no_alloc=*/true, /*ctx=*/nullptr };
    gguf_context_ptr ctx(gguf_init_from_file(path.c_str(), params));
    if (!ctx.get()) {
        throw std::runtime_error("failed to read gguf: " + path);
    }
    const int idx = gguf_find_key(ctx.get(), "general.architecture");
    if (idx < 0) {
        return "unknown";
    }
    const char * val = gguf_get_val_str(ctx.get(), idx);
    return val ? val : "unknown";
}

static llama_model_ptr load_model(const std::string & path, ggml_backend_dev_t dev) {
    llama_model_params model_params = llama_model_default_params();
    model_params.progress_callback = silent_model_load_progress;
    std::vector<ggml_backend_dev_t> devs = { dev, nullptr };
    model_params.devices = devs.data();
    model_params.split_mode = LLAMA_SPLIT_MODE_LAYER;

    llama_model_ptr model(llama_model_load_from_file(path.c_str(), model_params));
    if (!model) {
        throw std::runtime_error("failed to load model: " + path);
    }
    return model;
}

// a fresh context (fresh state) from an already-loaded model
static llama_context_ptr create_ctx(llama_model * model, int n_ubatch) {
    llama_context_params ctx_params = llama_context_default_params();
    ctx_params.n_ctx = 0;
    ctx_params.n_threads = 4;
    ctx_params.n_threads_batch = 4;
    ctx_params.n_ubatch = n_ubatch;
    ctx_params.n_batch = n_ubatch;

    llama_context_ptr lctx(llama_init_from_model(model, ctx_params));
    if (!lctx) {
        throw std::runtime_error("failed to init context");
    }
    return lctx;
}

// decode all tokens in one batch; returns the logits of every token
static std::vector<float> decode_prefill(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
    const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
    common_batch batch(lctx);
    for (size_t i = 0; i < tokens.size(); i++) {
        batch.add(tokens[i], i, 0, true);
    }
    if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
        throw std::runtime_error("prefill decode failed");
    }

    std::vector<float> ret;
    ret.reserve(tokens.size() * n_vocab);
    for (size_t i = 0; i < tokens.size(); i++) {
        const float * logits_ith = llama_get_logits_ith(lctx, i);
        for (uint32_t j = 0; j < n_vocab; j++) {
            ret.push_back(logits_ith[j]);
        }
    }
    return ret;
}

// decode one token at a time; returns the logits of the last token of each step
static std::vector<float> decode_gen(llama_model * model, llama_context * lctx, const std::vector<llama_token> & tokens) {
    const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model));
    common_batch batch(lctx);
    std::vector<float> ret;
    for (size_t i = 0; i < tokens.size(); i++) {
        batch.clear();
        batch.add(tokens[i], i, 0, true);
        if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch.get())) {
            throw std::runtime_error("decode failed");
        }
        const float * logits = llama_get_logits_ith(lctx, 0);
        for (uint32_t j = 0; j < n_vocab; j++) {
            ret.push_back(logits[j]);
        }
    }
    return ret;
}

static void read_counts(fusion_stats_get_t api_stats_get, ggml_backend_fusion_t finfo,
                        std::vector<const char *> & labels, std::vector<uint64_t> & counts) {
    const int n = api_stats_get(finfo, nullptr, nullptr, 0);
    labels.assign(n, nullptr);
    counts.assign(n, 0);
    api_stats_get(finfo, labels.data(), counts.data(), n);
}

// one row of the per-label report
struct fusion_row {
    std::string  arch;
    bool         moe;
    std::string  mode;
    std::string  label;
    uint64_t     count_fused;
    uint64_t     count_unfused;
    uint64_t     expected;
    double       nmse_fus;
    double       nmse_dev;
    bool         ok_count; // counts match the baseline
    bool         ok_nmse;  // nmse within epsilon
};

static void usage(const char * argv0) {
    printf("%s: verify fusion counts on a device against a per-device baseline\n\n", argv0);
    printf("usage: %s [options]\n\n", argv0);
    printf("options:\n");
    printf("  --models DIR   run over all .gguf models in a directory\n");
    printf("  --model FILE   run over a single model file (mutually exclusive with --models)\n");
    printf("  --device NAME  device to run on (e.g. MTL0, CPU)\n");
    printf("  --record CSV   write the golden baseline\n");
    printf("  --check  CSV   validate the counters against a baseline (default)\n");
    printf("  -h, --help     show this message and exit\n");
}

int main(int argc, char ** argv) {
    std::string models_dir;
    std::string model_file;
    std::string device_name;
    std::string record_path;
    std::string check_path;

    for (int i = 1; i < argc; i++) {
        const std::string arg = argv[i];
        const auto next = [&](const char * name) -> std::string {
            if (i + 1 >= argc) {
                LOG_ERR("%s: %s requires an argument\n", __func__, name);
                exit(1);
            }
            return argv[++i];
        };
        if (arg == "-h" || arg == "--help") {
            usage(argv[0]);
            exit(0);
        }
        if (arg == "--models")     { models_dir  = next("--models"); }
        else if (arg == "--model") { model_file  = next("--model"); }
        else if (arg == "--device"){ device_name = next("--device"); }
        else if (arg == "--record"){ record_path = next("--record"); }
        else if (arg == "--check") { check_path  = next("--check"); }
        else {
            LOG_ERR("%s: unknown argument: %s\n", __func__, arg.c_str());
            return 1;
        }
    }

    if (device_name.empty()) {
        LOG_ERR("%s: --device NAME is required\n", __func__);
        return 1;
    }
    if (models_dir.empty() && model_file.empty()) {
        LOG_ERR("%s: --models DIR or --model FILE is required\n", __func__);
        return 1;
    }
    if (!models_dir.empty() && !model_file.empty()) {
        LOG_ERR("%s: --models DIR and --model FILE are mutually exclusive\n", __func__);
        return 1;
    }
    if (!record_path.empty() && !check_path.empty()) {
        LOG_ERR("%s: --record and --check are mutually exclusive\n", __func__);
        return 1;
    }

    std::vector<std::string> models;
    if (!model_file.empty()) {
        if (!std::filesystem::is_regular_file(model_file)) {
            LOG_ERR("%s: model file '%s' does not exist\n", __func__, model_file.c_str());
            return 1;
        }
        models.push_back(model_file);
    } else {
        if (!std::filesystem::exists(models_dir) || !std::filesystem::is_directory(models_dir)) {
            LOG_ERR("%s: models directory '%s' does not exist\n", __func__, models_dir.c_str());
            return 1;
        }
        for (const auto & entry : std::filesystem::directory_iterator(models_dir)) {
            if (entry.is_regular_file() && entry.path().extension() == ".gguf") {
                models.push_back(entry.path().string());
            }
        }
        std::sort(models.begin(), models.end());

        if (models.empty()) {
            LOG_ERR("%s: no .gguf models found in '%s'\n", __func__, models_dir.c_str());
            return 1;
        }
    }

    common_init();
    llama_backend_init();

    ggml_backend_dev_t dev = ggml_backend_dev_by_name(device_name.c_str());
    if (!dev) {
        LOG_WRN("%s: device '%s' not found - skipping (baseline is device-specific)\n",
                __func__, device_name.c_str());
        return 0;
    }

    // resolve the generic fusion debugging functions through the ad-hoc get_proc_address
    // mechanism; a backend that does not adopt fusion debugging exports none of them
    auto * reg = ggml_backend_dev_backend_reg(dev);

    // output naming uses the backend base name (e.g. "MTL") rather than the specific device
    // name (e.g. "MTL0") the test was invoked with
    const std::string base_name = ggml_backend_reg_name(reg);

    auto api_get         = (fusion_get_t)         ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_get");
    auto api_stats_init  = (fusion_stats_init_t)  ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_init");
    auto api_stats_reset = (fusion_stats_reset_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_reset");
    auto api_stats_get   = (fusion_stats_get_t)   ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_stats_get");
    auto api_set_enabled = (fusion_set_enabled_t) ggml_backend_reg_get_proc_address(reg, "ggml_backend_fusion_set_enabled");

    if (!api_get || !api_stats_init || !api_set_enabled || !api_stats_reset || !api_stats_get) {
        LOG_ERR("%s: device '%s' does not export the generic fusion debugging API "
                "(ggml_backend_fusion_*) - cannot run the fusion regression test\n",
                __func__, device_name.c_str());
        return 1;
    }

    ggml_backend_fusion_t finfo = api_get(dev);

    // enable fusions stats
    api_stats_init(finfo);

    const bool has_counts = true;

    // load the baseline (if any): key arch|moe|mode|label -> expected count
    std::map<std::string, uint64_t> baseline;
    if (!check_path.empty()) {
        std::ifstream in(check_path);
        if (!in) {
            LOG_ERR("%s: cannot open baseline '%s'\n", __func__, check_path.c_str());
            return 1;
        }
        std::string line;
        while (std::getline(in, line)) {
            if (line.empty() || line[0] == '#') {
                continue;
            }
            std::vector<std::string> cols;
            size_t pos = 0;
            while ((pos = line.find(',')) != std::string::npos) {
                cols.push_back(trim(line.substr(0, pos)));
                line.erase(0, pos + 1);
            }
            cols.push_back(trim(line));
            if (cols.size() != 5) {
                continue;
            }
            baseline[cols[0] + "|" + cols[1] + "|" + cols[2] + "|" + cols[3]] = std::stoull(cols[4]);
        }
    }

    std::vector<fusion_row> rows;

    LOG_INF("%s: running fusion test over %zu models on '%s'\n", __func__, models.size(), base_name.c_str());

    const size_t seed = 1;

    for (const auto & model_path : models) {
        const std::string arch = get_arch(model_path);
        const bool moe = arch.find("moe") != std::string::npos;

        llama_model_ptr model;
        llama_model_ptr model_cpu;
        uint32_t n_vocab = 0;
        try {
            model = load_model(model_path, dev);
            model_cpu = load_model(model_path, ggml_backend_dev_by_name("CPU"));
            n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get()));
        } catch (const std::exception & e) {
            LOG_ERR("%s: %s: %s\n", __func__, model_path.c_str(), e.what());
            continue;
        }

        struct mode_cfg {
            std::string name;
            std::vector<float> (*decode)(llama_model *, llama_context *, const std::vector<llama_token> &);
            int n_tokens;
            int n_graphs;   // graph runs per mode (prefill=1, decode=16)
        };
        const mode_cfg modes[] = {
            { "prefill", decode_prefill, 32, 1  },
            { "decode",  decode_gen,     16, 16 },
        };

        // per-label, per-mode data for this model; prefill and decode are merged into a single
        // "any" row when their per-graph counts match
        struct mode_data {
            bool     present;
            uint64_t count_fused;   // per graph
            uint64_t count_unfused; // per graph
            double   nmse_fus;
            double   nmse_dev;
            bool     ok_nmse;
        };
        std::map<std::string, std::array<mode_data, 2>> mdata;

        for (int mi = 0; mi < 2; mi++) {
            const mode_cfg & mode = modes[mi];
            const auto tokens = get_tokens(mode.n_tokens, n_vocab, seed);

            // CPU reference for this mode (fresh context, fresh state)
            std::vector<float> logits_cpu;
            try {
                llama_context_ptr ctx = create_ctx(model_cpu.get(), 32);
                logits_cpu = mode.decode(model_cpu.get(), ctx.get(), tokens);
            } catch (const std::exception & e) {
                LOG_WRN("%s: %s: cpu reference: %s\n", __func__, model_path.c_str(), e.what());
            }

            // fused run on a fresh context (fresh state)
            std::vector<float> logits_fused;
            std::vector<const char *> labels;
            std::vector<uint64_t> counts_fused;
            {
                llama_context_ptr ctx = create_ctx(model.get(), 32);
                if (has_counts) {
                    api_set_enabled(finfo, true);
                    api_stats_reset(finfo);
                }
                logits_fused = mode.decode(model.get(), ctx.get(), tokens);
                if (has_counts) {
                    read_counts(api_stats_get, finfo, labels, counts_fused);
                }
            }

            // unfused run on another fresh context (fresh state)
            std::vector<float> logits_unfused;
            std::vector<uint64_t> counts_unfused;
            {
                llama_context_ptr ctx = create_ctx(model.get(), 32);
                if (has_counts) {
                    api_set_enabled(finfo, false);
                    api_stats_reset(finfo);
                }
                logits_unfused = mode.decode(model.get(), ctx.get(), tokens);
                if (has_counts) {
                    read_counts(api_stats_get, finfo, labels, counts_unfused);
                }
            }

            const double nmse_fus = nmse(logits_fused, logits_unfused);
            const double nmse_dev = logits_cpu.empty() ? 0.0 : nmse(logits_fused, logits_cpu);

            if (has_counts) {
                for (int i = 0; i < (int) labels.size(); i++) {
                    const uint64_t fused   = counts_fused[i]   / mode.n_graphs;
                    const uint64_t unfused = counts_unfused[i] / mode.n_graphs;
                    if (fused == 0 && unfused == 0) {
                        continue;
                    }
                    auto & d = mdata[labels[i]][mi];
                    d.present       = true;
                    d.count_fused   = fused;
                    d.count_unfused = unfused;
                    d.nmse_fus      = nmse_fus;
                    d.nmse_dev      = nmse_dev;
                    d.ok_nmse       = nmse_fus <= 1e-4;
                }
            } else {
                rows.push_back({ arch, moe, mode.name, "?", 0, 0, 0, nmse_fus, nmse_dev, true, nmse_fus <= 1e-4 });
            }
        }

        // build the per-label rows, merging prefill and decode into "any" when the per-graph
        // counts match (they always do for the deterministic fusion table)
        if (has_counts) {
            for (auto & kv : mdata) {
                const std::string & label = kv.first;
                const auto & d = kv.second;
                const bool both = d[0].present && d[1].present;
                const bool match = both && d[0].count_fused == d[1].count_fused;

                if (match) {
                    // one "any" row; use the worst NMSE across the two modes
                    const std::string any_key = arch + "|" + (moe ? "1" : "0") + "|any|" + label;
                    const uint64_t expected = baseline.count(any_key) ? baseline.at(any_key) : 0;
                    const bool ok_count = check_path.empty() || d[0].count_fused == expected;
                    const bool ok_nmse   = d[0].ok_nmse && d[1].ok_nmse;
                    const double nmse_fus = std::max(d[0].nmse_fus, d[1].nmse_fus);
                    const double nmse_dev = std::max(d[0].nmse_dev, d[1].nmse_dev);
                    rows.push_back({ arch, moe, "any", label, d[0].count_fused, d[0].count_unfused,
                                     expected, nmse_fus, nmse_dev, ok_count, ok_nmse });
                } else {
                    // counts differ - keep a separate row per mode
                    for (int mi = 0; mi < 2; mi++) {
                        if (!d[mi].present) {
                            continue;
                        }
                        const mode_data & a = d[mi];
                        const std::string mode_key = arch + "|" + (moe ? "1" : "0") + "|" + modes[mi].name + "|" + label;
                        const uint64_t expected = baseline.count(mode_key) ? baseline.at(mode_key) : 0;
                        const bool ok_count = check_path.empty() || a.count_fused == expected;
                        rows.push_back({ arch, moe, modes[mi].name, label, a.count_fused, a.count_unfused,
                                         expected, a.nmse_fus, a.nmse_dev, ok_count, a.ok_nmse });
                    }
                }
            }
        }

        LOG_INF("%s: %-20s (%s) done\n", __func__, arch.c_str(), model_path.c_str());
    }

    // print the report
    {
        std::ofstream out(record_path);
        std::ostream & os = record_path.empty() ? std::cout : out;
        if (!record_path.empty()) {
            os << "# test-fusion baseline for device " << base_name << "\n";
            os << "# " << std::left
               << std::setw(18) << "arch"  << ','
               << std::setw(4)  << "moe"   << ','
               << std::setw(8)  << "mode"  << ','
               << std::setw(28) << "label" << ','
               << std::right << std::setw(7) << "count" << '\n';
        }

        LOG_INF("%-20s %-4s %-8s %-22s %7s %7s %7s %10s %10s %s\n",
                "arch", "moe", "mode", "label", "fused", "unfused", "expected", "nmse_fus", "nmse_dev", "status");
        int n_ok = 0;
        int n_bad = 0;
        for (const auto & r : rows) {
            const bool ok = r.ok_count && r.ok_nmse;
            const char * status = ok ? "ok" : "FAIL";
            if (ok) { n_ok++; } else { n_bad++; }
            LOG_INF("%-20s %-4s %-8s %-22s %7llu %7llu %7llu %10.2e %10.2e %s\n",
                    r.arch.c_str(), r.moe ? "moe" : "dense", r.mode.c_str(), r.label.c_str(),
                    (unsigned long long) r.count_fused, (unsigned long long) r.count_unfused,
                    (unsigned long long) r.expected, r.nmse_fus, r.nmse_dev, status);
            if (!record_path.empty()) {
                os << std::left
                   << std::setw(20) << r.arch << ','
                   << std::setw(4)  << (r.moe ? "1" : "0") << ','
                   << std::setw(8)  << r.mode << ','
                   << std::setw(28) << r.label << ','
                   << std::right << std::setw(7) << r.count_fused << '\n';
            }
        }
        LOG_INF("summary: %d ok, %d failed\n", n_ok, n_bad);
        if (!record_path.empty()) {
            LOG_INF("%s: baseline written to '%s'\n", __func__, record_path.c_str());
        }

        if (n_bad && !models_dir.empty() && !check_path.empty()) {
            LOG_WRN("%s: if the fusion counts are expected to change, run with --record to update the baseline:\n"
                    "\n"
                    "./bin/test-llama-archs -o %s\n"
                    "%s --device %s --models %s --record %s\n",
                    __func__, models_dir.c_str(), argv[0], device_name.c_str(), models_dir.c_str(), check_path.c_str());
        }

        return n_bad;
    }
}