#!/usr/bin/env python3
# InvariantCore reproduction packet — standalone runner.
# Replicates ablation_final.py (2026-08-10) on the pinned 14-text corpus.
#
# Usage:
#   pip install sentence-transformers numpy
#   python3 invariant_repro.py            # run both models on corpus/, compare to expected JSON
#   python3 invariant_repro.py --model all-MiniLM-L6-v2   # single model
#
# Expected output: per-text delta_sigma values that match ablation_final_v1.json
# within float tolerance; cross-model Pearson/Spearman reported at the end.
import os, sys, json, random, hashlib
import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
CORPUS = os.path.join(HERE, "corpus")
EXPECTED = os.path.join(HERE, "expected_delta_sigma.json")
MAX_SENTS = 60
SHUFFLE_N = 3
MODELS = {"MiniLM": "all-MiniLM-L6-v2", "MPNet": "all-mpnet-base-v2"}


def load_raw(path):
    """Identical to ablation_final.load_raw: first 60 non-empty lines, >=3 words, sorted filenames."""
    sents = []
    for fname in sorted(os.listdir(path)):
        if not fname.endswith(".txt"):
            continue
        try:
            with open(os.path.join(path, fname), encoding="utf-8", errors="replace") as f:
                for line in f:
                    line = line.strip()
                    if line and len(line.split()) >= 3:
                        sents.append(line)
                        if len(sents) >= MAX_SENTS:
                            return sents
        except Exception:
            continue
    return sents


def mean_cos_dist(embeddings):
    if len(embeddings) < 2:
        return None
    dists = []
    for i in range(len(embeddings) - 1):
        a, b = embeddings[i], embeddings[i + 1]
        d = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-9)
        dists.append(float(1 - np.clip(d, -1, 1)))
    return float(np.mean(dists))


def compute_delta_sigma(raw_sents, model):
    if len(raw_sents) < 2:
        return None
    embs_real = model.encode(raw_sents, show_progress_bar=False)
    mu_real = mean_cos_dist(embs_real)
    if mu_real is None:
        return None
    ws_mus = []
    for seed in range(SHUFFLE_N):
        random.seed(seed + 1000)
        shuffled = []
        for s in raw_sents:
            words = s.split()
            random.shuffle(words)
            shuffled.append(" ".join(words))
        embs = model.encode(shuffled, show_progress_bar=False)
        mu = mean_cos_dist(embs)
        if mu is not None:
            ws_mus.append(mu)
    if len(ws_mus) < 2:
        return None
    mu_ws = float(np.mean(ws_mus))
    std_ws = float(np.std(ws_mus))
    if std_ws < 1e-9:
        return None
    return {"mu_real": mu_real, "mu_ws": mu_ws, "std_ws": std_ws,
            "delta_sigma": (mu_real - mu_ws) / std_ws, "n_sents": len(raw_sents)}


def pearson(x, y):
    x, y = np.array(x, float), np.array(y, float)
    if len(x) < 3 or np.std(x) == 0 or np.std(y) == 0:
        return float("nan"), float("nan")
    r = np.corrcoef(x, y)[0, 1]
    n = len(x)
    t = r * np.sqrt((n - 2) / (1 - r * r)) if abs(r) < 1 else float("inf")
    # two-sided p via Student-t survival (no scipy dependency)
    from math import erf, sqrt
    def norm_sf(z): return 0.5 * (1 - erf(z / sqrt(2)))
    p = 2 * norm_sf(abs(t)) if abs(t) < 1e7 else 0.0
    return float(r), float(p)


def spearman(x, y):
    from scipy import stats  # optional; falls back if unavailable
    r, p = stats.spearmanr(x, y)
    return float(r), float(p)


def main():
    random.seed(42)
    np.random.seed(42)
    only = None
    if "--model" in sys.argv:
        only = sys.argv[sys.argv.index("--model") + 1]
    from sentence_transformers import SentenceTransformer
    expected = json.load(open(EXPECTED)) if os.path.exists(EXPECTED) else {}
    results = {}
    for tag, mname in MODELS.items():
        if only and mname != only:
            continue
        print(f"# model {mname}")
        model = SentenceTransformer(mname)
        results[tag] = {}
        for fn in sorted(os.listdir(CORPUS)):
            if not fn.endswith(".txt"):
                continue
            key = fn[:-4]
            sents = load_raw(CORPUS)  # same folder semantics; corpus/ has one file per author
            # but load_raw consumes up to 60 lines across sorted files — for per-author
            # fidelity we read exactly the pinned file:
            with open(os.path.join(CORPUS, fn), encoding="utf-8") as f:
                sents = [ln.strip() for ln in f if ln.strip()]
            ds = compute_delta_sigma(sents, model)
            exp = (expected.get("MiniLM", {}).get(key) or expected.get("MPNet", {}).get(key) or {})
            if ds:
                results[tag][key] = ds
                tag_ok = "MiniLM" if mname.startswith("all-MiniLM") else "MPNet"
                e = expected.get(tag_ok, {}).get(key, {})
                diff = ""
                if e and e.get("delta_sigma") is not None:
                    diff = f"  (expected {e['delta_sigma']:.3f}, diff {ds['delta_sigma']-e['delta_sigma']:+.3f})"
                print(f"  {key:10s} delta_sigma={ds['delta_sigma']:9.3f}{diff}")
    # cross-model correlation on delta_sigma
    if "MiniLM" in results and "MPNet" in results:
        keys = [k for k in results["MiniLM"] if k in results["MPNet"]]
        x = [results["MiniLM"][k]["delta_sigma"] for k in keys]
        y = [results["MPNet"][k]["delta_sigma"] for k in keys]
        r, p = pearson(x, y)
        try:
            rho, p_rho = spearman(x, y)
        except Exception:
            rho, p_rho = float("nan"), float("nan")
        print(f"\nCross-model (n={len(keys)}): Pearson r={r:+.3f} p={p:.3f} | Spearman rho={rho:+.3f} p={p_rho:.3f}")
        print("Reported in ABLATION_REPORT_V1: r=+0.561 p=0.037, rho=+0.332 p=0.246 (n.s.)")
        print("Verdict if reproduced: shuffle-invariance hypothesis REJECTED (non-significant Spearman).")


if __name__ == "__main__":
    main()