// sa3-ggml-test.cpp: parity test CLI for the StableStep GGML SA3 modules.
//
// Runs GGML ports of the SA3 conditioning modules against golden vectors
// dumped from the validated ONNX graphs (see tools/onnx-export/), and reports
// cosine similarity + max abs diff per component.
//
// Usage:
//   sa3-ggml-test --models <dir-with-sa3-*.gguf> --goldens <dir-with-manifest.json>
//                 [--component text_enc|seconds|same_enc|same_dec|dit|all]
//
// Components:
//   text_enc: sa3-text-enc-BF16.gguf, T5Gemma encoder + learned padding
//             substitution. Inputs input_ids [1,S] i64 + attention_mask [1,S]
//             u8, expected embeddings [1,S,768] f32.
//   seconds:  sa3-dit-BF16.gguf (embedder tensors only), NumberConditioner
//             expo Fourier embedder. Input [1] f32, expected [1,768] f32.
//   same_enc: sa3-same-enc-F16.gguf, SAME-L encoder. Input audio
//             [1,2,524288] f32, expected latents [1,256,128] f32.
//   same_dec: sa3-same-dec-F16.gguf, SAME-L decoder. Input latents
//             [1,256,128] f32, expected audio [1,2,524288] f32.
//   dit:      sa3-dit-BF16.gguf, DiffusionTransformer single forward. Inputs
//             x [1,256,T] f32, t [1] f32, cross_attn_cond [1,S,768] f32,
//             global_embed [1,768] f32, local_add_cond [1,257,T] f32,
//             padding_mask [1,T] u8; expected v [1,256,T] f32.
//
// Exit code 0 only if every run component passes cosine > 0.999.

#include "sa3-dit-ggml.h"
#include "sa3-same-ggml.h"
#include "sa3-t5gemma-enc.h"
#include "yyjson.h"

#include <cmath>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <string>
#include <vector>

static const double PASS_COSINE = 0.999;

static bool read_file(const std::string & path, std::vector<uint8_t> & out) {
    FILE * f = fopen(path.c_str(), "rb");
    if (!f) {
        fprintf(stderr, "[Test] cannot open %s\n", path.c_str());
        return false;
    }
    fseek(f, 0, SEEK_END);
    long sz = ftell(f);
    fseek(f, 0, SEEK_SET);
    out.resize((size_t) sz);
    size_t rd = fread(out.data(), 1, (size_t) sz, f);
    fclose(f);
    if (rd != (size_t) sz) {
        fprintf(stderr, "[Test] short read on %s\n", path.c_str());
        return false;
    }
    return true;
}

struct Metrics {
    double cosine;
    double max_abs_diff;
};

static Metrics compare(const float * a, const float * b, size_t n) {
    double dot = 0, na = 0, nb = 0, mad = 0;
    for (size_t i = 0; i < n; i++) {
        dot += (double) a[i] * b[i];
        na += (double) a[i] * a[i];
        nb += (double) b[i] * b[i];
        double d = fabs((double) a[i] - b[i]);
        if (d > mad) {
            mad = d;
        }
    }
    Metrics m;
    m.cosine       = (na > 0 && nb > 0) ? dot / (sqrt(na) * sqrt(nb)) : 0.0;
    m.max_abs_diff = mad;
    return m;
}

// Manifest helpers: get golden["<component>"]["inputs"/"outputs"]["<name>"]["file"]
static std::string manifest_file(yyjson_val * root, const char * comp, const char * io, const char * name) {
    yyjson_val * c = yyjson_obj_get(root, comp);
    yyjson_val * g = c ? yyjson_obj_get(c, io) : NULL;
    yyjson_val * t = g ? yyjson_obj_get(g, name) : NULL;
    yyjson_val * f = t ? yyjson_obj_get(t, "file") : NULL;
    return (f && yyjson_is_str(f)) ? yyjson_get_str(f) : "";
}

static int64_t manifest_shape_prod(yyjson_val * root, const char * comp, const char * io, const char * name) {
    yyjson_val * c = yyjson_obj_get(root, comp);
    yyjson_val * g = c ? yyjson_obj_get(c, io) : NULL;
    yyjson_val * t = g ? yyjson_obj_get(g, name) : NULL;
    yyjson_val * s = t ? yyjson_obj_get(t, "shape") : NULL;
    if (!s || !yyjson_is_arr(s)) {
        return 0;
    }
    int64_t      prod = 1;
    size_t       idx, max;
    yyjson_val * d;
    yyjson_arr_foreach(s, idx, max, d) {
        prod *= yyjson_get_int(d);
    }
    return prod;
}

static bool run_text_enc(const std::string & models, const std::string & goldens, yyjson_val * root, Metrics * out) {
    std::string ids_f  = manifest_file(root, "text_enc", "inputs", "input_ids");
    std::string mask_f = manifest_file(root, "text_enc", "inputs", "attention_mask");
    std::string exp_f  = manifest_file(root, "text_enc", "outputs", "embeddings");
    if (ids_f.empty() || mask_f.empty() || exp_f.empty()) {
        fprintf(stderr, "[Test] text_enc: manifest missing entries\n");
        return false;
    }
    int64_t S       = manifest_shape_prod(root, "text_enc", "inputs", "input_ids");
    int64_t n_out   = manifest_shape_prod(root, "text_enc", "outputs", "embeddings");
    if (S <= 0 || n_out <= 0 || n_out % S != 0) {
        fprintf(stderr, "[Test] text_enc: bad shapes in manifest\n");
        return false;
    }
    int64_t H = n_out / S;

    std::vector<uint8_t> ids_raw, mask_raw, exp_raw;
    if (!read_file(goldens + "/" + ids_f, ids_raw) || !read_file(goldens + "/" + mask_f, mask_raw) ||
        !read_file(goldens + "/" + exp_f, exp_raw)) {
        return false;
    }
    if (ids_raw.size() != (size_t) S * 8 || mask_raw.size() != (size_t) S || exp_raw.size() != (size_t) n_out * 4) {
        fprintf(stderr, "[Test] text_enc: golden file sizes do not match manifest shapes\n");
        return false;
    }

    std::vector<int32_t> ids((size_t) S);
    const int64_t *      ids64 = (const int64_t *) ids_raw.data();
    for (int64_t i = 0; i < S; i++) {
        ids[(size_t) i] = (int32_t) ids64[i];
    }

    SA3T5GemmaEnc enc = {};
    if (!sa3_t5gemma_load(&enc, (models + "/sa3-text-enc-BF16.gguf").c_str())) {
        return false;
    }
    if ((int64_t) enc.cfg.hidden_size != H) {
        fprintf(stderr, "[Test] text_enc: model H=%d but golden H=%lld\n", enc.cfg.hidden_size, (long long) H);
        sa3_t5gemma_free(&enc);
        return false;
    }
    const char * env_layers = getenv("SA3_T5G_LAYERS");
    if (env_layers) {
        enc.debug_n_layers = atoi(env_layers);
        fprintf(stderr, "[Test] text_enc: DEBUG truncated to %d layers\n", enc.debug_n_layers);
    }

    std::vector<float> got((size_t) n_out);
    sa3_t5gemma_forward(&enc, ids.data(), mask_raw.data(), (int) S, got.data());
    sa3_t5gemma_free(&enc);

    *out = compare(got.data(), (const float *) exp_raw.data(), (size_t) n_out);
    return true;
}

static bool run_seconds(const std::string & models, const std::string & goldens, yyjson_val * root, Metrics * out) {
    std::string in_f  = manifest_file(root, "seconds", "inputs", "seconds");
    std::string exp_f = manifest_file(root, "seconds", "outputs", "embed");
    if (in_f.empty() || exp_f.empty()) {
        fprintf(stderr, "[Test] seconds: manifest missing entries\n");
        return false;
    }
    std::vector<uint8_t> in_raw, exp_raw;
    if (!read_file(goldens + "/" + in_f, in_raw) || !read_file(goldens + "/" + exp_f, exp_raw)) {
        return false;
    }
    if (in_raw.size() != 4) {
        fprintf(stderr, "[Test] seconds: bad input size\n");
        return false;
    }
    float   seconds = *(const float *) in_raw.data();
    int64_t n_out   = manifest_shape_prod(root, "seconds", "outputs", "embed");
    if (n_out <= 0 || exp_raw.size() != (size_t) n_out * 4) {
        fprintf(stderr, "[Test] seconds: golden size mismatch\n");
        return false;
    }

    SA3SecondsEmbedder emb;
    if (!sa3_seconds_embedder_load(&emb, (models + "/sa3-dit-BF16.gguf").c_str())) {
        return false;
    }
    if ((int64_t) emb.out_dim != n_out) {
        fprintf(stderr, "[Test] seconds: model out=%d but golden %lld\n", emb.out_dim, (long long) n_out);
        return false;
    }
    std::vector<float> got((size_t) n_out);
    sa3_seconds_embed(emb, seconds, got.data());

    *out = compare(got.data(), (const float *) exp_raw.data(), (size_t) n_out);
    return true;
}

// Shared runner for the SAME autoencoder halves. is_encoder selects the
// direction; the golden manifest supplies both tensors' shapes.
static bool run_same(const std::string & models, const std::string & goldens, yyjson_val * root,
                     bool is_encoder, Metrics * out) {
    const char * comp    = is_encoder ? "same_enc" : "same_dec";
    const char * in_name = is_encoder ? "audio" : "latents";
    const char * out_name = is_encoder ? "latents" : "audio";

    std::string in_f  = manifest_file(root, comp, "inputs", in_name);
    std::string exp_f = manifest_file(root, comp, "outputs", out_name);
    if (in_f.empty() || exp_f.empty()) {
        fprintf(stderr, "[Test] %s: manifest missing entries\n", comp);
        return false;
    }
    int64_t n_in  = manifest_shape_prod(root, comp, "inputs", in_name);
    int64_t n_out = manifest_shape_prod(root, comp, "outputs", out_name);
    int64_t n_lat = is_encoder ? n_out : n_in;
    if (n_in <= 0 || n_out <= 0 || n_lat % 256 != 0) {
        fprintf(stderr, "[Test] %s: bad shapes in manifest\n", comp);
        return false;
    }
    int n_latents = (int) (n_lat / 256);  // latent_dim 256

    std::vector<uint8_t> in_raw, exp_raw;
    if (!read_file(goldens + "/" + in_f, in_raw) || !read_file(goldens + "/" + exp_f, exp_raw)) {
        return false;
    }
    if (in_raw.size() != (size_t) n_in * 4 || exp_raw.size() != (size_t) n_out * 4) {
        fprintf(stderr, "[Test] %s: golden file sizes do not match manifest shapes\n", comp);
        return false;
    }

    SA3Same same = {};
    // F16 is the current conversion (see convert-sa3.py); fall back to the
    // older BF16 name if that is what is on disk.
    std::string gguf = models + (is_encoder ? "/sa3-same-enc-F16.gguf" : "/sa3-same-dec-F16.gguf");
    {
        FILE * f = fopen(gguf.c_str(), "rb");
        if (f) {
            fclose(f);
        } else {
            gguf = models + (is_encoder ? "/sa3-same-enc-BF16.gguf" : "/sa3-same-dec-BF16.gguf");
        }
    }
    if (!sa3_same_load(&same, gguf.c_str(), is_encoder)) {
        return false;
    }

    // Debug hooks: SA3_SAME_STAGE=<n> dumps the token sequence after n layers
    // (0 = folded input) to SA3_SAME_DUMP (default sa3_same_stage.bin).
    const char * env_stage = getenv("SA3_SAME_STAGE");
    if (env_stage) {
        same.debug_stage = atoi(env_stage);
        fprintf(stderr, "[Test] %s: DEBUG dumping stage %d\n", comp, same.debug_stage);
    }

    std::vector<float> got((size_t) n_out);
    sa3_same_forward(&same, (const float *) in_raw.data(), got.data(), n_latents);

    if (env_stage && !same.debug_out.empty()) {
        const char * dump = getenv("SA3_SAME_DUMP");
        std::string  path = dump ? dump : "sa3_same_stage.bin";
        FILE *       f    = fopen(path.c_str(), "wb");
        if (f) {
            fwrite(same.debug_out.data(), sizeof(float), same.debug_out.size(), f);
            fclose(f);
            fprintf(stderr, "[Test] %s: stage tensor (%zu floats) -> %s\n", comp, same.debug_out.size(),
                    path.c_str());
        }
    }
    sa3_same_free(&same);

    *out = compare(got.data(), (const float *) exp_raw.data(), (size_t) n_out);
    return true;
}

static bool run_dit(const std::string & models, const std::string & goldens, yyjson_val * root, Metrics * out) {
    const char * in_names[] = { "x", "t", "cross_attn_cond", "global_embed", "local_add_cond", "padding_mask" };
    std::vector<std::vector<uint8_t>> raw(7);
    for (int i = 0; i < 6; i++) {
        std::string f = manifest_file(root, "dit", "inputs", in_names[i]);
        if (f.empty() || !read_file(goldens + "/" + f, raw[(size_t) i])) {
            fprintf(stderr, "[Test] dit: missing input '%s'\n", in_names[i]);
            return false;
        }
    }
    std::string exp_f = manifest_file(root, "dit", "outputs", "v");
    if (exp_f.empty() || !read_file(goldens + "/" + exp_f, raw[6])) {
        fprintf(stderr, "[Test] dit: missing output 'v'\n");
        return false;
    }

    int64_t n_x     = manifest_shape_prod(root, "dit", "inputs", "x");
    int64_t n_cross = manifest_shape_prod(root, "dit", "inputs", "cross_attn_cond");
    int64_t n_glob  = manifest_shape_prod(root, "dit", "inputs", "global_embed");
    int64_t n_local = manifest_shape_prod(root, "dit", "inputs", "local_add_cond");
    int64_t T       = manifest_shape_prod(root, "dit", "inputs", "padding_mask");
    int64_t n_out   = manifest_shape_prod(root, "dit", "outputs", "v");
    if (T <= 0 || n_x != 256 * T || n_out != n_x || n_glob != 768 || n_cross % 768 != 0 ||
        n_local != 257 * T) {
        fprintf(stderr, "[Test] dit: bad shapes in manifest\n");
        return false;
    }
    int64_t S_c = n_cross / 768;
    if (raw[0].size() != (size_t) n_x * 4 || raw[1].size() != 4 || raw[2].size() != (size_t) n_cross * 4 ||
        raw[3].size() != (size_t) n_glob * 4 || raw[4].size() != (size_t) n_local * 4 ||
        raw[5].size() != (size_t) T || raw[6].size() != (size_t) n_out * 4) {
        fprintf(stderr, "[Test] dit: golden file sizes do not match manifest shapes\n");
        return false;
    }

    SA3DiT dit = {};
    if (!sa3_dit_load(&dit, (models + "/sa3-dit-BF16.gguf").c_str())) {
        return false;
    }

    // Debug hook: SA3_DIT_STAGE=<n> dumps the token sequence after n layers
    // (0 = memory+projected input) to SA3_DIT_DUMP (default sa3_dit_stage.bin).
    const char * env_stage = getenv("SA3_DIT_STAGE");
    if (env_stage) {
        dit.debug_stage = atoi(env_stage);
        fprintf(stderr, "[Test] dit: DEBUG dumping stage %d\n", dit.debug_stage);
    }

    std::vector<float> got((size_t) n_out);
    sa3_dit_forward(&dit, (const float *) raw[0].data(), *(const float *) raw[1].data(),
                    (const float *) raw[2].data(), S_c, (const float *) raw[3].data(),
                    (const float *) raw[4].data(), raw[5].data(), T, got.data());

    if (env_stage && !dit.debug_out.empty()) {
        const char * dump = getenv("SA3_DIT_DUMP");
        std::string  path = dump ? dump : "sa3_dit_stage.bin";
        FILE *       f    = fopen(path.c_str(), "wb");
        if (f) {
            fwrite(dit.debug_out.data(), sizeof(float), dit.debug_out.size(), f);
            fclose(f);
            fprintf(stderr, "[Test] dit: stage tensor (%zu floats) -> %s\n", dit.debug_out.size(),
                    path.c_str());
        }
    }
    sa3_dit_free(&dit);

    *out = compare(got.data(), (const float *) raw[6].data(), (size_t) n_out);
    return true;
}

int main(int argc, char ** argv) {
    std::string models, goldens, component = "all";
    for (int i = 1; i < argc; i++) {
        if (!strcmp(argv[i], "--models") && i + 1 < argc) {
            models = argv[++i];
        } else if (!strcmp(argv[i], "--goldens") && i + 1 < argc) {
            goldens = argv[++i];
        } else if (!strcmp(argv[i], "--component") && i + 1 < argc) {
            component = argv[++i];
        } else {
            fprintf(stderr,
                    "Usage: sa3-ggml-test --models <dir> --goldens <dir> [--component text_enc|seconds|same_enc|same_dec|dit|all]\n");
            return 2;
        }
    }
    if (models.empty() || goldens.empty()) {
        fprintf(stderr, "Usage: sa3-ggml-test --models <dir> --goldens <dir> [--component text_enc|seconds|same_enc|same_dec|dit|all]\n");
        return 2;
    }

    std::vector<uint8_t> manifest_raw;
    if (!read_file(goldens + "/manifest.json", manifest_raw)) {
        return 2;
    }
    yyjson_doc * doc = yyjson_read((const char *) manifest_raw.data(), manifest_raw.size(), 0);
    if (!doc) {
        fprintf(stderr, "[Test] cannot parse manifest.json\n");
        return 2;
    }
    yyjson_val * root = yyjson_doc_get_root(doc);

    bool all_pass = true;
    bool any_run  = false;

    if (component == "all" || component == "text_enc") {
        Metrics m;
        any_run = true;
        if (run_text_enc(models, goldens, root, &m)) {
            bool pass = m.cosine > PASS_COSINE;
            printf("text_enc: cosine=%.6f max_abs_diff=%.6f  %s\n", m.cosine, m.max_abs_diff,
                   pass ? "PASS" : "FAIL");
            all_pass = all_pass && pass;
        } else {
            printf("text_enc: ERROR\n");
            all_pass = false;
        }
    }
    if (component == "all" || component == "seconds") {
        Metrics m;
        any_run = true;
        if (run_seconds(models, goldens, root, &m)) {
            bool pass = m.cosine > PASS_COSINE;
            printf("seconds:  cosine=%.6f max_abs_diff=%.6f  %s\n", m.cosine, m.max_abs_diff,
                   pass ? "PASS" : "FAIL");
            all_pass = all_pass && pass;
        } else {
            printf("seconds:  ERROR\n");
            all_pass = false;
        }
    }
    if (component == "all" || component == "same_enc") {
        Metrics m;
        any_run = true;
        if (run_same(models, goldens, root, true, &m)) {
            bool pass = m.cosine > PASS_COSINE;
            printf("same_enc: cosine=%.6f max_abs_diff=%.6f  %s\n", m.cosine, m.max_abs_diff,
                   pass ? "PASS" : "FAIL");
            all_pass = all_pass && pass;
        } else {
            printf("same_enc: ERROR\n");
            all_pass = false;
        }
    }
    if (component == "all" || component == "same_dec") {
        Metrics m;
        any_run = true;
        if (run_same(models, goldens, root, false, &m)) {
            bool pass = m.cosine > PASS_COSINE;
            printf("same_dec: cosine=%.6f max_abs_diff=%.6f  %s\n", m.cosine, m.max_abs_diff,
                   pass ? "PASS" : "FAIL");
            all_pass = all_pass && pass;
        } else {
            printf("same_dec: ERROR\n");
            all_pass = false;
        }
    }
    if (component == "all" || component == "dit") {
        Metrics m;
        any_run = true;
        if (run_dit(models, goldens, root, &m)) {
            bool pass = m.cosine > PASS_COSINE;
            printf("dit:      cosine=%.6f max_abs_diff=%.6f  %s\n", m.cosine, m.max_abs_diff,
                   pass ? "PASS" : "FAIL");
            all_pass = all_pass && pass;
        } else {
            printf("dit:      ERROR\n");
            all_pass = false;
        }
    }
    yyjson_doc_free(doc);

    if (!any_run) {
        fprintf(stderr, "[Test] unknown component '%s'\n", component.c_str());
        return 2;
    }
    return all_pass ? 0 : 1;
}