#!/usr/bin/env python3
"""T^GPT v0.8.2 semantic shift and recovery-specific lift.

Primary changes from v0.8.1:
- no hard support threshold in confirmatory endpoint path;
- explicit UNMAPPED/UNRESOLVED states remain in denominator;
- continuous directional deficit + expansion;
- null calibration by independent reference replicate or label permutation;
- recovery is paired RECOVERY vs NEUTRAL_CONTROL from the same B_i;
- no averaging of per-task recovery ratios;
- generalized inference is task-level.
"""
from __future__ import annotations
import argparse, json, math, random, hashlib
from collections import Counter
from pathlib import Path
from typing import Dict, Iterable, List, Tuple

SPECIAL = {"UNMAPPED", "UNRESOLVED"}


def validate_labels(labels: Iterable[str], relevance: Dict[str, int], where: str) -> None:
    allowed = set(relevance) | SPECIAL
    bad = sorted({str(x) for x in labels if str(x) not in allowed})
    if bad:
        raise ValueError(f"{where}: unknown/final-forbidden labels: {bad}")


def probs(labels: List[str], relevance: Dict[str, int]) -> Dict[str, float]:
    validate_labels(labels, relevance, "labels")
    n = len(labels)
    if n == 0:
        raise ValueError("arm must contain at least one finalized generation")
    c = Counter(map(str, labels))
    return {cluster: c.get(cluster, 0) / n for cluster in relevance}


def rates(labels: List[str], relevance: Dict[str, int]) -> dict:
    validate_labels(labels, relevance, "labels")
    n = len(labels)
    if not n:
        raise ValueError("arm must contain at least one finalized generation")
    c = Counter(map(str, labels))
    known = sum(c.get(k, 0) for k in relevance)
    return {
        "N_total": n,
        "known_semantic_mass": known / n,
        "UNMAPPED_rate": c.get("UNMAPPED", 0) / n,
        "UNRESOLVED_rate": c.get("UNRESOLVED", 0) / n,
    }


def relevant_ids(relevance: Dict[str, int]) -> List[str]:
    bad = [k for k, v in relevance.items() if int(v) not in (0, 1)]
    if bad:
        raise ValueError(f"relevance values must be 0/1: {bad}")
    return [k for k, v in relevance.items() if int(v) == 1]


def directional(pa: Dict[str, float], pb: Dict[str, float], rel: List[str]) -> Tuple[float, float, Dict[str, float], Dict[str, float]]:
    d = {c: max(0.0, pa.get(c, 0.0) - pb.get(c, 0.0)) for c in rel}
    e = {c: max(0.0, pb.get(c, 0.0) - pa.get(c, 0.0)) for c in rel}
    return sum(d.values()), sum(e.values()), d, e


def closure(pb: Dict[str, float], pr: Dict[str, float], deficit_by_cluster: Dict[str, float]) -> float:
    total = 0.0
    for c, d in deficit_by_cluster.items():
        inc = max(0.0, pr.get(c, 0.0) - pb.get(c, 0.0))
        total += min(d, inc)
    return total


def perm_null(a: List[str], b: List[str], relevance: Dict[str, int], reps: int, rng: random.Random) -> dict:
    rel = relevant_ids(relevance)
    pool = list(a) + list(b)
    na = len(a)
    ds, es = [], []
    for _ in range(reps):
        x = pool[:]
        rng.shuffle(x)
        aa, bb = x[:na], x[na:]
        d, e, _, _ = directional(probs(aa, relevance), probs(bb, relevance), rel)
        ds.append(d)
        es.append(e)
    return {"deficit_null_mean": sum(ds) / len(ds), "expansion_null_mean": sum(es) / len(es), "_deficit_draws": ds, "_expansion_draws": es}


def paired_swap_null(pairs: List[dict], relevance: Dict[str, int], a: List[str], reps: int, rng: random.Random) -> dict:
    rel = relevant_ids(relevance)
    b = [str(x["base"]) for x in pairs]
    r = [str(x["recovery"]) for x in pairs]
    c = [str(x["control"]) for x in pairs]
    for name, arr in [("paired.base", b), ("paired.recovery", r), ("paired.control", c)]:
        validate_labels(arr, relevance, name)
    pa, pb = probs(a, relevance), probs(b, relevance)
    _, _, d_by, _ = directional(pa, pb, rel)
    obs = closure(pb, probs(r, relevance), d_by) - closure(pb, probs(c, relevance), d_by)
    draws = []
    for _ in range(reps):
        rr, cc = [], []
        for x in pairs:
            if rng.random() < 0.5:
                rr.append(str(x["recovery"])); cc.append(str(x["control"]))
            else:
                rr.append(str(x["control"])); cc.append(str(x["recovery"]))
        draws.append(closure(pb, probs(rr, relevance), d_by) - closure(pb, probs(cc, relevance), d_by))
    p = (1 + sum(abs(x) >= abs(obs) - 1e-15 for x in draws)) / (len(draws) + 1)
    return {"observed": obs, "two_sided_randomization_p": p, "null_mean": sum(draws) / len(draws)}


def one_task(task: dict, perm_reps: int, swap_reps: int, seed: int) -> dict:
    tid = str(task["task_id"])
    relevance = {str(k): int(v) for k, v in task["relevance"].items()}
    rel = relevant_ids(relevance)
    a = [str(x) for x in task["reference_A"]]
    validate_labels(a, relevance, f"{tid}.reference_A")
    pairs = task.get("paired_probes")
    if pairs is not None:
        ids = [str(x["pair_id"]) for x in pairs]
        if len(ids) != len(set(ids)):
            raise ValueError(f"{tid}: duplicate pair_id")
        b = [str(x["base"]) for x in pairs]
        r = [str(x["recovery"]) for x in pairs]
        c = [str(x["control"]) for x in pairs]
    else:
        b = [str(x) for x in task["contracted_B"]]
        r, c = [], []
    validate_labels(b, relevance, f"{tid}.contracted_B")
    pa, pb = probs(a, relevance), probs(b, relevance)
    d, e, d_by, e_by = directional(pa, pb, rel)
    stable = int(hashlib.sha256(tid.encode("utf-8")).hexdigest()[:8], 16)
    rng = random.Random(seed ^ stable)
    null_mode = "permutation"
    if task.get("reference_A_prime") is not None:
        ap = [str(x) for x in task["reference_A_prime"]]
        validate_labels(ap, relevance, f"{tid}.reference_A_prime")
        pap = probs(ap, relevance)
        d1, e1, _, _ = directional(pa, pap, rel)
        d2, e2, _, _ = directional(pap, pa, rel)
        null_floor_d, null_floor_e = (d1 + d2) / 2, (e1 + e2) / 2
        null_mode = "reference_replicate"
        null_info = {"A_prime_rates": rates(ap, relevance)}
    else:
        pn = perm_null(a, b, relevance, perm_reps, rng)
        null_floor_d = float(pn["deficit_null_mean"])
        null_floor_e = float(pn["expansion_null_mean"])
        dd = pn.pop("_deficit_draws"); ee = pn.pop("_expansion_draws")
        pn["deficit_p_ge"] = (1 + sum(x >= d - 1e-15 for x in dd)) / (perm_reps + 1)
        pn["expansion_p_ge"] = (1 + sum(x >= e - 1e-15 for x in ee)) / (perm_reps + 1)
        null_info = pn
    out = {
        "task_id": tid,
        "mapping": {"A": rates(a, relevance), "B": rates(b, relevance)},
        "shift": {
            "DEFICIT_RAW": d, "EXPANSION_RAW": e,
            "DEFICIT_EXCESS": d - null_floor_d, "EXPANSION_EXCESS": e - null_floor_e,
            "null_mode": null_mode, "null_floor_deficit": null_floor_d, "null_floor_expansion": null_floor_e,
            "deficit_by_cluster": d_by, "expansion_by_cluster": e_by, "null": null_info,
        },
    }
    if pairs is not None:
        for name, arr in [("R", r), ("CONTROL", c)]:
            validate_labels(arr, relevance, f"{tid}.{name}")
        cr = closure(pb, probs(r, relevance), d_by)
        cc = closure(pb, probs(c, relevance), d_by)
        out["mapping"]["R"] = rates(r, relevance); out["mapping"]["CONTROL"] = rates(c, relevance)
        out["recovery"] = {
            "CLOSURE_RECOVERY": cr, "CLOSURE_CONTROL": cc, "RECOVERY_LIFT": cr - cc,
            "paired_randomization": paired_swap_null(pairs, relevance, a, swap_reps, rng), "N_pairs": len(pairs)
        }
    return out


def percentile(xs: List[float], q: float):
    if not xs:
        return None
    ys = sorted(xs); pos = (len(ys) - 1) * q; lo = math.floor(pos); hi = math.ceil(pos)
    return ys[lo] if lo == hi else ys[lo] + (ys[hi] - ys[lo]) * (pos - lo)


def task_bootstrap(values: List[float], reps: int, seed: int):
    if len(values) < 2:
        return {"n_tasks": len(values), "mean": values[0] if values else None, "ci95": None}
    rng = random.Random(seed); draws = []
    for _ in range(reps):
        s = [values[rng.randrange(len(values))] for _ in values]
        draws.append(sum(s) / len(s))
    return {"n_tasks": len(values), "mean": sum(values) / len(values), "ci95": [percentile(draws, .025), percentile(draws, .975)]}


def analyze(data: dict) -> dict:
    perm_reps = int(data.get("permutation_reps", 5000)); swap_reps = int(data.get("paired_swap_reps", 5000))
    task_boot_reps = int(data.get("task_bootstrap_reps", 10000)); seed = int(data.get("seed", 8082))
    tasks = [one_task(t, perm_reps, swap_reps, seed + i * 7919) for i, t in enumerate(data["tasks"])]
    lifts = [t["recovery"]["RECOVERY_LIFT"] for t in tasks if "recovery" in t]
    deficits = [t["shift"]["DEFICIT_EXCESS"] for t in tasks]
    minimum_tasks = int(data.get("minimum_tasks_for_generalization", 30))
    return {
        "schema": "T_GPT_RECOVERY_RESULT_V0_8_2",
        "parameters": {"permutation_reps": perm_reps, "paired_swap_reps": swap_reps, "task_bootstrap_reps": task_boot_reps, "seed": seed},
        "tasks": tasks,
        "across_tasks": {
            "DEFICIT_EXCESS": task_bootstrap(deficits, task_boot_reps, seed + 1),
            "RECOVERY_LIFT": task_bootstrap(lifts, task_boot_reps, seed + 2) if lifts else None,
            "generalization_gate": {"minimum_tasks": minimum_tasks, "n_tasks": len(tasks), "open": len(tasks) >= minimum_tasks}
        },
        "guardrails": [
            "No hard support threshold is used in primary shift/recovery metrics.",
            "Task bootstrap concerns generalization over tasks, not within-task bootstrap validity of positive-part functionals.",
            "UNMAPPED and UNRESOLVED remain in denominators and are reported separately.",
            "RECOVERY_LIFT is primary for E2; REC_GAIN and IRREV_OBS are retired from the confirmatory path."
        ]
    }


def main():
    ap = argparse.ArgumentParser(); ap.add_argument("input", type=Path); ap.add_argument("-o", "--output", type=Path)
    args = ap.parse_args(); data = json.loads(args.input.read_text(encoding="utf-8")); result = analyze(data)
    payload = json.dumps(result, ensure_ascii=False, indent=2) + "\n"
    args.output.write_text(payload, encoding="utf-8") if args.output else print(payload, end="")


if __name__ == "__main__":
    main()