/* Closed-form unit tests for the REASONING layer (engram_reason.c). All inputs are
 * hand-built synthetic descriptors whose answers are known in closed form. Every
 * reasoning MODE is proven, not declared. ASan/UBSan target. */
#include "engram_reason.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>

static int failures = 0, checks = 0;
static void ok(const char* what, int cond) {
    checks++;
    if (!cond) { failures++; printf("  FAIL: %s\n", what); }
    else        printf("  ok:   %s\n", what);
}
static void approx(const char* what, double got, double exp, double tol) {
    ok(what, fabs(got - exp) <= tol);
    if (fabs(got - exp) > tol) printf("        got=%.9g exp=%.9g\n", got, exp);
}

/* ── descriptor builders (mirror scratchpad/test_geo_ops.c) ─────────────────── */
static float* vec(const double* v, int dim) {
    float* f = malloc((size_t)dim * sizeof(float));
    for (int i = 0; i < dim; i++) f[i] = (float)v[i];
    return f;
}
static GeoDescriptor* mk(int dim, const double* centroid,
                         int n_axes, const double* axis_flat, const double* extents,
                         int n_members, const char** ids, double total_var) {
    GeoDescriptor* g = calloc(1, sizeof(GeoDescriptor));
    g->dim = dim;
    g->centroid = centroid ? vec(centroid, dim) : NULL;
    g->global_mean = NULL;
    g->n_axes = n_axes;
    g->axes = n_axes ? calloc((size_t)n_axes, sizeof(GeoAxis)) : NULL;
    double tr = 0;
    for (int k = 0; k < n_axes; k++) {
        g->axes[k].axis = vec(&axis_flat[(size_t)k * dim], dim);
        g->axes[k].extent = extents[k];
        tr += extents[k] * extents[k];
    }
    g->total_variance = (total_var >= 0) ? total_var : tr;
    g->radius = sqrt(g->total_variance > 0 ? g->total_variance : 0);
    g->n_members = n_members; g->n_embedded = n_members;
    g->members = n_members ? calloc((size_t)n_members, sizeof(GeoMember)) : NULL;
    for (int i = 0; i < n_members; i++) {
        g->members[i].id = strdup(ids[i]);
        g->members[i].membership = 1.0;
        g->members[i].centrality = (double)(n_members - i);
        g->members[i].embedded = 1;
    }
    g->hub_id = n_members ? strdup(ids[0]) : strdup("");
    g->k_core = 1; g->co_registration = 0.0; g->n_edges = 0; g->edges = NULL;
    return g;
}

int main(void) {
    printf("== REASONING layer unit tests ==\n");

    /* ══════════════════ ANALOGY — recover an affine A→B, apply to C ══════════ */
    /* A→B is a +90° rotation in the e0-e1 plane ((x,y)→(-y,x)) plus a +5 shift in e2.
     * A frame = (e0,e1); B frame = rotated (e1,-e0); cB = R·cA + t. Predict D from C. */
    {
        int dim = 4;
        double cA[4] = {1,0,0,0};
        double cB[4] = {0,1,5,0};                 /* R·(1,0,0,0)=(0,1,0,0) + (0,0,5,0) */
        double cC[4] = {2,0,0,0};
        double axA[8] = {1,0,0,0,  0,1,0,0}; double exA[2] = {1,1};
        double axB[8] = {0,1,0,0, -1,0,0,0}; double exB[2] = {1,1};  /* R·e0, R·e1 */
        double axC[8] = {1,0,0,0,  0,1,0,0}; double exC[2] = {1,1};
        const char* idA[1] = {"A"}, *idB[1] = {"B"}, *idC[1] = {"C"};
        GeoDescriptor* A = mk(dim, cA, 2, axA, exA, 1, idA, -1);
        GeoDescriptor* B = mk(dim, cB, 2, axB, exB, 1, idB, -1);
        GeoDescriptor* C = mk(dim, cC, 2, axC, exC, 1, idC, -1);
        /* candidates: the true D + two distractors. true D = R·cC + t = (0,2,5,0). */
        double d_true[4] = {0,2,5,0}, d_far1[4] = {9,9,9,9}, d_far2[4] = {0,0,0,0};
        const char* idD[1] = {"Dt"}, *idF1[1] = {"F1"}, *idF2[1] = {"F2"};
        GeoDescriptor* Dt = mk(dim, d_true, 0, NULL, NULL, 1, idD, 0.0);
        GeoDescriptor* F1 = mk(dim, d_far1, 0, NULL, NULL, 1, idF1, 0.0);
        GeoDescriptor* F2 = mk(dim, d_far2, 0, NULL, NULL, 1, idF2, 0.0);
        const GeoDescriptor* cand[3] = {F1, Dt, F2};   /* true one at index 1 */
        GeoAnalogyResult res;
        int rc = engram_reason_analogy(A, B, C, cand, 3, &res);
        ok("analogy returns 0", rc == 0);
        printf("[analogy] residual=%.6f mapped=(%.4f,%.4f,%.4f,%.4f) best=%d bd=%.5f\n",
               res.analogy_residual, res.mapped_point[0], res.mapped_point[1],
               res.mapped_point[2], res.mapped_point[3], res.best, res.best_distance);
        approx("procrustes residual ~0", res.analogy_residual, 0.0, 1e-4);
        approx("mapped.x=0", res.mapped_point[0], 0.0, 1e-4);
        approx("mapped.y=2", res.mapped_point[1], 2.0, 1e-4);
        approx("mapped.z(e2)=5", res.mapped_point[2], 5.0, 1e-4);
        ok("nearest candidate = true D (idx 1)", res.best == 1);
        approx("best distance ~0", res.best_distance, 0.0, 1e-3);
        engram_reason_analogy_free(&res);
        engram_geo_free(A); engram_geo_free(B); engram_geo_free(C);
        engram_geo_free(Dt); engram_geo_free(F1); engram_geo_free(F2);
    }

    /* ══════════════════ INDUCTION — recover a shared subspace + membership ═══ */
    /* 3 examples all spread over span(e0,e1) (ext 1 & 0.8), each with a small
     * idiosyncratic axis (e2 or e3, ext 0.2). Centroids all 0. The induced rule's
     * top-2 axes must lie in span(e0,e1); a held-out in-plane point fits, an
     * off-subspace point does not. */
    {
        int dim = 4;
        double c0[4] = {0,0,0,0};
        double axsh[8] = {1,0,0,0,  0,1,0,0}; double exsh[2] = {1.0, 0.8};
        double ax1[12] = {1,0,0,0, 0,1,0,0, 0,0,1,0}; double ex1[3] = {1.0,0.8,0.2}; /* +e2 */
        double ax2[12] = {1,0,0,0, 0,1,0,0, 0,0,0,1}; double ex2[3] = {1.0,0.8,0.2}; /* +e3 */
        const char* i1[2] = {"e1a","e1b"}, *i2[2] = {"e2a","e2b"}, *i3[2] = {"e3a","e3b"};
        GeoDescriptor* E1 = mk(dim, c0, 3, ax1,  ex1,  2, i1, -1);
        GeoDescriptor* E2 = mk(dim, c0, 3, ax2,  ex2,  2, i2, -1);
        GeoDescriptor* E3 = mk(dim, c0, 2, axsh, exsh, 2, i3, -1);
        const GeoDescriptor* ex[3] = {E1, E2, E3};
        GeoInduction ind;
        int rc = engram_reason_induce(ex, 3, 8, 1.0, &ind);
        ok("induce returns 0", rc == 0);
        printf("[induction] rule n_axes=%d ext0=%.4f ext1=%.4f\n",
               ind.rule->n_axes, ind.rule->n_axes > 0 ? ind.rule->axes[0].extent : 0,
               ind.rule->n_axes > 1 ? ind.rule->axes[1].extent : 0);
        /* top-2 axes lie in span(e0,e1): their e2,e3 components ~0. */
        int inplane = 1;
        for (int k = 0; k < 2 && k < ind.rule->n_axes; k++) {
            const float* a = ind.rule->axes[k].axis;
            printf("        axis%d=(%.3f,%.3f,%.3f,%.3f) ext=%.4f\n", k, a[0],a[1],a[2],a[3], ind.rule->axes[k].extent);
            if (fabs(a[2]) > 0.06 || fabs(a[3]) > 0.06) inplane = 0;
        }
        ok("induced top-2 axes lie in shared span(e0,e1)", inplane);
        approx("dominant extent ~1.0", ind.rule->axes[0].extent, 1.0, 0.06);
        approx("second extent ~0.8",   ind.rule->axes[1].extent, 0.8, 0.06);
        /* membership: in-plane near-centroid positive fits; off-subspace negative doesn't. */
        float xpos[4] = {0.3f, -0.2f, 0, 0};
        float xneg[4] = {0, 0, 3.0f, 0};       /* large along e2 — outside the rule */
        float xfar[4] = {5.0f, 0, 0, 0};       /* in-plane but far — Mahalanobis blows up */
        double mp = engram_reason_membership(&ind, xpos);
        double mn = engram_reason_membership(&ind, xneg);
        double mf = engram_reason_membership(&ind, xfar);
        printf("[induction] membership pos=%.4f neg=%.4f far=%.4f\n", mp, mn, mf);
        ok("held-out positive fits (>0.5)", mp > 0.5);
        ok("off-subspace negative rejected (<0.3)", mn < 0.3);
        ok("in-plane-but-far rejected (<0.3)", mf < 0.3);
        ok("positive fits far better than negative", mp > mn + 0.4);
        engram_reason_induction_free(&ind);
        engram_geo_free(E1); engram_geo_free(E2); engram_geo_free(E3);
    }

    /* ══════════════════ ABDUCTION — pick the best-explaining structure ═══════ */
    /* obs planted near H1's centroid among 3 candidate structures. */
    {
        int dim = 4;
        double h0[4] = {0,0,0,0}, h1[4] = {5,0,0,0}, h2[4] = {0,5,0,0};
        double ax[8] = {1,0,0,0, 0,1,0,0}; double ex[2] = {1,1};
        const char* n0[1] = {"H0"}, *n1[1] = {"H1"}, *n2[1] = {"H2"};
        GeoDescriptor* H0 = mk(dim, h0, 2, ax, ex, 1, n0, -1);
        GeoDescriptor* H1 = mk(dim, h1, 2, ax, ex, 1, n1, -1);
        GeoDescriptor* H2 = mk(dim, h2, 2, ax, ex, 1, n2, -1);
        const GeoDescriptor* H[3] = {H0, H1, H2};
        float obs[4] = {5.2f, 0.1f, 0, 0};       /* sits inside H1 */
        GeoAbduction ab;
        int rc = engram_reason_abduce(obs, dim, H, 3, 1.0, &ab);
        ok("abduce returns 0", rc == 0);
        printf("[abduction] best=%d best_score=%.4f rank=[%d,%d,%d] d=[%.3f,%.3f,%.3f]\n",
               ab.best, ab.best_score, ab.rank[0], ab.rank[1], ab.rank[2],
               ab.distances[0], ab.distances[1], ab.distances[2]);
        ok("best explanation = H1", ab.best == 1);
        ok("rank[0] = H1", ab.rank[0] == 1);
        ok("H1 has smallest distance", ab.distances[1] < ab.distances[0] && ab.distances[1] < ab.distances[2]);
        engram_reason_abduction_free(&ab);
        engram_geo_free(H0); engram_geo_free(H1); engram_geo_free(H2);
    }

    /* ══════════════════ CAUSAL — direction + confounder flag ═════════════════ */
    /* Chain A→B→C along e0 (temporal 1<2<3). Confounder Z (e1) injects into A and
     * drives D (t=4). A–D correlate only via Z ⇒ must be flagged CONFOUNDED. */
    {
        int dim = 4;
        double cA[4] = {1,1,0,0};   /* e0 (chain) + e1 (confounder leak) */
        double cB[4] = {1,0,0,0};   /* e0 */
        double cC[4] = {2,0,0,0};   /* e0 */
        double cD[4] = {0,1,0,0};   /* e1 only — driven by Z */
        double cZ[4] = {0,1,0,0};   /* confounder centroid */
        double axZ[4] = {0,1,0,0}; double exZ[1] = {1};   /* Z's subspace = e1 */
        const char* idA[1]={"A"},*idB[1]={"B"},*idC[1]={"C"},*idD[1]={"D"},*idZ[1]={"Z"};
        GeoDescriptor* A = mk(dim, cA, 0, NULL, NULL, 1, idA, 0.0);
        GeoDescriptor* B = mk(dim, cB, 0, NULL, NULL, 1, idB, 0.0);
        GeoDescriptor* C = mk(dim, cC, 0, NULL, NULL, 1, idC, 0.0);
        GeoDescriptor* D = mk(dim, cD, 0, NULL, NULL, 1, idD, 0.0);
        GeoDescriptor* Z = mk(dim, cZ, 1, axZ, exZ, 1, idZ, -1);
        const GeoDescriptor* conf[1] = {Z};

        GeoCausal ab, bc, ad, bd;
        engram_reason_causal(A, B, conf, 1, /*t*/1, 2, 0.5, &ab);
        engram_reason_causal(B, C, conf, 1, 2, 3, 0.5, &bc);
        engram_reason_causal(A, D, conf, 1, 1, 4, 0.5, &ad);
        engram_reason_causal(B, D, conf, 1, 2, 4, 0.5, &bd);
        printf("[causal] A->B: raw=%.3f ctrl=%.3f dir=%d verdict=%d strength=%.3f\n",
               ab.assoc_raw, ab.assoc_controlled, ab.temporal_dir, ab.verdict, ab.strength);
        printf("[causal] B->C: raw=%.3f ctrl=%.3f dir=%d verdict=%d\n", bc.assoc_raw, bc.assoc_controlled, bc.temporal_dir, bc.verdict);
        printf("[causal] A--D: raw=%.3f ctrl=%.3f dir=%d verdict=%d confounded=%d\n",
               ad.assoc_raw, ad.assoc_controlled, ad.temporal_dir, ad.verdict, ad.confounded);
        printf("[causal] B--D: raw=%.3f verdict=%d\n", bd.assoc_raw, bd.verdict);
        ok("A->B DIRECTED", ab.verdict == GEO_CAUSAL_DIRECTED);
        ok("A->B direction A precedes B", ab.temporal_dir == 1);
        ok("A->B association survives control (ctrl high)", ab.assoc_controlled > 0.6);
        ok("B->C DIRECTED", bc.verdict == GEO_CAUSAL_DIRECTED);
        ok("A--D CONFOUNDED (flagged)", ad.verdict == GEO_CAUSAL_CONFOUNDED && ad.confounded == 1);
        ok("A--D raw correlated but control kills it", ad.assoc_raw > 0.6 && ad.assoc_controlled < 0.2);
        ok("B--D NONE (no association at all)", bd.verdict == GEO_CAUSAL_NONE);
        engram_geo_free(A); engram_geo_free(B); engram_geo_free(C); engram_geo_free(D); engram_geo_free(Z);
    }

    /* ══════════════════ PLANNING — geodesic path along a curved manifold ═════ */
    /* 6 neighborhoods on a semicircle (radius 10). Consecutive chord ~6.18,
     * skip-one ~11.76, endpoints ~20. neighbor_radius=7 admits only consecutive
     * hops ⇒ the plan must traverse the whole arc 0→1→2→3→4→5. */
    {
        int dim = 4; int N = 6; double R = 10.0;
        GeoDescriptor* nodes[6];
        char nm[6][8];
        for (int k = 0; k < N; k++) {
            double th = M_PI * (double)k / (double)(N - 1);
            double c[4] = { R * cos(th), R * sin(th), 0, 0 };
            snprintf(nm[k], sizeof nm[k], "n%d", k);
            const char* id[1] = { nm[k] };
            nodes[k] = mk(dim, c, 0, NULL, NULL, 1, id, 0.0);
        }
        const GeoDescriptor* cn[6];
        for (int k = 0; k < N; k++) cn[k] = nodes[k];
        GeoPlan plan;
        int rc = engram_reason_plan(cn, N, 0, 5, 7.0, 0, &plan);
        ok("plan returns 0", rc == 0);
        printf("[planning] reached=%d len=%d cost=%.4f path=[", plan.reached, plan.path_len, plan.total_cost);
        for (int i = 0; i < plan.path_len; i++) printf("%s%d", i ? "," : "", plan.path[i]);
        printf("]\n");
        ok("goal reached", plan.reached == 1);
        ok("path length = 6 (full arc)", plan.path_len == 6);
        int monotone = (plan.path_len == 6);
        for (int i = 0; i < plan.path_len; i++) if (plan.path[i] != i) monotone = 0;
        ok("path = 0,1,2,3,4,5 (the geodesic)", monotone);
        /* arc cost ~ 5 * 6.18 = 30.9, and strictly longer than the 20-unit chord. */
        approx("arc cost ~30.9", plan.total_cost, 30.9, 0.6);
        ok("arc longer than straight chord (20)", plan.total_cost > 20.0);
        engram_reason_plan_free(&plan);

        /* negative control: radius too small to connect anything ⇒ unreachable. */
        GeoPlan p2;
        engram_reason_plan(cn, N, 0, 5, 1.0, 0, &p2);
        ok("unreachable when radius < min edge", p2.reached == 0);
        engram_reason_plan_free(&p2);
        for (int k = 0; k < N; k++) engram_geo_free(nodes[k]);
    }

    printf("\n== %d checks, %d failures ==\n", checks, failures);
    return failures ? 1 : 0;
}