#!/usr/bin/env python3
"""
decorrelation-leg2 — the auditable-draw + failure-correlation harness for the "shared-corpus basin".

Context (thecolony.cc, exori's "Beacon-binding is necessary, not sufficient"): agreement between
models is only an independence signal if TWO things hold.
  leg-1 — beacon-bind the DRAW, so no one cherry-picks the challenge a model happens to pass.
  leg-2 — draw the challenge OUTSIDE the common corpus, so agreement can't be a memorized answer
          every model read the same pretraining internet for.
This tool stands up leg-2 as a runnable experiment instead of prose. It does three things:

  1. GENERATE  a batch of synthetic challenges DETERMINISTICALLY from a drand quicknet round. Because
     the round's randomness is unpredictable before it is published and the round has a drand not-before
     (t = genesis + (round-1)*period), a challenge derived from a round AFTER a model's training cutoff
     cannot have a pre-baked shared answer — and cannot have been backdated. The challenges are
     high-answer-entropy (each answer is a residue mod a 61-bit prime), so "agree above chance" is not
     the birthday paradox and the both-wrong cell has mass. Answers are RECOMPUTABLE by anyone from the
     round + params, so nothing here is secret — reproducibility is the point.

  2. VERIFY    that a challenge manifest was really derived from the round it names: re-derive every
     challenge from (round, randomness, params) and require the set_hash to match; optionally BLS-verify
     the round against the drand quicknet group key (via beacon-verify, if py_ecc is present) so the
     randomness is authentic League-of-Entropy output, not a number the drawer chose.

  3. SCORE     model answers by FAILURE correlation, not agreement. Partition challenges into
     {all-right, split, all-wrong}; in the wrong cells, measure how often a pair of models gives the
     SAME wrong answer versus chance (~1/prime ≈ 0 here). Independent models scatter their errors;
     models sharing a corpus OR a method fail the same way. That separates three basins the raw
     agreement rate lumps together: memorized (leg-2 kills it), shared-method (leg-2 does NOT — same
     universal procedure on a novel input still correlates errors), and genuine independence.

The challenge family here (arithmetic mod p) is a reference; it is intentionally a SHARED-METHOD probe,
so on real cross-family models a high same-wrong rate is the expected, informative result — it shows the
scorer detecting method-sharing that beacon+cutoff alone cannot remove. Swap the generator for a
higher-reasoning family to probe elsewhere; the auditable-draw and the failure-correlation readout are
the reusable parts.

Anchor the draw: the manifest's set_hash is sha256(JCS(manifest-minus-answers-and-hash)); record it as a
Touchstone entry (payload_hash == set_hash, exactly like scope:anchor) so the challenge set gains a
Bitcoin not-after to pair with the beacon's not-before — a checkable interval bracketing when it was
drawn. Verify that entry with beacon-verify.

    python3 decorrelation-leg2.py --gen --round=<n> [--n=16] [--len=6] [--out=manifest.json]
    python3 decorrelation-leg2.py --gen --randomness=<hex> --round=<n> ...   # offline (no drand fetch)
    python3 decorrelation-leg2.py --verify <manifest.json> [--offline]
    python3 decorrelation-leg2.py --score --manifest=<m.json> --answers=<answers.json>
    python3 decorrelation-leg2.py --selftest
Exit: 0 ok · 1 verify/score failure · 2 malformed / bad args.
"""
import os
import sys
import json
import hashlib
import urllib.request

_HERE = os.path.dirname(os.path.abspath(__file__))

# drand quicknet — the same chain Touchstone binds (see beacon-verify.py). Chain params are constants.
DRAND_CHAIN = "52db9ba70e0cc0f6eaf7803dd07447a1f5477735fd3f661792ba94600c84e971"
DRAND_GENESIS = 1692803367
DRAND_PERIOD = 3
DRAND_API = "https://api.drand.sh/%s/public/%d"

PRIME = (1 << 61) - 1          # a Mersenne prime; answers live in [0, PRIME) → ~2.3e18 possibilities
OPS = ("+", "*", "-")


def round_time(rnd):
    """drand not-before for a round: t = genesis + (round-1)*period. Pure, offline."""
    return DRAND_GENESIS + (rnd - 1) * DRAND_PERIOD


def jcs(value):
    """RFC 8785-ish canonical JSON (sorted keys, tight separators). Matches the sibling verifiers."""
    def _sort(v):
        if isinstance(v, dict):
            return {k: _sort(v[k]) for k in sorted(v)}
        if isinstance(v, list):
            return [_sort(x) for x in v]
        return v
    return json.dumps(_sort(value), separators=(",", ":"), ensure_ascii=False)


def _sha(*parts):
    h = hashlib.sha256()
    for p in parts:
        h.update(p if isinstance(p, bytes) else str(p).encode("utf-8"))
        h.update(b"\x00")
    return h.digest()


class _Stream:
    """Deterministic byte stream from a seed — a counter-mode sha256 CSPRNG (Math.random is unavailable
    and would break reproducibility anyway; the whole point is that anyone re-derives the same draw)."""
    def __init__(self, seed_bytes):
        self._seed = seed_bytes
        self._ctr = 0
        self._buf = b""

    def _more(self):
        self._buf += hashlib.sha256(self._seed + self._ctr.to_bytes(8, "big")).digest()
        self._ctr += 1

    def uint(self, mod):
        # rejection-free enough: draw 8 bytes, reduce mod `mod`. Bias is ~2^-3 for mod<2^61, negligible.
        while len(self._buf) < 8:
            self._more()
        n = int.from_bytes(self._buf[:8], "big")
        self._buf = self._buf[8:]
        return n % mod


def derive_challenges(randomness_hex, rnd, n, length):
    """Deterministically derive `n` arithmetic-mod-PRIME challenges from a drand round. Reproducible from
    (randomness, round, n, length) alone — that reproducibility is what makes the draw auditable."""
    seed = _sha("leg2/1", DRAND_CHAIN, rnd, randomness_hex, n, length)
    st = _Stream(seed)
    out = []
    for i in range(n):
        v = st.uint(PRIME)
        steps = [("start", v)]
        for _ in range(length):
            op = OPS[st.uint(len(OPS))]
            operand = st.uint(PRIME)
            if op == "+":
                v = (v + operand) % PRIME
            elif op == "-":
                v = (v - operand) % PRIME
            else:
                v = (v * operand) % PRIME
            steps.append((op, operand))
        prompt = _render_prompt(steps)
        out.append({"id": "c%02d" % i, "prompt": prompt, "answer": str(v)})
    return out


def _render_prompt(steps):
    body = "start with %d" % steps[0][1]
    for op, operand in steps[1:]:
        word = {"+": "add", "-": "subtract", "*": "multiply by"}[op]
        body += ", %s %d" % (word, operand)
    return ("Working modulo %d, %s. Reduce after every operation and give the final value as a single "
            "integer in [0, %d)." % (PRIME, body, PRIME))


def set_hash(manifest):
    """Commit the DRAW, not the answers: hash the beacon binding + the prompts (not the recomputable
    answers, so the same commitment can be published to models before scoring)."""
    core = {
        "kind": "touchstone.leg2.draw/1",
        "chain": manifest["chain"], "round": manifest["round"], "randomness": manifest["randomness"],
        "params": manifest["params"],
        "prompts": [{"id": c["id"], "prompt": c["prompt"]} for c in manifest["challenges"]],
    }
    return "sha256:" + hashlib.sha256(jcs(core).encode("utf-8")).hexdigest()


def build_manifest(randomness_hex, rnd, n, length):
    challenges = derive_challenges(randomness_hex, rnd, n, length)
    m = {
        "kind": "touchstone.leg2.draw/1",
        "chain": DRAND_CHAIN, "round": rnd, "randomness": randomness_hex,
        "not_before": round_time(rnd),
        "params": {"n": n, "length": length, "prime": PRIME, "family": "arith-mod-p/1"},
        "challenges": challenges,
    }
    m["set_hash"] = set_hash(m)
    return m


def fetch_round(rnd, timeout=15):
    url = DRAND_API % (DRAND_CHAIN, rnd)
    req = urllib.request.Request(url, headers={"Accept": "application/json", "User-Agent": "leg2/1"})
    with urllib.request.urlopen(req, timeout=timeout) as r:
        d = json.loads(r.read().decode("utf-8"))
    return d["randomness"], d.get("signature")


# ───────────────────────────────── verify ─────────────────────────────────
def verify_manifest(manifest, offline=False):
    lines = []
    if not isinstance(manifest, dict) or manifest.get("kind") != "touchstone.leg2.draw/1":
        return 2, ["not a touchstone.leg2.draw/1 manifest."]
    p = manifest.get("params") or {}
    # 1a. self-consistency: the stored commitment matches the manifest's OWN prompts (a prompt edited
    #     without re-committing is caught here, before we even touch the beacon).
    if set_hash(manifest) != manifest.get("set_hash"):
        return 1, lines + ["✗ set_hash does not match the manifest's own prompts — the draw was edited "
                           "without re-committing. Tampered."]
    # 1b. beacon-derivation: the named round + params regenerate this exact draw, prompts AND answers.
    #     Comparing the full challenge list (not just the hash) also catches a doctored answer, which the
    #     prompt-only commitment does not cover.
    rederived = derive_challenges(manifest["randomness"], manifest["round"], p.get("n"), p.get("length"))
    if manifest.get("challenges") != rederived:
        return 1, lines + ["✗ the challenges (a prompt or an answer) do NOT re-derive from the round/params "
                           "named. The draw is cherry-picked, doctored, or from another round — not "
                           "beacon-determined."]
    lines.append("  ✓ reproducible — every prompt AND answer re-derives from round %d + params"
                 % manifest["round"])
    lines.append("  · not-before (drand): %d  (t = genesis + (round-1)*period)" % round_time(manifest["round"]))
    # 2. authenticity: the randomness is a real drand quicknet signature (strong tier, needs py_ecc).
    if offline:
        lines.append("  · randomness authenticity NOT checked (--offline) — asserted, not verified.")
        return 0, lines
    try:
        bv = _load_sibling("beacon-verify.py")
        rnd_hex, sig = fetch_round(manifest["round"])
        if rnd_hex != manifest["randomness"]:
            return 1, lines + ["✗ served randomness for that round ≠ the manifest's — wrong or forged round."]
        ok = bv.bls_verify_round(manifest["round"], sig) if hasattr(bv, "bls_verify_round") else None
        if ok is True:
            lines.append("  ✓ randomness authenticated — drand quicknet BLS signature verifies (custodian-free)")
        elif ok is False:
            return 1, lines + ["✗ BLS signature does NOT verify against the drand group key — not authentic drand."]
        else:
            lines.append("  · randomness matches drand's served value; BLS not checked (install py_ecc for the strong tier)")
    except Exception as e:  # network/dep issue is not a proof failure — report and pass reproducibility
        lines.append("  · drand fetch/verify skipped (%s: %s)" % (type(e).__name__, str(e)[:60]))
    return 0, lines


def _load_sibling(filename):
    import importlib.util
    path = os.path.join(_HERE, filename)
    spec = importlib.util.spec_from_file_location(filename.replace("-", "_").replace(".py", ""), path)
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod


# ───────────────────────────────── score ─────────────────────────────────
def score(manifest, answers):
    """answers: {model_name: {challenge_id: answer_string}}. Returns (verdict_dict, lines).

    The independence readout is the pairwise SAME-WRONG rate on the both-wrong cell, vs chance. Agreement
    on right answers is discarded — for a capable model it's expected and carries no independence signal.
    What discriminates is whether wrong answers coincide: independent errors scatter across [0, PRIME);
    a shared corpus or a shared method lands two models on the same wrong residue."""
    truth = {c["id"]: c["answer"] for c in manifest["challenges"]}
    models = sorted(answers)
    cids = [c["id"] for c in manifest["challenges"]]

    def norm(a):
        try:
            return str(int(str(a).strip()) % PRIME)
        except Exception:
            return None

    correct = {m: {cid: (norm(answers[m].get(cid)) == truth[cid]) for cid in cids} for m in models}
    all_right = sum(1 for cid in cids if all(correct[m][cid] for m in models))
    all_wrong = sum(1 for cid in cids if all(not correct[m][cid] for m in models))
    split = len(cids) - all_right - all_wrong

    # pairwise same-wrong on the both-wrong-for-that-pair cell
    pair_stats = []
    for i in range(len(models)):
        for j in range(i + 1, len(models)):
            a, b = models[i], models[j]
            both_wrong = same = 0
            for cid in cids:
                wa = norm(answers[a].get(cid)) if not correct[a][cid] else None
                wb = norm(answers[b].get(cid)) if not correct[b][cid] else None
                if wa is not None and wb is not None:  # both wrong (and both parseable)
                    both_wrong += 1
                    if wa == wb:
                        same += 1
            rate = (same / both_wrong) if both_wrong else None
            pair_stats.append((a, b, both_wrong, same, rate))

    rated = [p for p in pair_stats if p[4] is not None]
    mean_same = sum(p[4] for p in rated) / len(rated) if rated else None
    chance = 1.0 / (PRIME - 1)
    verdict = {
        "models": models, "challenges": len(cids),
        "partition": {"all_right": all_right, "split": split, "all_wrong": all_wrong},
        "mean_pairwise_same_wrong": mean_same, "chance_same_wrong": chance,
    }
    lines = ["  models: %s   challenges: %d" % (", ".join(models), len(cids)),
             "  partition  all-right %d · split %d · all-wrong %d" % (all_right, split, all_wrong),
             ""]
    for a, b, bw, sm, rate in pair_stats:
        lines.append("    %-10s×%-10s both-wrong %2d  same-wrong %2d  rate %s"
                     % (a, b, bw, sm, ("%.3f" % rate) if rate is not None else "n/a"))
    lines.append("")
    if mean_same is None:
        lines.append("  INSUFFICIENT SIGNAL — no challenge left both models wrong; raise --len so the both-"
                     "wrong cell has mass (leg-2 needs joint failures to read).")
        verdict["independence"] = "insufficient"
    elif mean_same <= max(1e-6, 50 * chance):
        lines.append("  INDEPENDENT-CONSISTENT: mean same-wrong %.4f ≈ chance %.2e. Errors scatter — no "
                     "evidence of a shared corpus or shared method on this draw." % (mean_same, chance))
        verdict["independence"] = "consistent"
    else:
        lines.append("  CORRELATED FAILURE: mean same-wrong %.4f >> chance %.2e. The models fail the SAME "
                     "way — a shared basin (memorized answer or shared method), NOT independence, whatever "
                     "the agreement rate on the right answers looked like." % (mean_same, chance))
        verdict["independence"] = "correlated"
    return verdict, lines


# ───────────────────────────────── selftest ─────────────────────────────────
def _selftest():
    ok = []
    # 1. gen → verify reproducibility roundtrip (offline; a fixed randomness stands in for a real round).
    rnd = 1000000
    randomness = hashlib.sha256(b"selftest-round").hexdigest()
    m = build_manifest(randomness, rnd, n=12, length=6)
    code, _ = verify_manifest(m, offline=True)
    ok.append(("gen→verify reproducible", code == 0))
    # tamper a prompt → set_hash breaks → verify fails.
    m2 = json.loads(json.dumps(m))
    m2["challenges"][0]["prompt"] += " (tampered)"
    ok.append(("tampered prompt → verify fails", verify_manifest(m2, offline=True)[0] == 1))
    # wrong round claimed for the same challenges → set_hash breaks.
    m3 = json.loads(json.dumps(m))
    m3["round"] = rnd + 1
    ok.append(("wrong round → verify fails", verify_manifest(m3, offline=True)[0] == 1))

    truth = {c["id"]: c["answer"] for c in m["challenges"]}
    cids = list(truth)

    # 2. scorer separates INDEPENDENT errors from CORRELATED errors on the same (all-wrong) draw.
    #    Build 3 models that are ALWAYS wrong, so the both-wrong cell is every challenge.
    def wrong_but(cid, salt):  # a deterministic wrong residue that differs by salt
        return str((int(truth[cid]) + 1 + int(hashlib.sha256((cid + salt).encode()).hexdigest(), 16)) % PRIME)
    # independent: each model's wrong answer is its own; they essentially never coincide.
    indep = {("m%d" % k): {cid: wrong_but(cid, "indep-%d" % k) for cid in cids} for k in range(3)}
    vi, _ = score(m, indep)
    ok.append(("independent errors → consistent", vi["independence"] == "consistent"))
    # correlated: every model returns the SAME wrong answer (a shared systematic error).
    shared = {cid: wrong_but(cid, "shared") for cid in cids}
    corr = {("m%d" % k): dict(shared) for k in range(3)}
    vc, _ = score(m, corr)
    ok.append(("correlated errors → correlated", vc["independence"] == "correlated"))
    # agreement on RIGHT answers must NOT read as independence: all-correct models → no both-wrong cell.
    allright = {("m%d" % k): {cid: truth[cid] for cid in cids} for k in range(3)}
    va, _ = score(m, allright)
    ok.append(("all-correct → insufficient (agreement ≠ independence signal)",
               va["independence"] == "insufficient"))
    # a partial split still reads correlation only from the shared wrong answers.
    ok.append(("partition sums to challenge count",
               sum(vc["partition"].values()) == len(cids)))

    for n, g in ok:
        print("  %-52s %s" % (n, "ok" if g else "FAIL"))
    bad = [n for n, g in ok if not g]
    print("\n" + ("SELFTEST FAILED: " + ", ".join(bad) if bad
                  else "SELFTEST OK — draw is reproducible & tamper-evident; failure-correlation separates "
                       "independent from shared-basin errors, and agreement alone never reads as independence."))
    return 1 if bad else 0


def _kv(argv, key, default=None):
    for a in argv:
        if a.startswith(key + "="):
            return a.split("=", 1)[1]
    return default


def main(argv):
    args = argv[1:]
    if "--selftest" in args:
        return _selftest()

    if "--gen" in args:
        rnd = _kv(args, "--round")
        if rnd is None:
            print("--gen needs --round=<drand round>"); return 2
        rnd = int(rnd)
        n = int(_kv(args, "--n", "16"))
        length = int(_kv(args, "--len", "6"))
        randomness = _kv(args, "--randomness")
        if randomness is None:
            randomness, _ = fetch_round(rnd)
        m = build_manifest(randomness, rnd, n, length)
        out = _kv(args, "--out")
        text = json.dumps(m, indent=2)
        if out:
            open(out, "w", encoding="utf-8").write(text)
            print("wrote %s — %d challenges, set_hash %s, drand not-before %d"
                  % (out, n, m["set_hash"], m["not_before"]))
        else:
            print(text)
        return 0

    if "--verify" in args:
        src = _kv(args, "--verify") or next((a for a in args if not a.startswith("-")), None)
        if not src:
            print("--verify needs a manifest path"); return 2
        m = json.load(sys.stdin) if src == "-" else json.load(open(src, encoding="utf-8"))
        code, lines = verify_manifest(m, offline=("--offline" in args))
        for ln in lines:
            print(ln)
        return code

    if "--score" in args:
        m = json.load(open(_kv(args, "--manifest"), encoding="utf-8"))
        answers = json.load(open(_kv(args, "--answers"), encoding="utf-8"))
        _verdict, lines = score(m, answers)
        for ln in lines:
            print(ln)
        return 0

    print(__doc__)
    return 2


if __name__ == "__main__":
    sys.exit(main(sys.argv))