import json
import os
from pathlib import Path

ROOT = Path(".")
INPUT_JSON = ROOT / "results" / "row20_h0_attribution" / "row20_thin_slice_attribution.json"
OUTPUT_JSON = ROOT / "results" / "row20_h0_attribution" / "row20_SN_offset_probe.json"
OUTPUT_REPORT = ROOT / "results" / "row20_h0_attribution" / "row20_SN_offset_probe_REPORT.md"


def load_json(path):
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def find_block(data, tag):
    for block in data.get("blocks", []):
        if block.get("tag") == tag:
            return block
    raise ValueError(f"Block not found: {tag}")


def get_z(node):
    return float(node["slice"].split("z=")[1])


def mean(values):
    return sum(values) / len(values)


def main():
    data = load_json(INPUT_JSON)
    sn = find_block(data, "SN")
    nodes = sorted(sn["nodes"], key=get_z)

    raw = []
    for n in nodes:
        raw.append({
            "slice": n["slice"],
            "z": get_z(n),
            "data": n["data"],
            "theory": n["theory"],
            "residual": n["residual_after_gate_handling"],
            "sigma": n["sigma_diag"],
            "pull": n["diag_pull_sigma"],
            "chi2": n["white_chi2_contribution"],
        })

    residuals = [n["residual"] for n in raw]

    # Recenter by subtracting the mean residual.
    # This tests whether the pivot survives after removing a global offset.
    global_offset = mean(residuals)

    recentered = []
    for n in raw:
        r = n["residual"] - global_offset
        pull = r / n["sigma"]
        recentered.append({
            **n,
            "recentered_residual": r,
            "recentered_pull": pull,
            "recentered_chi2_diag": pull ** 2,
        })

    # Recenter by endpoints only.
    # This tests whether z=0.35 is anomalous relative to the low/high-z anchor line.
    low = raw[0]
    mid = raw[1]
    high = raw[2]

    endpoint_offset = mean([low["residual"], high["residual"]])

    endpoint_recentered = []
    for n in raw:
        r = n["residual"] - endpoint_offset
        pull = r / n["sigma"]
        endpoint_recentered.append({
            **n,
            "endpoint_recentered_residual": r,
            "endpoint_recentered_pull": pull,
            "endpoint_recentered_chi2_diag": pull ** 2,
        })

    # Linear endpoint trend test after offset handling
    z0, r0 = low["z"], low["residual"]
    z1, r1 = mid["z"], mid["residual"]
    z2, r2 = high["z"], high["residual"]

    interp_mid = r0 + (r2 - r0) * ((z1 - z0) / (z2 - z0))
    pivot_residual = r1 - interp_mid
    pivot_pull = pivot_residual / mid["sigma"]

    # Recentered pivot is unchanged by global additive offset, but report explicitly.
    r0g = low["residual"] - global_offset
    r1g = mid["residual"] - global_offset
    r2g = high["residual"] - global_offset
    interp_mid_g = r0g + (r2g - r0g) * ((z1 - z0) / (z2 - z0))
    pivot_residual_global_recentered = r1g - interp_mid_g
    pivot_pull_global_recentered = pivot_residual_global_recentered / mid["sigma"]

    interpretation = []

    interpretation.append(
        "This probe removes global SN offset structure and tests whether the z=0.35 pivot survives."
    )

    if abs(pivot_pull_global_recentered) > 2:
        interpretation.append(
            "The z=0.35 pivot survives global recentering; it is not explained by the existing nuisance offset alone."
        )
    else:
        interpretation.append(
            "The z=0.35 pivot weakens after global recentering; nuisance offset may be a major contributor."
        )

    if abs(pivot_pull_global_recentered - pivot_pull) < 1e-9:
        interpretation.append(
            "The pivot curvature is invariant under a global additive SN offset, as expected."
        )

    result = {
        "signature": "ROW20_SN_OFFSET_PROBE_v1_NATIVE_DIAGNOSTIC",
        "context": "DIAGNOSTIC_NON_GATING",
        "rules": data.get("rules", {}),
        "existing_nuisance": sn.get("nuisance", {}),
        "raw_nodes": raw,
        "global_offset_removed": {
            "global_offset": global_offset,
            "nodes": recentered,
        },
        "endpoint_offset_removed": {
            "endpoint_offset": endpoint_offset,
            "nodes": endpoint_recentered,
        },
        "pivot_geometry": {
            "midpoint_slice": mid["slice"],
            "raw_pivot_residual": pivot_residual,
            "raw_pivot_pull": pivot_pull,
            "global_recentered_pivot_residual": pivot_residual_global_recentered,
            "global_recentered_pivot_pull": pivot_pull_global_recentered,
        },
        "interpretation": interpretation,
        "verdict": (
            "SN_pivot_survival_under_offset_probe_complete; diagnostic only; "
            "no physics, thresholds, datasets, parameters, or gate verdicts changed."
        ),
    }

    OUTPUT_JSON.parent.mkdir(parents=True, exist_ok=True)

    with open(OUTPUT_JSON, "w", encoding="utf-8") as f:
        json.dump(result, f, indent=2)

    with open(OUTPUT_REPORT, "w", encoding="utf-8") as f:
        f.write("# UTFANSWF v21 — Row 20 SN Offset/Pivot Probe\n\n")
        f.write("**Context:** DIAGNOSTIC_NON_GATING. No physics, thresholds, datasets, parameters, or gate verdicts changed.\n\n")

        f.write("## Existing SN nuisance\n\n")
        f.write("```json\n")
        f.write(json.dumps(sn.get("nuisance", {}), indent=2))
        f.write("\n```\n\n")

        f.write("## Raw SN nodes\n\n")
        f.write("| Slice | z | residual | sigma | pull | chi2 |\n")
        f.write("|---|---:|---:|---:|---:|---:|\n")
        for n in raw:
            f.write(
                f"| {n['slice']} | {n['z']:.3f} | {n['residual']:.6f} | "
                f"{n['sigma']:.6f} | {n['pull']:.6f} | {n['chi2']:.6f} |\n"
            )

        f.write("\n## Global offset removal\n\n")
        f.write(f"Global residual offset removed: `{global_offset:.6f}`\n\n")
        f.write("| Slice | recentered residual | recentered pull | recentered chi2 diag |\n")
        f.write("|---|---:|---:|---:|\n")
        for n in recentered:
            f.write(
                f"| {n['slice']} | {n['recentered_residual']:.6f} | "
                f"{n['recentered_pull']:.6f} | {n['recentered_chi2_diag']:.6f} |\n"
            )

        f.write("\n## Endpoint offset removal\n\n")
        f.write(f"Endpoint residual offset removed: `{endpoint_offset:.6f}`\n\n")
        f.write("| Slice | endpoint-recentered residual | endpoint-recentered pull | endpoint-recentered chi2 diag |\n")
        f.write("|---|---:|---:|---:|\n")
        for n in endpoint_recentered:
            f.write(
                f"| {n['slice']} | {n['endpoint_recentered_residual']:.6f} | "
                f"{n['endpoint_recentered_pull']:.6f} | "
                f"{n['endpoint_recentered_chi2_diag']:.6f} |\n"
            )

        f.write("\n## Pivot geometry\n\n")
        f.write("| Quantity | Value |\n")
        f.write("|---|---:|\n")
        f.write(f"| raw pivot residual | {pivot_residual:.6f} |\n")
        f.write(f"| raw pivot pull | {pivot_pull:.6f} |\n")
        f.write(f"| global-recentered pivot residual | {pivot_residual_global_recentered:.6f} |\n")
        f.write(f"| global-recentered pivot pull | {pivot_pull_global_recentered:.6f} |\n")

        f.write("\n## Interpretation\n\n")
        for item in interpretation:
            f.write(f"- {item}\n")

        f.write("\n## Verdict\n\n")
        f.write(result["verdict"] + "\n")

    print("Wrote:")
    print(os.path.abspath(OUTPUT_JSON))
    print(os.path.abspath(OUTPUT_REPORT))
    print()
    print("SN offset/pivot probe:")
    print(f"global_offset_removed={global_offset:.6f}")
    print(f"raw_pivot_pull={pivot_pull:.6f} sigma")
    print(f"global_recentered_pivot_pull={pivot_pull_global_recentered:.6f} sigma")
    print(f"verdict={result['verdict']}")


if __name__ == "__main__":
    main()