#!/usr/bin/env python3
"""
QUANTUM ENGINE VERIFICATION — against quantum mechanics, not against itself.

The physics engine was checked against published Lorenz constants. Same standard here.

  PART A  the PUBLISHED ARCHIVE — free, offline, no QPU time
     A1  conservation: counts must sum to the declared shot total
     A2  Born-rule sanity: measured marginals in [0,1], no impossible outcomes
     A3  THE WEIGHT-BIRTH MECHANISM. cosmos_born.pt was made by mapping measured
         bitstrings to uniforms u = int(bits)/2^n, then to weights through the inverse
         normal CDF z = sqrt(2)*erfinv(2u-1). If that pipeline is correct, z over real
         archived shots must come out standard normal: mean 0, sd 1, and the right
         tail fractions. This is checked on her ACTUAL archive, so it verifies the
         mechanism that literally created her weights.
     A4  serial independence is explicitly NOT inferred from histogram counts;
         the archive does not retain within-job shot order.

  PART B  LIVE HARDWARE — the CHSH Bell test
     A local hidden-variable theory cannot exceed |S| = 2. Quantum mechanics permits
     up to 2*sqrt(2) = 2.8284 (Tsirelson's bound). Running CHSH on the backend her
     engine actually uses is the one measurement that cannot be reproduced by any
     classical process, however clever. If S > 2 on her hardware, the quantum path is
     genuinely quantum. If S <= 2, everything downstream is classical randomness with
     a receipt.

     Cost: one job, four circuits.
"""
import math
import os
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
sys.stdout.reconfigure(encoding="utf-8", errors="replace")

# Ingestion lives in archive_io so the loader can be benchmarked on its own
# (`python benchmarks/archive_io.py --bench`) rather than only through this verifier.
# A3 used to expand every archived shot into its own Python float -- 2.4 million of them,
# ~79 MB -- purely to hand a list to statistics.pstdev. The archive is already a
# histogram, so weighted moments give the identical mean, sd and tail fractions from the
# (bitstring, count) pairs directly: measured 22.7x faster, 5.0x less memory, and equal
# to the old path to 1.1e-16.
from archive_io import (JSON_BACKEND, born_z_moments, is_labeled_hardware,  # noqa: E402
                        load_records)

ROOT = Path(__file__).resolve().parents[1]
ARCHIVE = Path(
    os.getenv(
        "COSMOS_QUANTUM_ARCHIVE",
        str(ROOT / "data" / "quantum_measurements_public.jsonl"),
    )
)
R = []


def chk(name, ok, detail):
    R.append((name, ok))
    print(f"  [{'PASS' if ok else 'FAIL'}] {name}\n         {detail}\n")


def load_archive():
    return load_records(ARCHIVE)


print("=" * 80)
print("  PART A — PUBLISHED QUANTUM ARCHIVE (no QPU time)")
print("=" * 80 + "\n")

runs = load_archive()

hardware_runs = [record for record in runs if is_labeled_hardware(record)]
legacy_runs = [
    record for record in runs
    if str(record.get("provider_class") or "") == "legacy_unlabelled"
]
simulator_runs = [
    record for record in runs
    if str(record.get("provider_class") or "") == "classical_simulator"
]
print(f"  archive: {ARCHIVE}")
print(f"  json backend: {JSON_BACKEND}")
print(f"  loaded {len(runs):,} records with measurement counts")
print(f"    labeled IBM hardware: {len(hardware_runs):,}")
print(f"    legacy unlabelled:     {len(legacy_runs):,}")
print(f"    classical simulator:   {len(simulator_runs):,}\n")

# A1 conservation
bad, checked, tot_shots = 0, 0, 0
for d in runs:
    s = sum(d["counts"].values())
    tot_shots += s
    dec = d.get("total_shots")
    if isinstance(dec, int) and dec > 0:
        checked += 1
        if s != dec:
            bad += 1
chk("shot conservation (counts sum to declared total)", bad == 0,
    f"{checked:,} records carried a declared total; {bad} mismatched; "
    f"{tot_shots:,} samples archived across all provider classes")

# A2 Born-rule sanity
neg = sum(1 for d in runs for v in d["counts"].values() if v < 0)
widths = {len(k) for d in runs for k in d["counts"]}
chk("no negative counts; consistent register width", neg == 0 and len(widths) <= 3,
    f"negative counts {neg}; bitstring widths present {sorted(widths)}")

# A3 THE WEIGHT-BIRTH MECHANISM on explicitly labeled IBM hardware shots.
# Weighted moments over the histogram; the shot cap of 40 that stops one high-shot job
# from dominating is part of the estimator and is preserved exactly.
mom = born_z_moments(hardware_runs, cap=40)
m, sd = mom.mean, mom.sd
frac1, frac2 = mom.tail(1.0), mom.tail(2.0)
ok = abs(m) < 0.05 and abs(sd - 1.0) < 0.06 and abs(frac1 - 0.6827) < 0.03 and abs(frac2 - 0.9545) < 0.02
chk("weight-birth pipeline yields a standard normal", ok,
    f"n={mom.n:,.0f}  mean {m:+.4f} (want 0)  sd {sd:.4f} (want 1)  "
    f"|z|<=1 {frac1:.4f} (want 0.6827)  |z|<=2 {frac2:.4f} (want 0.9545)")

# A4 cannot be measured from histograms. Iterating count-dictionary keys creates
# an arbitrary outcome ordering, not the original sequence of hardware shots.
R.append(("serial independence (shot order not retained)", None))
print(
    "  [SKIP] serial independence\n"
    "         archive records are histograms; within-job shot order was not retained\n"
)

# ── PART B — CHSH on real hardware ─────────────────────────────────────────
print("=" * 80)
print("  PART B — CHSH BELL TEST ON HER LIVE HARDWARE")
print("=" * 80 + "\n")


def token():
    # Public release boundary: credentials come only from the current process
    # environment. The verifier never searches local files or vault paths.
    return (
        os.getenv("IBM_QUANTUM_TOKEN")
        or os.getenv("QISKIT_IBM_TOKEN")
        or None
    )


def run_chsh():
    from qiskit import QuantumCircuit
    from qiskit_ibm_runtime import QiskitRuntimeService, SamplerV2
    from qiskit.transpiler.preset_passmanagers import generate_preset_pass_manager

    tok = token()
    if not tok:
        return None, "no IBM token"
    svc = QiskitRuntimeService(channel="ibm_quantum_platform", token=tok)
    be = svc.least_busy(operational=True, simulator=False)

    def bell(theta_a, theta_b):
        qc = QuantumCircuit(2, 2)
        qc.h(0)
        qc.cx(0, 1)                     # |Phi+>
        qc.ry(-2 * theta_a, 0)          # rotate measurement basis
        qc.ry(-2 * theta_b, 1)
        qc.measure([0, 1], [0, 1])
        return qc

    # the standard CHSH angles that reach Tsirelson's bound
    a, ap = 0.0, math.pi / 4
    b, bp = math.pi / 8, 3 * math.pi / 8
    settings = [("AB", a, b), ("AB'", a, bp), ("A'B", ap, b), ("A'B'", ap, bp)]
    pm = generate_preset_pass_manager(backend=be, optimization_level=2, seed_transpiler=5)
    circs = [pm.run(bell(x, y)) for _, x, y in settings]
    job = SamplerV2(mode=be).run(circs, shots=4096)
    jid = job.job_id() if callable(getattr(job, "job_id", None)) else "?"
    res = job.result()

    E = {}
    for i, (name, _, _) in enumerate(settings):
        data = res[i].data
        counts = (data.c if hasattr(data, "c") else data.meas).get_counts()
        tot = sum(counts.values())
        same = sum(v for k, v in counts.items() if k.count("1") % 2 == 0)
        E[name] = (2 * same - tot) / tot
    S = E["AB"] - E["AB'"] + E["A'B"] + E["A'B'"]
    return {"backend": be.name, "job": jid, "E": E, "S": S}, None


try:
    out, err = run_chsh()
except Exception as e:
    out, err = None, f"{type(e).__name__}: {e}"

if out:
    print(f"  backend {out['backend']}   job {out['job']}   4096 shots x 4 settings\n")
    for k, v in out["E"].items():
        print(f"    E({k:<4s}) = {v:+.4f}")
    S = out["S"]
    print(f"\n    S = E(AB) - E(AB') + E(A'B) + E(A'B') = {S:+.4f}")
    print(f"    classical limit          2.0000")
    print(f"    Tsirelson (quantum max)  2.8284\n")
    viol = abs(S) > 2.0
    chk("CHSH violates the classical bound", viol,
        f"|S| = {abs(S):.4f} vs classical 2.0000 — "
        f"{'no local hidden-variable theory can produce this' if viol else 'within classical reach'}")
else:
    R.append(("CHSH on hardware", None))
    print(f"  [SKIP] {err} (offline archive checks remain valid)\n")

print("=" * 80)
passed = sum(1 for _, value in R if value is True)
failed = sum(1 for _, value in R if value is False)
skipped = sum(1 for _, value in R if value is None)
print(f"  {passed} PASSED / {failed} FAILED / {skipped} SKIPPED")
for name, value in R:
    status = "PASS" if value is True else ("FAIL" if value is False else "SKIP")
    print(f"     {status}  {name}")
print("=" * 80)
sys.exit(0 if failed == 0 else 1)