"""Check a Spooky Bell test yourself. Python 3.8+, standard library only: nothing of ours is imported.

    python3 verify_bell.py test-1.json                                   # a raw file from the site
    python3 verify_bell.py https://<site>/data/test-1.json               # or straight from the site

What it checks (two-chain tests):
  1. The questions: qubit A's come from Solana's seed, qubit B's from Robinhood Chain's.
  2. Each chain's half: the answers are exactly the bits each chain records.
  3. The combined score, the odds without entanglement, and the circuit fingerprint.
Older tests:
  1. The questions: every shot's (x, y) is the one the public seed dictates, and so is every holder pairing.
  2. The score: wins, Bell's S, and the odds that anything without entanglement (75% per shot at best) does as well.
  3. Holder pairs: every pair and question follows from the seed and the holders frozen at the draw. (Older tests:
     the Merkle root over every shot.)
  4. The circuit fingerprint: sha256 of the seed, settings, machine and exact circuit matches the published one.
What it can't check: that IBM produced the answers. The IBM job ID and the exact circuit are included, so anyone with an
IBM Quantum account can rerun it.
"""
import hashlib
import json
import math
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-verify"}), timeout=60))
    with open(arg) as f:
        return json.load(f)


def tail(k, n, p):
    """P(Binomial(n, p) >= k), in log space."""
    if k <= 0:
        return 1.0
    logs = [math.lgamma(n + 1) - math.lgamma(i + 1) - math.lgamma(n - i + 1) + i * math.log(p) + (n - i) * math.log1p(-p) for i in range(k, n + 1)]
    top = max(logs)
    return math.exp(top) * sum(math.exp(x - top) for x in logs)


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


def house_questions(seed, shots):
    s, out, chunk = seed32(seed), [], b""
    for i in range(shots):
        bit = 2 * i
        if bit % 256 == 0:
            chunk = hashlib.sha256(s + b"q" + (bit // 256).to_bytes(4, "big")).digest()
        v = int.from_bytes(chunk, "big") >> (254 - bit % 256) & 3
        out.append([v >> 1, v & 1])
    return out


def pairs_for(seed, members, prefs, max_pairs):
    n, s = len(members), seed32(seed)
    if n < 2:
        return []
    step = 1 + int.from_bytes(hashlib.sha256(s + b"step").digest(), "big") % (n - 1)
    while math.gcd(step, n) != 1:
        step = 1 if step + 1 >= n else step + 1
    shift = int.from_bytes(hashlib.sha256(s + b"shift").digest(), "big") % n
    q = lambda pos: prefs[pos] - 1 if prefs[pos] else hashlib.sha256(s + b"m" + pos.to_bytes(4, "big")).digest()[31] & 1
    out = []
    for k in range(min(n // 2, max_pairs)):
        ia, ib = (step * 2 * k + shift) % n, (step * (2 * k + 1) + shift) % n
        out.append((members[ia].lower(), members[ib].lower(), q(ia), q(ib)))
    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_v2(d):
    seed, recs = d["seed"], d["records"]
    ok = True
    print(f"Test on {d['label']} ({d['backend']}), seed {seed}")
    house = house_questions(seed, d["house_shots"])
    q_ok = len(recs) == len(house) and all(r[:2] == q for r, q in zip(recs, house))
    ok &= q_ok
    print(f"  questions .............. {'every shot matches the seed' if q_ok else 'DO NOT MATCH THE SEED'}")
    wins, S = score(recs)
    n = len(recs)
    p = tail(wins, n, 0.75)
    same = wins == d["score"]["wins"] and abs(S - d["score"]["S"]) < 1e-9
    ok &= same
    print(f"  shots won .............. {wins:,} of {n:,} ({wins / n:.2%}); without entanglement, 75% at most")
    print(f"  Bell's S ............... {S:.4f}; Einstein's limit is 2")
    print(f"  odds without entanglement {p:.3g}" if p > 0 else "  odds without entanglement < 1e-300")
    print(f"  matches published score  {'yes' if same else 'NO'}")
    pairs = d.get("pairs") or []
    if pairs:
        if d.get("members") is not None:
            want = pairs_for(seed, d["members"], d["prefs"], len(pairs))
            p_ok = [(x["a"].lower(), x["b"].lower(), x["qa"], x["qb"]) for x in pairs] == want
            ok &= p_ok
            print(f"  holder pairing ......... {'every pair and question matches the seed' if p_ok else 'DOES NOT MATCH'}")
        w = sum(1 for x in pairs if (x["ra"] ^ x["rb"]) == (x["qa"] & x["qb"]))
        print(f"  holder pairs won ....... {w} of {len(pairs)}")
    commit = hashlib.sha256(json.dumps({"seed": seed, "house": house, "pairs": [[x["qa"], x["qb"]] for x in pairs], "backend": d["backend"], "circuit": d["circuit_qasm"]}, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
    c_ok = commit == d["commit"]
    ok &= c_ok
    print(f"  circuit fingerprint .... {'matches' if c_ok else 'DOES NOT MATCH'} ({commit[:16]}…)")
    if d.get("job_id"):
        print(f"  IBM job ................ {d['job_id']}")
    if d.get("txs"):
        for name, tx in d["txs"].items():
            print(f"  on-chain {name:<14} https://solscan.io/tx/{tx}")
    return ok


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 main_v3(d):
    recs, n = d["records"], d["shots"]
    ok = True
    print(f"Test on {d['label']} ({d['backend']})")
    print(f"  Solana's seed .......... {d['seeds']['a']}")
    print(f"  Robinhood Chain's seed . {d['seeds']['b']}")
    xs, ys = side_questions(d["seeds"]["a"], "a", n), side_questions(d["seeds"]["b"], "b", n)
    qa, qb = all(r[0] == x for r, x in zip(recs, xs)), all(r[1] == y for r, y in zip(recs, ys))
    ok &= qa and qb and len(recs) == n
    print(f"  questions .............. {'Solana' if qa else 'SOLANA WRONG'}'s and {'Robinhood' if qb else 'ROBINHOOD WRONG'}'s both match their seeds")
    ha, hb = bytes.fromhex(d["half_a"]), bytes.fromhex(d["half_b"])
    bit = lambda b, i: b[i // 8] >> (7 - i % 8) & 1
    halves = all(r[2] == bit(ha, i) and r[3] == bit(hb, i) for i, r in enumerate(recs))
    ok &= halves
    print(f"  each chain's half ...... {'matches every shot' if halves else 'DOES NOT MATCH'} (qubit A {sum(r[2] for r in recs) / n:.1%} ones, qubit B {sum(r[3] for r in recs) / n:.1%} ones: noise alone)")
    wins, S = score(recs)
    p = tail(wins, n, 0.75)
    same = wins == d["score"]["wins"] and abs(S - d["score"]["S"]) < 1e-9
    ok &= same
    print(f"  together ............... {wins:,} of {n:,} won ({wins / n:.2%}); without entanglement, 75% at most")
    print(f"  Bell's S ............... {S:.4f}; Einstein's limit is 2")
    print(f"  odds without entanglement {p:.3g}" if p > 0 else "  odds without entanglement < 1e-300")
    commit = hashlib.sha256(json.dumps({"seeds": [d["seeds"]["a"], d["seeds"]["b"]], "shots": n, "backend": d["backend"], "circuit": d["circuit_qasm"]}, sort_keys=True, separators=(",", ":")).encode()).hexdigest()
    c_ok = commit == d["commit"]
    ok &= c_ok
    print(f"  circuit fingerprint .... {'matches' if c_ok else 'DOES NOT MATCH'} ({commit[:16]}…)")
    if d.get("job_id"):
        print(f"  IBM job ................ {d['job_id']}")
    for name, tx in (d.get("txs") or {}).items():
        print(f"  {name:<22} {'https://robin.etherscan.io/tx/' if name.startswith('rh') else 'https://solscan.io/tx/'}{tx}")
    return ok


def main():
    d = load(sys.argv[1] if len(sys.argv) > 1 else input("raw file or URL: ").strip())
    if d.get("version") == 3:
        ok = main_v3(d)
        print("ALL CHECKS PASS" if ok else "SOMETHING DOESN'T MATCH")
        sys.exit(0 if ok else 1)
    if d.get("version") == 2:
        ok = main_v2(d)
        print("ALL CHECKS PASS" if ok else "SOMETHING DOESN'T MATCH")
        sys.exit(0 if ok else 1)
    seed, recs, per = d["seed"], d["records"], d["shots_per_block"]
    ok = True
    print(f"Test on {d['label']} ({d['backend']}), seed {seed!r}")

    want = []
    for i in range(d["blocks"]):
        b = hashlib.sha256(f"q:{seed}:{i}".encode()).digest()[0]
        want.append((b & 1, (b >> 1) & 1))
    q_ok = all((r[0], r[1]) == want[k // per] for k, r in enumerate(recs)) and len(recs) == d["blocks"] * per
    ok &= q_ok
    print(f"  questions .............. {'every shot matches the seed' if q_ok else 'DO NOT MATCH THE SEED'}")

    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)
    S = E["00"] + E["01"] + E["10"] - E["11"]
    n = len(recs)
    p = tail(wins, n, 0.75)
    s = d["score"]
    same = wins == s["wins"] and abs(S - s["S"]) < 1e-9
    ok &= same
    print(f"  shots won .............. {wins:,} of {n:,} ({wins / n:.2%}); without entanglement, 75% at most")
    print(f"  Bell's S ............... {S:.4f}; Einstein's limit is 2")
    print(f"  odds without entanglement {p:.3g}" if p > 0 else "  odds without entanglement < 1e-300")
    print(f"  matches published score  {'yes' if same else 'NO'}")

    sh = hashlib.sha256(seed.encode()).digest()
    level = [hashlib.sha256(b"pair" + sh + i.to_bytes(4, "big") + bytes(r)).digest() for i, r in enumerate(recs)]
    while len(level) > 1:
        level = [hashlib.sha256(min(l, r) + max(l, r)).digest() for l, r in ((level[k], level[k + 1] if k + 1 < len(level) else level[k]) for k in range(0, len(level), 2))]
    r_ok = level[0].hex() == d["pairs_root"]
    ok &= r_ok
    print(f"  pairs root ............. {'matches' if r_ok else 'DOES NOT MATCH'} ({d['pairs_root'][:16]}…)")

    commit = hashlib.sha256(json.dumps({"seed": seed, "blocks": d["blocks"], "shots_per_block": per, "backend": d["backend"], "circuit": d["circuit_qasm"]}, sort_keys=True).encode()).hexdigest()
    c_ok = commit == d["commit"]
    ok &= c_ok
    print(f"  circuit fingerprint .... {'matches' if c_ok else 'DOES NOT MATCH'} ({commit[:16]}…)")
    if d.get("job_id"):
        print(f"  IBM job ................ {d['job_id']}")
    print("ALL CHECKS PASS" if ok else "SOMETHING DOESN'T MATCH")
    sys.exit(0 if ok else 1)


if __name__ == "__main__":
    main()