"""3a. SAT/SMT preimage scaling on r-step SHA-256 compression (standard IV, no feed-forward), question C2.

Cell = (r, k, seed, solver). A seeded random 512-bit message M* is drawn; target = compress(H0, M*, r, ff=False),
so a preimage exists. The k *free* bits are the first k message bits in (word, LSB-first) order: W_0 bits 0..31,
then W_1, ... ; all other message bits are fixed to M* by unit clauses. Solvers: CaDiCaL 1.9.5 (pysat Cadical195),
Kissat 4.0.4 (subprocess, --time), and z3 (bit-vector theory) on r <= 16, k <= 32. Every SAT answer is verified
by recomputing the hash with wp34.sha256.compress.

Modes:
  --list-cells            print one JSON cell spec per line (used by cluster_driver.py)
  --cell '<json>' --out f run one cell in this process (fresh process => peak RSS is per cell)
  --aggregate DIR --out f gather per-cell JSONs into the experiment result with provenance
"""
from __future__ import annotations
import argparse, json, os, pathlib, random, sys, time
from common import S, provenance, peak_rss_mb, MASK, mean_sd
import cnf_sha256 as C
import solvers as SV

R_LIST = [4, 8, 12, 16, 20, 24, 28, 32, 40, 48, 64]
K_LIST = [16, 24, 32, 40, 48, 56, 64]
TIMEOUT = 300.0


def free_positions(k: int, placement: str):
    """'first': the first k message bits (W_0 LSB-first, then W_1, ...).
    'spread': k bits spread evenly over the 16 words (bit i -> word i mod 16, position i // 16), so that no single
    word carries all the freedom (for r <= 16 a single free word makes the instance invertible step by step)."""
    if placement == "first":
        return [(i // 32, i % 32) for i in range(k)]
    if placement == "spread":
        return [(i % 16, i // 16) for i in range(k)]
    raise ValueError(placement)


def instance(r: int, k: int, seed: int, placement: str = "first"):
    rng = random.Random(1_000_003 * seed + 1_000 * r + k)
    block = bytes(rng.getrandbits(8) for _ in range(64))
    words = S.bytes_to_words(block)
    target = S.compress(S.H0, block, r, feedforward=False)
    return block, words, target, free_positions(k, placement)


def list_cells(seeds, solvers, placements=("first", "spread")):
    cells = []
    for placement in placements:
        for r in R_LIST:
            for k in K_LIST:
                for seed in seeds:
                    for sol in solvers:
                        if sol == "z3" and not (r <= 16 and k <= 32):
                            continue
                        if placement == "spread" and sol != "cadical":
                            continue
                        cells.append({"r": r, "k": k, "seed": seed, "solver": sol, "placement": placement})
    return cells


def run_cell(cell: dict, timeout: float) -> dict:
    r, k, seed, sol = cell["r"], cell["k"], cell["seed"], cell["solver"]
    placement = cell.setdefault("placement", "first")
    block, words, target, free = instance(r, k, seed, placement)
    freeset = set(free)
    t0 = time.time()
    res = {"cell": cell, "timeout_s": timeout, "target_hex": "".join("%08x" % v for v in target),
           "fixed_message_hex": block.hex(), "free_bits": "%d bits, placement=%s" % (k, placement)}
    if sol in ("cadical", "kissat"):
        cnf, W, IV, OUT = C.build_compression(r, feedforward=False)
        for j, w in enumerate(W):
            for i in range(32):
                if (j, i) not in freeset:
                    cnf.assert_eq(w[i], bool((words[j] >> i) & 1))
        for o, v in zip(OUT, target):
            cnf.fix_word(o, v)
        res.update(cnf.stats())
        res["encode_s"] = round(time.time() - t0, 3)
        if sol == "cadical":
            out = SV.run_cadical(cnf.clauses, timeout, n_vars=cnf.nv, seed=seed)
        else:
            out = SV.run_kissat(cnf.clauses, cnf.nv, timeout, seed=seed)
        model = out.pop("model", None)
        res.update(out)
        if out["status"] == "sat":
            found = [C.decode_word(model, w) for w in W]
            res["preimage_hex"] = "".join("%08x" % v for v in found)
            res["verified"] = S.compress(S.H0, S.words_to_bytes(found), r, feedforward=False) == target
    elif sol == "z3":
        def build(z3, s):
            Wv = [z3.BitVec("W%d" % j, 32) for j in range(16)]
            for j in range(16):
                for i in range(32):
                    if (j, i) not in freeset:
                        s.add(z3.Extract(i, i, Wv[j]) == z3.BitVecVal((words[j] >> i) & 1, 1))
            OUT = SV.z3_compress(z3, list(S.H0), Wv, r, False)
            for o, v in zip(OUT, target):
                s.add(o == z3.BitVecVal(v, 32))
            return {"W%d" % j: Wv[j] for j in range(16)}
        res["encode_s"] = None
        out = SV.run_z3(build, timeout, seed=seed)
        model = out.pop("model", None)
        res.update(out)
        if out["status"] == "sat":
            found = [model["W%d" % j] for j in range(16)]
            res["preimage_hex"] = "".join("%08x" % v for v in found)
            res["verified"] = S.compress(S.H0, S.words_to_bytes(found), r, feedforward=False) == target
    res["total_wall_s"] = round(time.time() - t0, 3)
    res["peak_rss_mb"] = max(peak_rss_mb(), peak_rss_mb(children=True))
    res["hostname"] = provenance.host_info()["hostname"]
    return res


def aggregate(d: str, out: str, started: float):
    cells = []
    for p in sorted(pathlib.Path(d).glob("*.json")):
        doc = json.loads(p.read_text())
        if isinstance(doc, dict) and "cell" in doc:  # skip the driver's skipped_cells.json
            cells.append(doc)
    summary = {}
    for placement in ("first", "spread"):
      for sol in ("cadical", "kissat", "z3"):
        for r in R_LIST:
            for k in K_LIST:
                cs = [c for c in cells if c["cell"]["solver"] == sol and c["cell"]["r"] == r and c["cell"]["k"] == k
                      and c["cell"].get("placement", "first") == placement]
                if not cs:
                    continue
                times = sorted(c["wall_s"] for c in cs)
                solved = [c for c in cs if c["status"] == "sat"]
                conf = sorted(c["conflicts"] for c in solved if c.get("conflicts") is not None)
                key = f"{sol}:r{r}:k{k}" + ("" if placement == "first" else ":spread")
                summary[key] = {
                    "solver": sol, "placement": placement, "r": r, "k": k, "n_runs": len(cs), "n_sat": len(solved),
                    "n_timeout": sum(c["status"] == "timeout" for c in cs),
                    "solved_fraction": len(solved) / len(cs),
                    "median_wall_s": times[len(times) // 2],
                    "median_wall_s_solved": sorted(c["wall_s"] for c in solved)[len(solved) // 2] if solved else None,
                    "median_conflicts_solved": conf[len(conf) // 2] if conf else None,
                    "all_verified": all(c.get("verified") for c in solved) if solved else None,
                    "n_vars": cs[0].get("n_vars"), "n_clauses": cs[0].get("n_clauses"),
                    "max_peak_rss_mb": max(c["peak_rss_mb"] for c in cs),
                    "hosts": sorted({c["hostname"] for c in cs}),
                }
    provenance.record("sha256_preimage_sat_scaling",
                      {"r_list": R_LIST, "k_list": K_LIST, "timeout_s": TIMEOUT, "feedforward": False,
                       "iv": "standard H0", "free_bits": {"first": "first k message bits, W_0 LSB-first", "spread": "bit i -> word i mod 16, position i//16 (cadical only)"},
                       "solvers": {"cadical": "pysat Cadical195 (CaDiCaL 1.9.5)", "kissat": "kissat 4.0.4 binary",
                                   "z3": "z3 bit-vector, r<=16 & k<=32"}},
                      {"summary": summary, "cells": cells}, seed="per-cell (see cells)",
                      extra={"n_cells": len(cells)}, out=out, started=started)
    print("aggregated", len(cells), "cells ->", out)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--list-cells", action="store_true")
    ap.add_argument("--seeds", default="0,1,2")
    ap.add_argument("--solvers", default="cadical,kissat,z3")
    ap.add_argument("--placements", default="first,spread")
    ap.add_argument("--cell")
    ap.add_argument("--timeout", type=float, default=TIMEOUT)
    ap.add_argument("--aggregate")
    ap.add_argument("--out")
    ap.add_argument("--seed", type=int, default=0, help="(single-cell convenience) seed when --cell omits it")
    a = ap.parse_args()
    t0 = time.time()
    if a.list_cells:
        for c in list_cells([int(x) for x in a.seeds.split(",")], a.solvers.split(","), a.placements.split(",")):
            print(json.dumps(c))
    elif a.cell:
        cell = json.loads(a.cell)
        cell.setdefault("seed", a.seed)
        res = run_cell(cell, a.timeout)
        if a.out:
            pathlib.Path(a.out).parent.mkdir(parents=True, exist_ok=True)
            pathlib.Path(a.out).write_text(json.dumps(res, indent=1, default=str))
        print(json.dumps({k: v for k, v in res.items() if k not in ("fixed_message_hex",)}, default=str))
    elif a.aggregate:
        aggregate(a.aggregate, a.out, t0)