/* vindex_bench.c — standalone proof harness for the engram HNSW ANN index.
 *
 * Measures brute-force cosine top-k (the correctness ORACLE) vs vindex_search
 * (HNSW) on: (a) the REAL paged store harvested read-only, and (b) synthetic
 * clustered data at several sizes to trace the scaling curve. Reports build time,
 * per-query latency (brute vs HNSW), and recall@k (HNSW top-k vs brute top-k).
 *
 * Read-only: never opens a socket, never writes the store. Safe on an nsbx clone.
 *
 * Build:  cc -O2 -std=c11 vindex_bench.c engram_vindex.c -lm -o vindex_bench
 * Usage:  vindex_bench store <neuron.egm> <dim> [nqueries] [k] [ef_csv]
 *         vindex_bench synth <N> [dim] [clusters] [nqueries] [k] [ef_csv]
 */
#include "engram_vindex.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <stdint.h>
#include <time.h>

/* ── deterministic PRNG (splitmix64) so runs are reproducible ─────────────── */
static uint64_t g_seed = 0xD1B54A32D192ED03ULL;
static uint64_t sm(void){
    uint64_t z = (g_seed += 0x9E3779B97F4A7C15ULL);
    z = (z ^ (z >> 30)) * 0xBF58476D1CE4E5B9ULL;
    z = (z ^ (z >> 27)) * 0x94D049BB133111EBULL;
    return z ^ (z >> 31);
}
static double urand(void){ return (double)((sm() >> 11) + 1) * (1.0/9007199254740993.0); }
static double grand(void){ /* Box-Muller */
    double u1 = urand(), u2 = urand();
    return sqrt(-2.0*log(u1)) * cos(2.0*M_PI*u2);
}

static double now_s(void){
    struct timespec ts; clock_gettime(CLOCK_MONOTONIC, &ts);
    return (double)ts.tv_sec + (double)ts.tv_nsec*1e-9;
}

/* L2-normalise a row in place. */
static void l2norm(float* v, int dim){
    double ss = 0; for (int i=0;i<dim;i++) ss += (double)v[i]*v[i];
    if (ss > 0){ float inv = (float)(1.0/sqrt(ss)); for (int i=0;i<dim;i++) v[i]*=inv; }
}

/* Brute-force top-k by cosine distance (1 - dot on normalised vecs).
 * data is n*dim, already L2-normalised. Writes k node ids (row indices) into
 * out_ids ascending by distance. Returns nothing; assumes k<=n. */
static void brute_topk(const float* data, int n, int dim, const float* q,
                       int k, int* out_ids, float* out_d){
    /* maintain a small sorted array of the k best (ascending distance). */
    for (int i=0;i<k;i++){ out_ids[i]=-1; out_d[i]=2.0f+1.0f; }
    for (int i=0;i<n;i++){
        const float* r = data + (size_t)i*dim;
        float s0=0,s1=0,s2=0,s3=0; int j=0;
        for (; j+4<=dim; j+=4){ s0+=q[j]*r[j]; s1+=q[j+1]*r[j+1]; s2+=q[j+2]*r[j+2]; s3+=q[j+3]*r[j+3]; }
        float dot=(s0+s1)+(s2+s3); for (; j<dim; j++) dot+=q[j]*r[j];
        float d = 1.0f - dot;
        if (d >= out_d[k-1]) continue;
        int p = k-1;
        while (p>0 && out_d[p-1] > d){ out_d[p]=out_d[p-1]; out_ids[p]=out_ids[p-1]; p--; }
        out_d[p]=d; out_ids[p]=i;
    }
}

/* recall@k: |brute_topk ∩ hnsw_topk| / k. Both are id arrays of length k. */
static double recall_at_k(const int* gt, const uint64_t* ann, int nann, int k){
    int hit = 0;
    for (int i=0;i<k;i++){
        if (gt[i] < 0) continue;
        for (int j=0;j<nann;j++){ if ((int)ann[j] == gt[i]){ hit++; break; } }
    }
    return (double)hit / (double)k;
}

/* Parse "64,128,256" into an int array; returns count. */
static int parse_csv(const char* s, int* out, int maxo){
    int n=0; if(!s||!*s) return 0;
    const char* p=s;
    while(*p && n<maxo){ out[n++]=atoi(p); while(*p && *p!=',') p++; if(*p==',') p++; }
    return n;
}

/* Generate n unit vectors on a LOW-DIMENSIONAL MANIFOLD, the property that makes
 * real text embeddings tractable for ANN: each vector is a fixed random linear map
 * A (dim × LATENT) applied to a latent gaussian z ∈ R^LATENT, plus small ambient
 * noise, then L2-normalised. Points therefore lie near a `latent`-dim subspace, so
 * every point has a well-defined tight neighbourhood (high recall) and the HNSW
 * graph is cheap to build — unlike near-isotropic 768-d gaussians, where the curse
 * of dimensionality makes all points near-equidistant (no structure → slow build,
 * low recall) and unlike tight clusters (near-duplicates → artificial top-k ties).
 * `sigma` is the ambient-noise scale. This reproduces the intrinsic-dimensionality
 * regime of nomic embeddings, so the scaling curve reflects real-corpus behaviour. */
#define SYNTH_LATENT 48
static void gen_synth(float* data, int n, int dim, int clusters, double sigma){
    (void)clusters;
    float* A = malloc((size_t)dim*SYNTH_LATENT*sizeof(float));   /* fixed random basis */
    for (size_t i=0;i<(size_t)dim*SYNTH_LATENT;i++) A[i]=(float)grand();
    float z[SYNTH_LATENT];
    for (int i=0;i<n;i++){
        for (int l=0;l<SYNTH_LATENT;l++) z[l]=(float)grand();
        float* v = data+(size_t)i*dim;
        for (int j=0;j<dim;j++){
            float acc = (float)(sigma*grand());
            const float* row = A + (size_t)j*SYNTH_LATENT;
            for (int l=0;l<SYNTH_LATENT;l++) acc += row[l]*z[l];
            v[j]=acc;
        }
        l2norm(v, dim);
    }
    free(A);
}

/* Build M / ef_construction come from env (VIDX_M / VIDX_EFC) so the scaling
 * sweep can trade build cost against graph quality without a recompile. 0 = default. */
static int env_int(const char* k, int dflt){ const char* s=getenv(k); return (s&&*s)?atoi(s):dflt; }

/* Run the full brute-vs-HNSW comparison over an already-normalised dataset. */
static void run_bench(const char* label, float* data, int n, int dim,
                      int nq, int k, int* efs, int nef, double build_s){
    (void)build_s;
    int bM = env_int("VIDX_M", 0), bEFC = env_int("VIDX_EFC", 0);
    printf("\n=== %s : N=%d  dim=%d  k=%d  queries=%d ===\n", label, n, dim, k, nq);

    /* build the index once (shared across ef settings). */
    double t0 = now_s();
    VIndex* ix = vindex_create(dim, bM, bEFC);
    for (int i=0;i<n;i++) vindex_insert(ix, (uint64_t)i, data + (size_t)i*dim);
    double bt = now_s()-t0;
    printf("HNSW build: M=%d ef_construction=%d  ->  %.3f s  (%.1f k nodes/s)\n",
           bM?bM:VINDEX_DEFAULT_M, bEFC?bEFC:VINDEX_DEFAULT_EF_CONSTRUCTION, bt, n/1000.0/bt);

    /* choose query vectors: perturb random dataset rows (near-but-not-identical). */
    int* qidx = malloc((size_t)nq*sizeof(int));
    float* qv = malloc((size_t)nq*dim*sizeof(float));
    for (int i=0;i<nq;i++){
        int r = (int)(sm() % (uint64_t)n);
        qidx[i]=r;
        float* dst = qv+(size_t)i*dim; const float* src = data+(size_t)r*dim;
        for (int j=0;j<dim;j++) dst[j] = src[j] + (float)(0.01*grand());
        l2norm(dst, dim);
    }

    /* ground truth: brute-force top-k for every query (also the oracle latency). */
    int* gt = malloc((size_t)nq*k*sizeof(int));
    float* gd = malloc((size_t)k*sizeof(float));
    double tb0 = now_s();
    for (int i=0;i<nq;i++) brute_topk(data, n, dim, qv+(size_t)i*dim, k, gt+(size_t)i*k, gd);
    double brute_ms = (now_s()-tb0)*1000.0/nq;
    printf("BRUTE-FORCE  : %8.3f ms/query   (oracle; O(N*D))\n", brute_ms);

    /* HNSW at each ef. */
    uint64_t* aid = malloc((size_t)k*sizeof(uint64_t));
    float* ad = malloc((size_t)k*sizeof(float));
    printf("%-6s %14s %12s %10s\n", "ef", "HNSW ms/query", "speedup", "recall@k");
    for (int e=0;e<nef;e++){
        int ef = efs[e];
        double th0 = now_s();
        double rec_sum = 0;
        for (int i=0;i<nq;i++){
            int m = vindex_search(ix, qv+(size_t)i*dim, k, ef, aid, ad);
            rec_sum += recall_at_k(gt+(size_t)i*k, aid, m, k);
        }
        double hnsw_ms = (now_s()-th0)*1000.0/nq;
        printf("%-6d %14.4f %11.1fx %10.4f\n", ef, hnsw_ms, brute_ms/hnsw_ms, rec_sum/nq);
    }

    free(qidx); free(qv); free(gt); free(gd); free(aid); free(ad);
    vindex_free(ix);
}

int main(int argc, char** argv){
    setvbuf(stdout, NULL, _IOLBF, 0);   /* line-buffered so progress streams to a log */
    if (argc < 2){ fprintf(stderr,"usage: %s store <path> <dim> [nq] [k] [ef_csv] | synth <N> [dim] [clusters] [nq] [k] [ef_csv] | sweep <dim> <N_csv> [nq] [k] [ef_csv]\n", argv[0]); return 2; }
    int defef[8]; int ndef;

    if (strcmp(argv[1],"sweep")==0){
        if (argc < 4){ fprintf(stderr,"sweep needs <dim> <N_csv>\n"); return 2; }
        int dim = atoi(argv[2]);
        int Ns[16]; int nN = parse_csv(argv[3], Ns, 16);
        int nq = (argc>4)?atoi(argv[4]):200;
        int k  = (argc>5)?atoi(argv[5]):10;
        ndef = (argc>6)?parse_csv(argv[6],defef,8):parse_csv("64,128,200",defef,8);
        for (int s=0;s<nN;s++){
            int N = Ns[s];
            float* data = malloc((size_t)N*dim*sizeof(float));
            if (!data){ fprintf(stderr,"OOM at N=%d\n",N); continue; }
            int clusters = N/100; if (clusters < 64) clusters = 64;
            gen_synth(data, N, dim, clusters, 1.0);
            char lbl[64]; snprintf(lbl,sizeof lbl,"SYNTH N=%d", N);
            run_bench(lbl, data, N, dim, nq, k, defef, ndef, 0.0);
            free(data);
        }
        return 0;
    }

    if (strcmp(argv[1],"store")==0){
        if (argc < 4){ fprintf(stderr,"store needs <path> <dim>\n"); return 2; }
        const char* path = argv[2]; int dim = atoi(argv[3]);
        int nq = (argc>4)?atoi(argv[4]):500;
        int k  = (argc>5)?atoi(argv[5]):10;
        ndef = (argc>6)?parse_csv(argv[6],defef,8):parse_csv("32,64,128,200,400",defef,8);
        printf("Harvesting emb vectors from %s (dim=%d) ...\n", path, dim);
        float* data=NULL; int n=0;
        double t0=now_s();
        int h = vindex_harvest_from_store(path, dim, &data, NULL, &n);
        double harvest_s = now_s()-t0;
        if (h < 0 || n == 0){ fprintf(stderr,"harvest failed (h=%d n=%d) — wrong dim or path?\n", h, n); return 1; }
        printf("Harvested %d live embedded nodes in %.2f s\n", n, harvest_s);
        for (int i=0;i<n;i++) l2norm(data+(size_t)i*dim, dim);   /* oracle needs normalised */
        if (nq > n) nq = n;
        run_bench("REAL STORE", data, n, dim, nq, k, defef, ndef, 0.0);
        free(data);
        return 0;
    }

    if (strcmp(argv[1],"synth")==0){
        if (argc < 3){ fprintf(stderr,"synth needs <N>\n"); return 2; }
        int N = atoi(argv[2]);
        int dim = (argc>3)?atoi(argv[3]):768;
        int clusters = (argc>4)?atoi(argv[4]):200;
        int nq = (argc>5)?atoi(argv[5]):500;
        int k  = (argc>6)?atoi(argv[6]):10;
        ndef = (argc>7)?parse_csv(argv[7],defef,8):parse_csv("64,128,200",defef,8);
        printf("Generating %d synthetic clustered vectors (dim=%d clusters=%d) ...\n", N, dim, clusters);
        float* data = malloc((size_t)N*dim*sizeof(float));
        if (!data){ fprintf(stderr,"OOM allocating %zu bytes\n", (size_t)N*dim*sizeof(float)); return 1; }
        gen_synth(data, N, dim, clusters, 0.35);
        char lbl[64]; snprintf(lbl,sizeof lbl,"SYNTH");
        run_bench(lbl, data, N, dim, nq, k, defef, ndef, 0.0);
        free(data);
        return 0;
    }

    fprintf(stderr,"unknown mode '%s'\n", argv[1]);
    return 2;
}