"""Rerun a Spooky test on your own IBM Quantum account and compare your result with the published one.

    pip install qiskit qiskit-ibm-runtime
    export IBM_QUANTUM_TOKEN=...            # a free IBM Quantum account works
    python3 rerun_bell.py <raw file or URL> [--backend ibm_fez]

It rebuilds the same experiment from the raw file: entangle two qubits (H, then CNOT), then measure each one along the
published angles, with the same question for every shot (Solana's for qubit A, Robinhood Chain's for qubit B,
re-derived here from the two seeds). It runs on your account, then prints your win rate and Bell's S next to the
published ones. Your shots will differ one by one (that's quantum mechanics), but the totals should land close: a few
tenths of a percent, plus the noise of the machine you pick.
"""
import argparse
import hashlib
import json
import math
import os
import sys
import urllib.request


def load(arg):
    if arg.startswith("http"):
        return json.load(urllib.request.urlopen(urllib.request.Request(arg, headers={"user-agent": "spooky-rerun"}), timeout=60))
    with open(arg) as f:
        return json.load(f)


def seed32(seed):
    if seed.startswith("0x") and len(seed) == 66:
        return bytes.fromhex(seed[2:])
    return hashlib.sha256(seed.encode()).digest()


def side_questions(seed, side, shots):
    s, out, chunk = seed32(seed), [], b""
    for i in range(shots):
        if i % 256 == 0:
            chunk = hashlib.sha256(s + side.encode() + (i // 256).to_bytes(4, "big")).digest()
        out.append(chunk[(i % 256) // 8] >> (7 - i % 8) & 1)
    return out


def score(recs):
    wins = sum(1 for x, y, a, b in recs if (a ^ b) == (x & y))
    E = {}
    for xy in ("00", "01", "10", "11"):
        sub = [(a, b) for x, y, a, b in recs if f"{x}{y}" == xy]
        E[xy] = sum(1 if a == b else -1 for a, b in sub) / len(sub) if sub else 0.0
    return wins, E["00"] + E["01"] + E["10"] - E["11"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("raw")
    ap.add_argument("--backend", help="an IBM machine name; default: your account's least busy")
    ap.add_argument("--local", action="store_true", help="a noisy local simulator of IBM Fez instead (needs qiskit-aer; for trying the script)")
    args = ap.parse_args()
    d = load(args.raw)
    if d.get("version") != 3:
        sys.exit("this script reruns two-chain tests (version 3)")

    from qiskit import QuantumCircuit, transpile
    from qiskit.circuit import Parameter
    from qiskit_ibm_runtime import QiskitRuntimeService, SamplerV2

    n = d["shots"]
    xs, ys = side_questions(d["seeds"]["a"], "a", n), side_questions(d["seeds"]["b"], "b", n)
    A, B = d["angles"]["A"], d["angles"]["B"]

    ta, tb = Parameter("theta_a"), Parameter("theta_b")
    qc = QuantumCircuit(2, 2)
    qc.h(0)
    qc.cx(0, 1)
    qc.ry(-ta, 0)
    qc.ry(-tb, 1)
    qc.measure([0, 1], [0, 1])

    if args.local:
        from qiskit_aer import AerSimulator
        from qiskit_ibm_runtime.fake_provider import FakeFez
        backend = AerSimulator.from_backend(FakeFez())
        backend.name = "local simulator of IBM Fez"
    else:
        svc = QiskitRuntimeService(channel="ibm_quantum_platform", token=os.environ["IBM_QUANTUM_TOKEN"], instance=os.environ.get("IBM_QUANTUM_INSTANCE") or None)
        backend = svc.backend(args.backend) if args.backend else svc.least_busy(operational=True, simulator=False, min_num_qubits=2)
    t = transpile(qc, backend=backend, optimization_level=3)
    order = [p.name for p in t.parameters]
    values = [[A[x] if name == "theta_a" else B[y] for name in order] for x, y in zip(xs, ys)]
    print(f"running {n:,} shots on {backend.name}…", flush=True)
    job = SamplerV2(mode=backend).run([(t, values)], shots=1)
    if not args.local:
        print(f"IBM job {job.job_id()}", flush=True)
    data = job.result()[0].data
    bits = getattr(data, next(iter(data.keys())))
    recs = []
    for i, (x, y) in enumerate(zip(xs, ys)):
        s = bits.get_bitstrings(loc=i)[0]
        recs.append((x, y, int(s[-1]), int(s[-2])))
    wins, S = score(recs)
    print(f"yours:     {wins:,} of {n:,} won ({wins / n:.2%}), Bell's S = {S:.3f}")
    print(f"published: {d['score']['wins']:,} of {n:,} won ({d['score']['win_rate']:.2%}), Bell's S = {d['score']['S']:.3f} on {d['label']}")
    print("Without entanglement, nothing beats 75% (S = 2) on average.")


if __name__ == "__main__":
    main()