/* bench_discrimination.c — M9 REFINEMENT bench: measures whether mean-centering
 * the anisotropic nomic-embed-text space sharpens the §5 geometry operators on
 * REAL data. Read-only over a COPY of the live store (never the live file).
 *
 *   usage:  bench_discrimination [store.egm]
 *           (or set ENGRAM_BENCH_STORE). If no store is given/openable it prints
 *           SKIP and exits 0 — so it is safe in CI without live data.
 *
 * It picks two semantically distinct cohorts by keyword (domain A vs domain B),
 * computes the global mean over the embed-eligible set (via engram_geo_mean_build
 * — the same offset the descriptor uses), then reports BEFORE (raw unit space)
 * vs AFTER (mean-centered space):
 *   - cross-centroid cosine        (lower = better separated)
 *   - cross-centroid Euclid dist   (translation-invariant: a control)
 *   - intra-cohesion per domain    (member cos to own centroid)
 *   - overlap operator             (cross_cos / sqrt(intraA*intraB): ~1 = domains
 *                                   indistinguishable, ~0 = cleanly separated)
 *   - angular separation ratio z   (centroid angle / summed angular spread)
 *   - mean pairwise cosine sample  (the anisotropy headline; ~0.55 raw -> ~0 ctr)
 *
 * Pure C11; links engram_store.c + engram_geometry.c; -lm.
 */
#include "engram_store.h"
#include "engram_geometry.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <math.h>

#define CAP_DOMAIN 400
#define CAP_SAMPLE 800

typedef struct { float** v; int n, cap, dim; } VecSet;
static void vs_init(VecSet* s){ s->v=NULL; s->n=0; s->cap=0; s->dim=0; }
static void vs_push(VecSet* s, const float* e, int dim, int cap){
    if(s->n>=cap) return;
    if(s->dim==0) s->dim=dim;
    if(s->n==s->cap){ int nc=s->cap?s->cap*2:64; s->v=realloc(s->v,(size_t)nc*sizeof*s->v); s->cap=nc; }
    float* c=malloc((size_t)dim*sizeof(float));
    double nn=0; for(int d=0;d<dim;d++) nn+=(double)e[d]*e[d]; nn=sqrt(nn);
    if(nn<1e-12){ free(c); return; }
    for(int d=0;d<dim;d++) c[d]=(float)(e[d]/nn);   /* L2-normalized copy */
    s->v[s->n++]=c;
}
static void vs_free(VecSet* s){ for(int i=0;i<s->n;i++) free(s->v[i]); free(s->v); }

typedef struct { VecSet A, B, S; long idx; } Coh;
static int has(const char* h, const char* n){ return h && strcasestr(h,n)!=NULL; }
static void cb(const StoreNode* n, void* ctx){
    Coh* c=ctx;
    if(!(n->emb && n->emb_dim>0)) return;
    /* every 5th embedded node -> isotropy sample */
    if((c->idx++ % 5)==0) vs_push(&c->S, n->emb, n->emb_dim, CAP_SAMPLE);
    const char* t=n->content; const char* g=n->tags;
    int A = has(t,"quantiz")||has(g,"quantiz")||has(t,"lorablation")||has(t,"70B")||has(t,"LoRA merge");
    int B = has(t,"kubernetes")||has(t,"terraform")||has(t,"argo")||has(g,"infrastructure")||has(t,"vault")||has(t,"cloudflare");
    if(A && !B) vs_push(&c->A, n->emb, n->emb_dim, CAP_DOMAIN);
    else if(B && !A) vs_push(&c->B, n->emb, n->emb_dim, CAP_DOMAIN);
}

/* mean of a VecSet into out (dim doubles). */
static void mean_of(const VecSet* s, const float* gm, double* out){
    int dim=s->dim; for(int d=0;d<dim;d++) out[d]=0;
    for(int i=0;i<s->n;i++) for(int d=0;d<dim;d++) out[d]+=(double)s->v[i][d]-(gm?gm[d]:0.0);
    if(s->n) for(int d=0;d<dim;d++) out[d]/=s->n;
}
static double dnorm(const double* a, int dim){ double s=0; for(int d=0;d<dim;d++) s+=a[d]*a[d]; return sqrt(s); }
static double dcos(const double* a, const double* b, int dim){
    double na=dnorm(a,dim), nb=dnorm(b,dim); if(na<1e-12||nb<1e-12) return 0;
    double s=0; for(int d=0;d<dim;d++) s+=a[d]*b[d]; double c=s/(na*nb);
    if(c>1)c=1; if(c<-1)c=-1; return c;
}
static double deuclid(const double* a, const double* b, int dim){
    double s=0; for(int d=0;d<dim;d++){ double x=a[d]-b[d]; s+=x*x; } return sqrt(s);
}
/* mean cosine of members (minus gm) to centroid c (already gm-subtracted). */
static double cohesion(const VecSet* s, const float* gm, const double* c){
    int dim=s->dim; double nc=dnorm(c,dim); if(nc<1e-12||s->n==0) return 0;
    double acc=0; for(int i=0;i<s->n;i++){
        double dot=0, nv=0;
        for(int d=0;d<dim;d++){ double v=(double)s->v[i][d]-(gm?gm[d]:0.0); dot+=v*c[d]; nv+=v*v; }
        nv=sqrt(nv); if(nv<1e-12) continue; double cc=dot/(nv*nc);
        if(cc>1)cc=1; if(cc<-1)cc=-1; acc+=cc;
    }
    return acc/s->n;
}
/* mean pairwise cosine over a sample (isotropy metric). */
static double mean_pairwise_cos(const VecSet* s, const float* gm){
    int dim=s->dim; if(s->n<2) return 0; double acc=0; long np=0;
    for(int i=0;i<s->n;i++) for(int j=i+1;j<s->n;j++){
        double dot=0, na=0, nb=0;
        for(int d=0;d<dim;d++){ double a=(double)s->v[i][d]-(gm?gm[d]:0.0), b=(double)s->v[j][d]-(gm?gm[d]:0.0);
            dot+=a*b; na+=a*a; nb+=b*b; }
        na=sqrt(na); nb=sqrt(nb); if(na<1e-12||nb<1e-12) continue;
        double c=dot/(na*nb); if(c>1)c=1; if(c<-1)c=-1; acc+=c; np++;
    }
    return np? acc/np : 0;
}

static void report(const char* label, Coh* c, const float* gm){
    int dim=c->A.dim; double* ca=malloc((size_t)dim*sizeof(double)); double* cb=malloc((size_t)dim*sizeof(double));
    mean_of(&c->A, gm, ca); mean_of(&c->B, gm, cb);
    double xcos=dcos(ca,cb,dim), xeuc=deuclid(ca,cb,dim);
    double cohA=cohesion(&c->A,gm,ca), cohB=cohesion(&c->B,gm,cb);
    double overlap = (cohA>0&&cohB>0)? xcos/sqrt(cohA*cohB) : xcos;
    double theta = acos(xcos<-1?-1:(xcos>1?1:xcos));
    double sig = acos(cohA<-1?-1:(cohA>1?1:cohA)) + acos(cohB<-1?-1:(cohB>1?1:cohB));
    double z = (sig>1e-9)? theta/sig : 0;
    double mpc = mean_pairwise_cos(&c->S, gm);
    printf("  [%s]\n", label);
    printf("    cross-centroid cosine   = %+.4f   (lower = better separated)\n", xcos);
    printf("    cross-centroid Euclid   =  %.4f   (translation-invariant control)\n", xeuc);
    printf("    intra-cohesion A / B    =  %.4f / %.4f\n", cohA, cohB);
    printf("    OVERLAP operator        =  %.4f   (~1 = indistinguishable, ~0 = clean)\n", overlap);
    printf("    angular separation z    =  %.3f   (centroid-angle / summed spread; >1 = separated)\n", z);
    printf("    mean pairwise cosine    = %+.4f   (isotropy: ~0.55 anisotropic -> ~0 isotropic)\n", mpc);
    free(ca); free(cb);
}

int main(int argc, char** argv){
    const char* path = (argc>1)? argv[1] : getenv("ENGRAM_BENCH_STORE");
    if(!path){ printf("SKIP: no store path (arg or ENGRAM_BENCH_STORE)\n"); return 0; }
    EngramPagedStore* st=store_open(path);
    if(!st){ printf("SKIP: could not open %s\n", path); return 0; }

    Coh c; vs_init(&c.A); vs_init(&c.B); vs_init(&c.S); c.idx=0;
    store_scan_nodes(st, cb, &c);
    printf("=== two-domain discrimination bench (real store copy) ===\n");
    printf("domain A (quantization) n=%d ; domain B (infrastructure) n=%d ; sample n=%d ; dim=%d\n",
           c.A.n, c.B.n, c.S.n, c.A.dim);
    if(c.A.n<3 || c.B.n<3){ printf("SKIP: a cohort is too small to be meaningful\n");
        vs_free(&c.A); vs_free(&c.B); vs_free(&c.S); store_close(st); return 0; }

    GeoMeanCache* mc=engram_geo_mean_build(st);
    const float* gm=engram_geo_mean_vec(mc);
    printf("global-mean cache: dim=%d over %llu embedded nodes\n\n",
           engram_geo_mean_dim(mc), (unsigned long long)engram_geo_mean_count(mc));

    printf("BEFORE (raw anisotropic unit space):\n");
    report("RAW", &c, NULL);
    printf("\nAFTER (mean-centered isotropic space):\n");
    report("CENTERED", &c, gm);

    engram_geo_mean_free(mc);
    vs_free(&c.A); vs_free(&c.B); vs_free(&c.S);
    store_close(st);
    return 0;
}