/* engram_reason.c — the REASONING layer. Pure compositions over engram_geometry.h.
 * stdlib + libm only; READ-ONLY over its descriptor inputs; touches no store/index. */
#include "engram_reason.h"
#include <stdlib.h>
#include <string.h>
#include <math.h>

/* ── small float-vector helpers ─────────────────────────────────────────────── */
static double vdot(const float* a, const float* b, int dim) {
    double s = 0; for (int i = 0; i < dim; i++) s += (double)a[i] * (double)b[i]; return s;
}
static double vnorm(const float* a, int dim) { return sqrt(vdot(a, a, dim)); }
static double vcos(const float* a, const float* b, int dim) {
    double na = vnorm(a, dim), nb = vnorm(b, dim);
    if (na < 1e-12 || nb < 1e-12) return 0.0;          /* a null vector ⇒ no direction */
    double c = vdot(a, b, dim) / (na * nb);
    if (c > 1.0) c = 1.0; if (c < -1.0) c = -1.0;
    return c;
}
static double l2(const float* a, const float* b, int dim) {
    double s = 0; for (int i = 0; i < dim; i++) { double d = (double)a[i] - (double)b[i]; s += d * d; }
    return sqrt(s);
}

/* ═══════════════════════════════════════════ SHARED — point-to-manifold FIT ══ */
int engram_reason_point_fit(const GeoDescriptor* g, const float* x,
                            double ext_floor, GeoFit* out) {
    if (!g || !x || !out || g->dim <= 0 || !g->centroid) return -1;
    if (!(ext_floor > 0)) ext_floor = 1.0;
    int dim = g->dim;
    /* residual r = x − centroid */
    double rr = 0;                          /* ‖r‖²                                 */
    float* r = malloc((size_t)dim * sizeof(float));
    if (!r) return -1;
    for (int i = 0; i < dim; i++) { double d = (double)x[i] - (double)g->centroid[i]; r[i] = (float)d; rr += d * d; }
    double maha2 = 0, ss_in = 0;            /* Mahalanobis² and in-subspace energy  */
    for (int k = 0; k < g->n_axes; k++) {
        const float* ax = g->axes[k].axis; if (!ax) continue;
        double proj = vdot(r, ax, dim);     /* axes are orthonormal directions       */
        double den = g->axes[k].extent; if (den < ext_floor) den = ext_floor;
        maha2 += (proj / den) * (proj / den);
        ss_in += proj * proj;
    }
    double ortho2 = rr - ss_in; if (ortho2 < 0) ortho2 = 0;   /* off-subspace energy */
    double dist2 = maha2 + ortho2 / (ext_floor * ext_floor);
    out->mahalanobis    = sqrt(maha2);
    out->ortho_residual = sqrt(ortho2);
    out->distance       = sqrt(dist2);
    out->score          = 1.0 / (1.0 + dist2);
    free(r);
    return 0;
}

/* ═══════════════════════════════════════════════════════════════ ANALOGY ════ */
int engram_reason_analogy(const GeoDescriptor* A, const GeoDescriptor* B,
                          const GeoDescriptor* C,
                          const GeoDescriptor* const* candidates, int n_candidates,
                          GeoAnalogyResult* out) {
    if (!A || !B || !C || !out) return -1;
    if (!A->centroid || !B->centroid || !C->centroid) return -1;
    int dim = A->dim;
    if (B->dim != dim || C->dim != dim) return -1;
    memset(out, 0, sizeof *out);
    out->dim = dim; out->best = -1;

    /* Learn R_{A→B}. engram_geo_analogy(X,Y) yields R with apply(R, Y-axis) ≈ X-axis
     * (R maps Y's frame → X's frame); so R that maps A→B is engram_geo_analogy(B,A). */
    GeoAnalogy an;
    if (engram_geo_analogy(B, A, &an) != 0) return -1;
    out->analogy_residual = an.residual;

    /* mapped = R·c_C + (c_B − R·c_A) : the A→B affine (rotation + residual shift). */
    float* RcA = malloc((size_t)dim * sizeof(float));
    float* RcC = malloc((size_t)dim * sizeof(float));
    out->mapped_point = malloc((size_t)dim * sizeof(float));
    if (!RcA || !RcC || !out->mapped_point) { free(RcA); free(RcC); free(out->mapped_point); out->mapped_point = NULL; engram_geo_analogy_free(&an); return -1; }
    engram_geo_analogy_apply(&an, A->centroid, RcA);
    engram_geo_analogy_apply(&an, C->centroid, RcC);
    for (int i = 0; i < dim; i++)
        out->mapped_point[i] = (float)((double)RcC[i] + ((double)B->centroid[i] - (double)RcA[i]));
    free(RcA); free(RcC);
    engram_geo_analogy_free(&an);

    /* nearest candidate to the mapped point (centroid L2). */
    if (candidates && n_candidates > 0) {
        out->n_candidates = n_candidates;
        out->distances = malloc((size_t)n_candidates * sizeof(double));
        if (!out->distances) return -1;
        double best = -1; int bi = -1;
        for (int i = 0; i < n_candidates; i++) {
            const GeoDescriptor* cd = candidates[i];
            double d = (cd && cd->centroid && cd->dim == dim) ? l2(out->mapped_point, cd->centroid, dim) : INFINITY;
            out->distances[i] = d;
            if (bi < 0 || d < best) { best = d; bi = i; }
        }
        out->best = bi; out->best_distance = best;
    }
    return 0;
}
void engram_reason_analogy_free(GeoAnalogyResult* r) {
    if (!r) return;
    free(r->mapped_point); free(r->distances);
    r->mapped_point = NULL; r->distances = NULL;
}

/* ═══════════════════════════════════════════════════════════════ INDUCTION ══ */
int engram_reason_induce(const GeoDescriptor* const* examples, int n_examples,
                         int top_axes, double ext_floor, GeoInduction* out) {
    if (!examples || n_examples < 1 || !out) return -1;
    if (top_axes <= 0) top_axes = 8;
    memset(out, 0, sizeof *out);

    /* fold the examples left→right through the pooled-Gaussian combine. n==1 pools
     * the single example with itself (identical cov ⇒ same shape, id-union = itself). */
    GeoDescriptor* acc = engram_geo_combine(examples[0],
                                            examples[n_examples > 1 ? 1 : 0], top_axes);
    if (!acc) return -1;
    for (int i = 2; i < n_examples; i++) {
        GeoDescriptor* nxt = engram_geo_combine(acc, examples[i], top_axes);
        engram_geo_free(acc);
        if (!nxt) return -1;
        acc = nxt;
    }
    out->rule = acc;
    out->n_examples = n_examples;
    out->ext_floor = (ext_floor > 0) ? ext_floor
                   : (acc->radius > 0 ? acc->radius * 0.25 : 1.0);
    return 0;
}
double engram_reason_membership(const GeoInduction* ind, const float* x) {
    if (!ind || !ind->rule || !x) return -1;
    GeoFit f;
    if (engram_reason_point_fit(ind->rule, x, ind->ext_floor, &f) != 0) return -1;
    return f.score;
}
void engram_reason_induction_free(GeoInduction* out) {
    if (!out) return;
    if (out->rule) engram_geo_free(out->rule);
    out->rule = NULL;
}

/* ═══════════════════════════════════════════════════════════════ ABDUCTION ══ */
int engram_reason_abduce(const float* obs, int dim,
                         const GeoDescriptor* const* hypotheses, int n,
                         double ext_floor, GeoAbduction* out) {
    if (!obs || !hypotheses || n < 1 || dim <= 0 || !out) return -1;
    if (!(ext_floor > 0)) ext_floor = 1.0;
    memset(out, 0, sizeof *out);
    out->n = n; out->best = -1;
    out->scores    = malloc((size_t)n * sizeof(double));
    out->distances = malloc((size_t)n * sizeof(double));
    out->rank      = malloc((size_t)n * sizeof(int));
    if (!out->scores || !out->distances || !out->rank) { engram_reason_abduction_free(out); return -1; }

    double best = -1; int bi = -1;
    for (int i = 0; i < n; i++) {
        out->rank[i] = i;
        const GeoDescriptor* h = hypotheses[i];
        GeoFit f;
        if (!h || h->dim != dim || engram_reason_point_fit(h, obs, ext_floor, &f) != 0) {
            out->scores[i] = 0.0; out->distances[i] = INFINITY;
        } else {
            out->scores[i] = f.score; out->distances[i] = f.distance;
        }
        if (bi < 0 || out->scores[i] > best) { best = out->scores[i]; bi = i; }
    }
    out->best = bi; out->best_score = (bi >= 0) ? out->scores[bi] : 0.0;

    /* rank indices best→worst by score (insertion sort — n is small). */
    for (int i = 1; i < n; i++) {
        int key = out->rank[i]; int j = i - 1;
        while (j >= 0 && out->scores[out->rank[j]] < out->scores[key]) { out->rank[j + 1] = out->rank[j]; j--; }
        out->rank[j + 1] = key;
    }
    return 0;
}
void engram_reason_abduction_free(GeoAbduction* out) {
    if (!out) return;
    free(out->scores); free(out->distances); free(out->rank);
    out->scores = NULL; out->distances = NULL; out->rank = NULL;
}

/* ═══════════════════════════════════════════════════════════════════ CAUSAL ══ */
/* |cos| of two descriptors' centroids after removing confounder Z's subspace. */
static double controlled_assoc(const GeoDescriptor* x, const GeoDescriptor* y,
                               const GeoDescriptor* z) {
    GeoResidual rx, ry; double c = 0;
    int ox = engram_geo_subtract(x, z, 0, &rx);
    int oy = engram_geo_subtract(y, z, 0, &ry);
    if (ox == 0 && oy == 0 && rx.residual_centroid && ry.residual_centroid)
        c = fabs(vcos(rx.residual_centroid, ry.residual_centroid, x->dim));
    if (ox == 0) engram_geo_residual_free(&rx);
    if (oy == 0) engram_geo_residual_free(&ry);
    return c;
}
int engram_reason_causal(const GeoDescriptor* x, const GeoDescriptor* y,
                         const GeoDescriptor* const* confounders, int n_conf,
                         int64_t t_x, int64_t t_y,
                         double drop_frac, GeoCausal* out) {
    if (!x || !y || !out || !x->centroid || !y->centroid || x->dim != y->dim) return -1;
    if (!(drop_frac > 0 && drop_frac < 1)) drop_frac = 0.5;
    memset(out, 0, sizeof *out);
    const double assoc_floor = 0.2;   /* below this = no meaningful association       */

    out->assoc_raw = fabs(vcos(x->centroid, y->centroid, x->dim));
    /* control for each confounder; the strongest single explainer wins (min assoc). */
    double ctrl = out->assoc_raw;
    for (int i = 0; i < n_conf; i++) {
        if (!confounders[i]) continue;
        double c = controlled_assoc(x, y, confounders[i]);
        if (c < ctrl) ctrl = c;
    }
    out->assoc_controlled = ctrl;
    out->temporal_dir = (t_x < t_y) ? 1 : (t_x > t_y) ? -1 : 0;

    if (out->assoc_raw < assoc_floor) {
        out->verdict = GEO_CAUSAL_NONE;
    } else if (ctrl < (1.0 - drop_frac) * out->assoc_raw && ctrl < assoc_floor) {
        out->verdict = GEO_CAUSAL_CONFOUNDED; out->confounded = 1;
    } else if (out->temporal_dir != 0) {
        out->verdict = GEO_CAUSAL_DIRECTED; out->strength = ctrl;
    } else {
        out->verdict = GEO_CAUSAL_NONE;      /* associated + robust but unorientable  */
    }
    return 0;
}

/* ═══════════════════════════════════════════════════════════════════ PLANNING ══ */
int engram_reason_plan(const GeoDescriptor* const* nodes, int n,
                       int start, int goal, double neighbor_radius,
                       int use_wasserstein, GeoPlan* out) {
    if (!nodes || n < 1 || !out) return -1;
    if (start < 0 || start >= n || goal < 0 || goal >= n) return -1;
    if (!(neighbor_radius > 0)) return -1;
    memset(out, 0, sizeof *out);

    /* dense edge weights (i<j symmetric); INFINITY = not adjacent. */
    double* W = malloc((size_t)n * (size_t)n * sizeof(double));
    if (!W) return -1;
    for (int i = 0; i < n; i++) for (int j = 0; j < n; j++) W[(size_t)i * n + j] = (i == j) ? 0.0 : INFINITY;
    for (int i = 0; i < n; i++) {
        for (int j = i + 1; j < n; j++) {
            GeoDistance d;
            if (nodes[i] && nodes[j] && engram_geo_distance(nodes[i], nodes[j], &d) == 0) {
                double w = use_wasserstein ? d.wasserstein2 : d.centroid_distance;
                if (w <= neighbor_radius) { W[(size_t)i * n + j] = w; W[(size_t)j * n + i] = w; }
            }
        }
    }

    /* O(n²) Dijkstra. */
    double* dist = malloc((size_t)n * sizeof(double));
    int*    prev = malloc((size_t)n * sizeof(int));
    char*   done = calloc((size_t)n, 1);
    if (!dist || !prev || !done) { free(W); free(dist); free(prev); free(done); return -1; }
    for (int i = 0; i < n; i++) { dist[i] = INFINITY; prev[i] = -1; }
    dist[start] = 0;
    for (int it = 0; it < n; it++) {
        int u = -1; double bd = INFINITY;
        for (int i = 0; i < n; i++) if (!done[i] && dist[i] < bd) { bd = dist[i]; u = i; }
        if (u < 0) break;
        done[u] = 1;
        if (u == goal) break;
        for (int v = 0; v < n; v++) {
            double w = W[(size_t)u * n + v];
            if (w < INFINITY && !done[v] && dist[u] + w < dist[v]) { dist[v] = dist[u] + w; prev[v] = u; }
        }
    }

    if (dist[goal] < INFINITY) {
        int len = 0; for (int v = goal; v != -1; v = prev[v]) len++;
        out->path = malloc((size_t)len * sizeof(int));
        if (out->path) {
            out->path_len = len;
            int idx = len - 1;
            for (int v = goal; v != -1; v = prev[v]) out->path[idx--] = v;
            out->total_cost = dist[goal];
            out->reached = 1;
        }
    }
    free(W); free(dist); free(prev); free(done);
    return 0;
}
void engram_reason_plan_free(GeoPlan* out) {
    if (!out) return;
    free(out->path); out->path = NULL;
}