"""Tseitin CNF encoder for r-step SHA-256 compression (research question C2).

Representation. A *bit* is either a DIMACS literal (non-zero int) or a Python bool constant; a *word* is a list of 32
bits, index 0 = least significant bit. Rotations and shifts are free (list re-indexing; shifted-in bits are the
constant False). Gates with constant inputs are folded. Gates:
  xor(a,b)      : 4 clauses       and(a,b) : 3 clauses      not : literal negation
  ch(e,f,g)     : 4 clauses (z = e ? f : g)   maj(a,b,c) : 6 clauses
  full adder    : sum = xor(xor(x,y),c) (8 clauses), carry = maj(x,y,c) (6 clauses) -> 14 clauses per bit
  add(x, y)     : 32-bit ripple-carry (carry out of bit 31 dropped = mod 2^32); multi-operand sums are chained
                  2-operand adders.
The step and schedule follow `wp34.sha256` exactly (T1 = h + S1(e) + Ch(e,f,g) + K_t + W_t; T2 = S0(a) + Maj(a,b,c);
a' = T1 + T2, e' = d + T1). Message words are always fresh variables (fixing bits = unit clauses), so the
variable/clause counts of an instance do not depend on how many bits are fixed. The IV is either the constant H0
(folded) or 8 free words (semi-free-start).

`build_compression(r, feedforward, free_iv)` returns (cnf, W, IV, OUT) where W = 16 message words, IV = 8 words
(constants or variables) and OUT = 8 output words (after r steps, with or without the IV feed-forward).
`decode_word(model, word)` reads a word back from a solver model. The self-test (`python cnf_sha256.py --selftest`)
fixes all 512 message bits + IV and checks the solved output against `wp34.sha256.compress`.
"""
from __future__ import annotations
import argparse, random, sys, time
from common import S, MASK

K = S.K
H0 = S.H0


class CNF:
    def __init__(self):
        self.nv = 0
        self.clauses: list[list[int]] = []
        self._true = None

    # ---- variables -------------------------------------------------------------------------------------------
    def var(self) -> int:
        self.nv += 1
        return self.nv

    def word(self) -> list[int]:
        return [self.var() for _ in range(32)]

    @staticmethod
    def const_word(v: int) -> list[bool]:
        return [bool((v >> i) & 1) for i in range(32)]

    def true_lit(self) -> int:
        if self._true is None:
            self._true = self.var()
            self.clauses.append([self._true])
        return self._true

    def lit(self, b) -> int:
        """Any bit -> literal (constants mapped onto a pinned TRUE variable)."""
        if b is True:
            return self.true_lit()
        if b is False:
            return -self.true_lit()
        return b

    def add(self, *lits):
        self.clauses.append(list(lits))

    # ---- gates -------------------------------------------------------------------------------------------------
    @staticmethod
    def neg(a):
        if isinstance(a, bool):
            return not a
        return -a

    def xor(self, a, b):
        if isinstance(a, bool):
            return self.neg(b) if a else b
        if isinstance(b, bool):
            return self.neg(a) if b else a
        if a == b:
            return False
        if a == -b:
            return True
        z = self.var()
        self.add(-a, -b, -z); self.add(a, b, -z); self.add(a, -b, z); self.add(-a, b, z)
        return z

    def and_(self, a, b):
        if isinstance(a, bool):
            return b if a else False
        if isinstance(b, bool):
            return a if b else False
        if a == b:
            return a
        if a == -b:
            return False
        z = self.var()
        self.add(-z, a); self.add(-z, b); self.add(z, -a, -b)
        return z

    def or_(self, a, b):
        return self.neg(self.and_(self.neg(a), self.neg(b)))

    def ch(self, e, f, g):
        """z = (e and f) xor (not e and g) = e ? f : g."""
        if isinstance(e, bool):
            return f if e else g
        if isinstance(f, bool) and isinstance(g, bool):
            if f == g:
                return f
            return e if f else self.neg(e)
        if not isinstance(f, bool) and not isinstance(g, bool) and f == g:
            return f  # (never compare a literal with a bool: in Python 1 == True)
        f, g = self.lit(f), self.lit(g)
        z = self.var()
        self.add(-e, -f, z); self.add(-e, f, -z); self.add(e, -g, z); self.add(e, g, -z)
        return z

    def maj(self, a, b, c):
        for x, y, w in ((a, b, c), (b, c, a), (c, a, b)):
            if isinstance(x, bool):
                return self.or_(y, w) if x else self.and_(y, w)
        if a == b or a == c:
            return a
        if b == c:
            return b
        if a == -b:
            return c
        if a == -c:
            return b
        if b == -c:
            return a
        z = self.var()
        self.add(-a, -b, z); self.add(-a, -c, z); self.add(-b, -c, z)
        self.add(a, b, -z); self.add(a, c, -z); self.add(b, c, -z)
        return z

    # ---- words -------------------------------------------------------------------------------------------------
    def xor_w(self, x, y):
        return [self.xor(p, q) for p, q in zip(x, y)]

    def and_w(self, x, y):
        return [self.and_(p, q) for p, q in zip(x, y)]

    def not_w(self, x):
        return [self.neg(p) for p in x]

    @staticmethod
    def rotr(x, k):
        return [x[(i + k) % 32] for i in range(32)]

    @staticmethod
    def shr(x, k):
        return [x[i + k] if i + k < 32 else False for i in range(32)]

    def Sigma0(self, x):
        return self.xor_w(self.xor_w(self.rotr(x, 2), self.rotr(x, 13)), self.rotr(x, 22))

    def Sigma1(self, x):
        return self.xor_w(self.xor_w(self.rotr(x, 6), self.rotr(x, 11)), self.rotr(x, 25))

    def sigma0(self, x):
        return self.xor_w(self.xor_w(self.rotr(x, 7), self.rotr(x, 18)), self.shr(x, 3))

    def sigma1(self, x):
        return self.xor_w(self.xor_w(self.rotr(x, 17), self.rotr(x, 19)), self.shr(x, 10))

    def ch_w(self, e, f, g):
        return [self.ch(p, q, s) for p, q, s in zip(e, f, g)]

    def maj_w(self, a, b, c):
        return [self.maj(p, q, s) for p, q, s in zip(a, b, c)]

    def add_w(self, x, y):
        """Ripple-carry modular adder."""
        out, carry = [], False
        for i in range(32):
            s = self.xor(self.xor(x[i], y[i]), carry)
            if i < 31:
                carry = self.maj(x[i], y[i], carry)
            out.append(s)
        return out

    def sum_w(self, *words):
        acc = words[0]
        for w in words[1:]:
            acc = self.add_w(acc, w)
        return acc

    def eq_w(self, x, y):
        """Assert x == y bitwise."""
        for p, q in zip(x, y):
            self.assert_eq(p, q)

    def assert_eq(self, p, q):
        if isinstance(p, bool) and isinstance(q, bool):
            if p != q:
                self.add(-self.true_lit())
            return
        if isinstance(p, bool):
            p, q = q, p
        if isinstance(q, bool):
            self.add(p if q else -p)
            return
        self.add(-p, q); self.add(p, -q)

    def fix_word(self, w, value: int):
        for i in range(32):
            self.assert_eq(w[i], bool((value >> i) & 1))

    def neq_w_any(self, pairs):
        """Assert that at least one bit differs among the given (word, word) pairs: OR of XORs."""
        diffs = []
        for x, y in pairs:
            for p, q in zip(x, y):
                d = self.xor(p, q)
                if d is True:
                    return
                if d is not False:
                    diffs.append(d)
        if not diffs:
            self.add(-self.true_lit())
        else:
            self.add(*diffs)

    # ---- SHA-256 -----------------------------------------------------------------------------------------------
    def step(self, state, k: int, w):
        a, b, c, d, e, f, g, h = state
        T1 = self.sum_w(h, self.Sigma1(e), self.ch_w(e, f, g), self.const_word(k), w)
        T2 = self.add_w(self.Sigma0(a), self.maj_w(a, b, c))
        return (self.add_w(T1, T2), a, b, c, self.add_w(d, T1), e, f, g)

    def schedule(self, W16, rounds: int):
        W = list(W16)
        for t in range(16, rounds):
            W.append(self.sum_w(self.sigma1(W[t - 2]), W[t - 7], self.sigma0(W[t - 15]), W[t - 16]))
        return W

    def compress(self, IV, W16, rounds: int, feedforward: bool, trace=None):
        W = self.schedule(W16, rounds)
        s = tuple(IV)
        for t in range(rounds):
            s = self.step(s, K[t], W[t])
            if trace is not None:
                trace.append(s)
        if not feedforward:
            return list(s)
        return [self.add_w(x, y) for x, y in zip(IV, s)]

    def stats(self):
        return {"n_vars": self.nv, "n_clauses": len(self.clauses)}

    def to_dimacs(self) -> str:
        lines = ["p cnf %d %d" % (self.nv, len(self.clauses))]
        lines += [" ".join(map(str, c)) + " 0" for c in self.clauses]
        return "\n".join(lines) + "\n"


def build_compression(rounds: int, feedforward: bool = False, free_iv: bool = False, cnf: CNF | None = None):
    cnf = cnf or CNF()
    W = [cnf.word() for _ in range(16)]
    IV = [cnf.word() for _ in range(8)] if free_iv else [CNF.const_word(h) for h in H0]
    OUT = cnf.compress(IV, W, rounds, feedforward)
    return cnf, W, IV, OUT


def decode_word(model_set, word) -> int:
    """model_set: set of true literals (positive ints)."""
    v = 0
    for i, b in enumerate(word):
        if b is True or (not isinstance(b, bool) and ((b > 0 and b in model_set) or (b < 0 and -b not in model_set))):
            v |= 1 << i
    return v


def model_to_set(model) -> set:
    return {l for l in model if l > 0}


def selftest(rounds_list=(1, 4, 16, 64), n_msgs=3, seed=0, verbose=True):
    from pysat.solvers import Cadical195
    rng = random.Random(seed)
    ok = True
    for r in rounds_list:
        for ff in (False, True):
            for _ in range(n_msgs):
                block = bytes(rng.getrandbits(8) for _ in range(64))
                words = S.bytes_to_words(block)
                cnf, W, IV, OUT = build_compression(r, feedforward=ff)
                for w, v in zip(W, words):
                    cnf.fix_word(w, v)
                t0 = time.time()
                with Cadical195(bootstrap_with=cnf.clauses) as s:
                    assert s.solve(), "unexpected UNSAT"
                    m = model_to_set(s.get_model())
                got = [decode_word(m, w) for w in OUT]
                ref = S.compress(S.H0, block, r, feedforward=ff)
                good = got == ref
                ok &= good
                if verbose:
                    print(f"r={r:2d} ff={ff!s:5} vars={cnf.nv:7d} clauses={len(cnf.clauses):8d} "
                          f"solve={time.time()-t0:.3f}s match={good}")
    return ok


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--selftest", action="store_true")
    ap.add_argument("--dimacs", help="write the r-step CNF (free message, fixed IV) to this file")
    ap.add_argument("--rounds", type=int, default=64)
    a = ap.parse_args()
    if a.selftest:
        sys.exit(0 if selftest() else 1)
    if a.dimacs:
        cnf, *_ = build_compression(a.rounds)
        open(a.dimacs, "w").write(cnf.to_dimacs())
        print(cnf.stats())