"""Verify the paper's tables against the measured dataset.

Reads the decoded per-port counts and the job manifest from `../dataset/` and
recomputes, from the measured counts, the quantities reported in the
manuscript: per-row dump/syndrome probabilities and the headline suppression
ratios.  Each manuscript row is resolved to its job through
`data/manifest/job_manifest.csv` (so the extra E4/E4b campaign jobs are
ignored in favour of the manifest's E4f set).

Usage:
    python verify_package.py [path-to-dataset]
"""

from __future__ import annotations

import csv
import os
import sys

from stats import suppression_ratio, wilson_ci, fmt_pct

PKG = sys.argv[1] if len(sys.argv) > 1 else os.path.join(
    os.path.dirname(__file__), "..", "dataset")
PER_PORT = os.path.join(PKG, "data", "processed", "per_port_counts.csv")
MANIFEST = os.path.join(PKG, "data", "manifest", "job_manifest.csv")


def load_per_port() -> dict[str, dict[str, int]]:
    """job_name -> {'total': single_photon_total, 'p0'..'p7': counts}."""
    jobs: dict[str, dict[str, int]] = {}
    with open(PER_PORT, newline="") as fh:
        for r in csv.DictReader(fh):
            j = jobs.setdefault(r["job_name"], {"total": int(r["single_photon_total"])})
            j[f"p{r['port']}"] = int(r["count"])
    return jobs


def load_manifest() -> list[dict[str, str]]:
    with open(MANIFEST, newline="") as fh:
        return list(csv.DictReader(fh))


def job_key(row: dict[str, str], jobs: dict) -> str | None:
    cand = [row["raw_result_file"].removesuffix(".result.json"),
            row["expected_job_name"]]
    for c in cand:
        if c in jobs:
            return c
    return None


def metric_count(exp: str, table: str, config: str, j: dict) -> tuple[int, str]:
    """Return (event_count_for_metric, metric_name) for a job's port dict."""
    if exp == "5" and table == "Table stab":
        return j["p0"] + j["p1"] + j["p2"] + j["p3"], "syndrome(0-3)"
    if exp == "5" and table == "Suppression control":
        return j["p0"] + j["p1"] + j["p2"] + j["p3"], "syndrome(0-3)"
    if exp == "5" and table == "Table parity":
        port = {"Sx": 1, "Sy": 2, "Sz": 3}[
            next(s for s in ("Sx", "Sy", "Sz") if s in config)]
        return j[f"p{port}"], f"port {port}"
    return j["p0"], "dump(port 0)"


def main() -> None:
    if not os.path.exists(PER_PORT):
        raise SystemExit(f"per_port_counts.csv not found at {PER_PORT}")
    jobs = load_per_port()
    manifest = load_manifest()

    print("=" * 78)
    print("Per-row real-data values (from manifest-resolved jobs)")
    print("=" * 78)
    print(f"{'Exp':>3} {'configuration':<34}{'n':>7}  {'metric':<14}{'p (95% CI)'}")
    missing = []
    by = {}  # (exp) -> list of (config, job, count, total)
    for row in manifest:
        key = job_key(row, jobs)
        if key is None:
            missing.append(row["configuration"])
            continue
        j = jobs[key]
        if row["experiment"] == "5" and row["table"] == "Table parity":
            # selectivity = expected syndrome port / syndrome-subspace ports (0-3)
            port = {"Sx": 1, "Sy": 2, "Sz": 3}[
                next(s for s in ("Sx", "Sy", "Sz") if s in row["configuration"])]
            syn = j["p0"] + j["p1"] + j["p2"] + j["p3"]
            cnt = j[f"p{port}"]
            ci = wilson_ci(cnt, syn)
            metric, n_disp = f"select p{port}/(0-3)", syn
        else:
            cnt, metric = metric_count(row["experiment"], row["table"],
                                       row["configuration"], j)
            ci = wilson_ci(cnt, j["total"])
            n_disp = j["total"]
        print(f"{row['experiment']:>3} {row['configuration']:<34}{n_disp:>7}  "
              f"{metric:<18}{ci.point:.4f} [{ci.low:.4f}, {ci.high:.4f}]")
        by.setdefault(row["experiment"], []).append(
            (row["configuration"], cnt, j["total"]))

    if missing:
        print("\nMISSING jobs for:", ", ".join(missing))

    print("\n" + "=" * 78)
    print("Headline results from REAL data vs paper")
    print("=" * 78)

    # Exp 1: mean / std / best dump across modes
    e1 = by.get("1", [])
    if e1:
        ps = [c / n for _, c, n in e1]
        import statistics
        print(f"Exp1 BALANCE: mean dump {statistics.mean(ps):.3f}, "
              f"std {statistics.pstdev(ps):.3f}, best {min(ps):.3f}   "
              f"(paper: 0.188, 0.030, 0.140)")

    # Exp 2: pooled neutral leakage + suppression vs control
    e2 = {c: (cnt, n) for c, cnt, n in by.get("2", [])}
    if e2:
        ctrl = e2["Control non-neutral"]
        neu = [e2[k] for k in ("|0>-|1> neutral", "|0>-|4> neutral",
                               "Balanced 4+4-")]
        nx = sum(c for c, _ in neu); nn = sum(n for _, n in neu)
        pooled = wilson_ci(nx, nn)
        sup = suppression_ratio(ctrl[0], ctrl[1], nx, nn)
        print(f"Exp2 neutral pooled leakage {fmt_pct(pooled.point)} "
              f"[{fmt_pct(pooled.low)}, {fmt_pct(pooled.high)}]  (paper 0.6%)")
        print(f"Exp2 suppression {sup.point:.1f}x [{sup.low:.1f}, {sup.high:.1f}]"
              f"   (paper 31.6x [27.2, 36.7])")

    # Exp 5: code mean syndrome + suppression vs control
    e5 = {c: (cnt, n) for c, cnt, n in by.get("5", [])}
    code_keys = ["Code b0 xor b1", "Code b0 xor b2", "Code b1 xor b2",
                 "Code b0 xor b1 xor b2"]
    if all(k in e5 for k in code_keys):
        cx = sum(e5[k][0] for k in code_keys); cn = sum(e5[k][1] for k in code_keys)
        mean = wilson_ci(cx, cn)
        ctrl5 = e5["Mode |0> control"]
        sup5 = suppression_ratio(ctrl5[0], ctrl5[1], cx, cn)
        print(f"Exp5 mean syndrome leakage {fmt_pct(mean.point)} "
              f"[{fmt_pct(mean.low)}, {fmt_pct(mean.high)}]  (paper 2.32%)")
        print(f"Exp5 suppression {sup5.point:.1f}x [{sup5.low:.1f}, {sup5.high:.1f}]"
              f"   (paper 23.7x [22.2, 25.3])")


if __name__ == "__main__":
    main()
