// DFly (AngelSpec) encoder context fusion: numerical parity + memory-layout contract.
//
// The draft-time graph in src/models/dflash.cpp mixes the raw per-target-layer features into
// one context PER DRAFT LAYER, which needs two reshape/permute round trips to put the
// contracted axis first. Both the axis handling and the resulting flat layout are easy to get
// subtly wrong in a way that still loads and still produces plausible drafts, so this test
// pins them against an independent scalar reference.
//
// Checks:
//   1. fused context == softmax-weighted mix of the raw features + shared base projection
//   2. the flattened [n_embd*n_layer, n_tokens] output is LAYER-MAJOR within each token,
//      i.e. the decoder's ggml_view_2d(offset = il*n_embd, stride = nb[1]) recovers layer il

#include "ggml.h"
#include "ggml-cpu.h"

#include <cmath>
#include <cstdio>
#include <cstdlib>
#include <vector>

static void graph_compute(ggml_cgraph * gf, int n_threads) {
    std::vector<uint8_t> buf;
    ggml_cplan plan = ggml_graph_plan(gf, n_threads, nullptr);
    if (plan.work_size > 0) {
        buf.resize(plan.work_size);
        plan.work_data = buf.data();
    }
    ggml_graph_compute(gf, &plan);
}

int main() {
    // deliberately non-square and n_feat != n_layer so a transposed axis cannot pass
    const int64_t n_embd   = 8;
    const int64_t n_feat   = 5;   // target capture layers (DFly Qwen3-8B ships 5)
    const int64_t n_layer  = 3;   // draft layers
    const int64_t n_tokens = 4;
    const float   eps      = 1e-6f;

    std::vector<uint8_t> mem(64u*1024u*1024u);
    ggml_init_params ip = { mem.size(), mem.data(), false };
    ggml_context * ctx = ggml_init(ip);

    ggml_tensor * inp    = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_feat*n_embd, n_tokens);
    ggml_tensor * fc     = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_feat*n_embd, n_embd);
    ggml_tensor * fusion = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_feat, n_layer);
    ggml_tensor * cnorm  = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd);

    srand(1234);
    auto fill = [](ggml_tensor * t) {
        float * d = (float *) t->data;
        for (int64_t i = 0; i < ggml_nelements(t); ++i) {
            d[i] = 2.0f*((float) rand()/(float) RAND_MAX) - 1.0f;
        }
    };
    fill(inp); fill(fc); fill(fusion); fill(cnorm);

    // ---- the graph under test: mirrors llama_model_dflash::graph<true> (DFly branch) ----
    ggml_tensor * base = ggml_mul_mat(ctx, fc, inp);                       // [n_embd, n_tokens]

    ggml_tensor * probs = ggml_soft_max(ctx, fusion);                      // [n_feat, n_layer]

    ggml_tensor * feats = ggml_cont(ctx, ggml_permute(ctx,
                ggml_reshape_3d(ctx, inp, n_embd, n_feat, n_tokens), 1, 0, 2, 3));

    ggml_tensor * resid = ggml_mul_mat(ctx, probs,
                ggml_reshape_2d(ctx, feats, n_feat, n_embd*n_tokens));     // [n_layer, n_embd*n_tokens]

    resid = ggml_cont(ctx, ggml_permute(ctx,
                ggml_reshape_3d(ctx, resid, n_layer, n_embd, n_tokens), 1, 0, 2, 3));

    ggml_tensor * cur = ggml_add(ctx, resid, ggml_reshape_3d(ctx, base, n_embd, 1, n_tokens));
    cur = ggml_mul(ctx, ggml_rms_norm(ctx, cur, eps), cnorm);
    cur = ggml_reshape_2d(ctx, cur, n_embd*n_layer, n_tokens);

    // the decoder's per-layer slice of the round-tripped row (layer 1, arbitrary interior pick)
    const int64_t il_probe = 1;
    ggml_tensor * slice = ggml_cont(ctx, ggml_view_2d(ctx, cur, n_embd, n_tokens,
                cur->nb[1], (size_t) il_probe*n_embd*cur->nb[0]));

    ggml_cgraph * gf = ggml_new_graph(ctx);
    ggml_build_forward_expand(gf, cur);
    ggml_build_forward_expand(gf, slice);
    graph_compute(gf, 2);

    // ---- independent scalar reference ----
    const float * pinp = (const float *) inp->data;
    const float * pfc  = (const float *) fc->data;
    const float * pfus = (const float *) fusion->data;
    const float * pcn  = (const float *) cnorm->data;

    std::vector<float> ref((size_t) n_embd*n_layer*n_tokens);

    for (int64_t n = 0; n < n_tokens; ++n) {
        // shared base projection: base[h] = sum_r fc[r,h] * inp[r,n]
        std::vector<float> b(n_embd, 0.0f);
        for (int64_t h = 0; h < n_embd; ++h) {
            for (int64_t r = 0; r < n_feat*n_embd; ++r) {
                b[h] += pfc[r + h*n_feat*n_embd] * pinp[r + n*n_feat*n_embd];
            }
        }

        for (int64_t l = 0; l < n_layer; ++l) {
            // softmax over this draft layer's n_feat mixing logits
            float mx = -INFINITY;
            for (int64_t t = 0; t < n_feat; ++t) mx = std::fmax(mx, pfus[t + l*n_feat]);
            float sum = 0.0f;
            std::vector<float> w(n_feat);
            for (int64_t t = 0; t < n_feat; ++t) { w[t] = std::exp(pfus[t + l*n_feat] - mx); sum += w[t]; }
            for (int64_t t = 0; t < n_feat; ++t) w[t] /= sum;

            // residual mix of the raw features, then base + residual
            std::vector<float> v(n_embd);
            for (int64_t h = 0; h < n_embd; ++h) {
                float acc = 0.0f;
                for (int64_t t = 0; t < n_feat; ++t) {
                    acc += w[t] * pinp[(h + t*n_embd) + n*n_feat*n_embd];
                }
                v[h] = b[h] + acc;
            }

            // RMS norm over n_embd, scaled by context_norm
            float ss = 0.0f;
            for (int64_t h = 0; h < n_embd; ++h) ss += v[h]*v[h];
            const float scale = 1.0f/std::sqrt(ss/(float) n_embd + eps);
            for (int64_t h = 0; h < n_embd; ++h) {
                ref[(size_t) (h + l*n_embd) + (size_t) n*n_embd*n_layer] = v[h]*scale*pcn[h];
            }
        }
    }

    // ---- compare ----
    const float * got = (const float *) cur->data;
    double max_err = 0.0;
    for (size_t i = 0; i < ref.size(); ++i) {
        max_err = std::fmax(max_err, std::fabs((double) got[i] - (double) ref[i]));
    }
    printf("fused context: max abs err = %.3e\n", max_err);

    // layer-major layout: the decoder's strided view must equal reference layer il_probe
    const float * pslice = (const float *) slice->data;
    double max_err_slice = 0.0;
    for (int64_t n = 0; n < n_tokens; ++n) {
        for (int64_t h = 0; h < n_embd; ++h) {
            const double r = ref[(size_t) (h + il_probe*n_embd) + (size_t) n*n_embd*n_layer];
            max_err_slice = std::fmax(max_err_slice, std::fabs((double) pslice[h + n*n_embd] - r));
        }
    }
    printf("layer-%lld slice: max abs err = %.3e\n", (long long) il_probe, max_err_slice);

    // a transposed fusion axis would still land within this bound only by coincidence;
    // guard against it explicitly by requiring the per-layer contexts to actually differ
    double min_sep = INFINITY;
    for (int64_t n = 0; n < n_tokens; ++n) {
        for (int64_t l = 1; l < n_layer; ++l) {
            double d = 0.0;
            for (int64_t h = 0; h < n_embd; ++h) {
                const double a = got[(size_t) (h)            + (size_t) n*n_embd*n_layer];
                const double c = got[(size_t) (h + l*n_embd) + (size_t) n*n_embd*n_layer];
                d += (a-c)*(a-c);
            }
            min_sep = std::fmin(min_sep, std::sqrt(d));
        }
    }
    printf("min per-layer context separation = %.3e\n", min_sep);

    ggml_free(ctx);

    const bool ok = max_err < 1e-4 && max_err_slice < 1e-4 && min_sep > 1e-3;
    printf("%s\n", ok ? "PASS" : "FAIL");
    return ok ? 0 : 1;
}