#!/usr/bin/env python3
import collections
import lzma
import os
import sys
from pathlib import Path


def open_dimacs(path: Path):
    if path.suffix == ".xz":
        return lzma.open(path, "rt", encoding="ascii")
    return path.open("rt", encoding="ascii")


def read_blocks(path: Path):
    nvars = 0
    curz = None
    block = []
    with open_dimacs(path) as f:
        pending = []
        for line in f:
            if not line or line[0] == "c":
                continue
            if line[0] == "p":
                nvars = int(line.split()[2])
                continue
            for tok in line.split():
                lit = int(tok)
                if lit == 0:
                    if pending:
                        m = max(abs(x) for x in pending)
                        if curz is not None and m != curz:
                            yield nvars, curz, block
                            block = []
                        curz = m
                        block.append(tuple(pending))
                    pending = []
                else:
                    pending.append(lit)
    if block:
        yield nvars, curz, block


class Expr:
    def __init__(self):
        self.ids = {}
        self.keys = []
        self.not_cache = {}
        self.const0 = self.raw(("const", 0))
        self.const1 = self.raw(("const", 1))

    def raw(self, key):
        old = self.ids.get(key)
        if old is not None:
            return old
        idx = len(self.keys)
        self.ids[key] = idx
        self.keys.append(key)
        return idx

    def var(self, v):
        return self.raw(("var", v))

    def neg(self, h):
        if h == self.const0:
            return self.const1
        if h == self.const1:
            return self.const0
        key = self.keys[h]
        if key[0] == "not":
            return key[1]
        old = self.not_cache.get(h)
        if old is not None:
            return old
        out = self.raw(("not", h))
        self.not_cache[h] = out
        self.not_cache[out] = h
        return out

    def lit(self, hashes, lit):
        h = hashes[abs(lit)]
        return self.neg(h) if lit < 0 else h

    def andor(self, op, args):
        args = list(args)
        if op == "and":
            if self.const0 in args:
                return self.const0
            args = [a for a in args if a != self.const1]
        else:
            if self.const1 in args:
                return self.const1
            args = [a for a in args if a != self.const0]
        flat = []
        for a in args:
            k = self.keys[a]
            if k[0] == op:
                flat.extend(k[1])
            else:
                flat.append(a)
        args = tuple(sorted(set(flat)))
        if not args:
            return self.const1 if op == "and" else self.const0
        if len(args) == 1:
            return args[0]
        return self.raw((op, args))

    def xor(self, args):
        parity = 0
        vals = []
        for a in args:
            if a == self.const0:
                continue
            if a == self.const1:
                parity ^= 1
                continue
            k = self.keys[a]
            if k[0] == "xor":
                vals.extend(k[1])
                parity ^= k[2]
            else:
                vals.append(a)
        vals.sort()
        out = []
        i = 0
        while i < len(vals):
            j = i + 1
            while j < len(vals) and vals[j] == vals[i]:
                j += 1
            if (j - i) & 1:
                out.append(vals[i])
            i = j
        if not out:
            return self.const1 if parity else self.const0
        if len(out) == 1:
            return self.neg(out[0]) if parity else out[0]
        return self.raw(("xor", tuple(out), parity))

    def tt(self, table, args):
        return self.raw(("tt", table, tuple(args)))

    def sboxbit(self, tid, bit, inp, bitfunc):
        return self.raw(("sboxbit", tid, bit, tuple(inp), bitfunc))


class DSU:
    def __init__(self, n):
        self.p = list(range(n))

    def find(self, x):
        while self.p[x] != x:
            self.p[x] = self.p[self.p[x]]
            x = self.p[x]
        return x

    def union(self, a, b):
        ra, rb = self.find(a), self.find(b)
        if ra != rb:
            if ra > rb:
                ra, rb = rb, ra
            self.p[rb] = ra


def lid(lit):
    return 2 * (abs(lit) - 1) + (1 if lit < 0 else 0)


def lit_from_id(i):
    v = i // 2 + 1
    return -v if (i & 1) else v


def eqdet(block, z):
    if any(len(c) == 1 for c in block):
        return None
    bins = [c for c in block if len(c) == 2 and -z in c]
    big = [c for c in block if len(c) > 2 and z in c]
    if len(big) != 1 or len(bins) + 1 != len(big[0]):
        return None
    pats = []
    for c in bins:
        other = [x for x in c if abs(x) != z]
        if len(other) != 1:
            return None
        pats.append(other[0])
    big_lits = [x for x in big[0] if abs(x) != z]
    if sorted(big_lits) != sorted(-x for x in pats):
        return None
    return tuple(sorted(pats, key=lambda x: abs(x)))


def implication_from_unit_gate(block, z):
    units = [c[0] for c in block if len(c) == 1 and abs(c[0]) == z]
    if len(units) != 1:
        return None
    unit = units[0]
    active = []
    for c in block:
        if len(c) == 1:
            continue
        if -unit in c:
            rest = tuple(x for x in c if x != -unit)
            active.append(rest)
    if len(active) == 1 and len(active[0]) == 2:
        return active[0]
    return None


def decode_big_gate(block, z):
    block = [c for c in block if len(c) != 1]
    if not block:
        return None
    bins = [c for c in block if len(c) == 2 and any(abs(x) == z for x in c)]
    bigs = [c for c in block if len(c) > 2 and any(abs(x) == z for x in c)]
    if len(bigs) != 1 or len(bins) + 1 != len(bigs[0]):
        return None
    lits = []
    zsign = None
    for c in bins:
        zl = [x for x in c if abs(x) == z]
        other = [x for x in c if abs(x) != z]
        if len(zl) != 1 or len(other) != 1:
            return None
        if zsign is None:
            zsign = 1 if zl[0] > 0 else -1
        elif zsign != (1 if zl[0] > 0 else -1):
            return None
        lits.append(other[0])
    big_other = [x for x in bigs[0] if abs(x) != z]
    big_z = [x for x in bigs[0] if abs(x) == z]
    if len(big_z) != 1 or (1 if big_z[0] > 0 else -1) != -zsign:
        return None
    if sorted(big_other) != sorted(-x for x in lits):
        return None
    if zsign < 0:
        return "and", tuple(lits)
    return "or", tuple(-x for x in lits)


def brute_func(block, z, fanins):
    block = [c for c in block if len(c) != 1]
    if not block:
        return None
    if len(fanins) > 4:
        return None
    pos = {v: i for i, v in enumerate(fanins)}
    table = []
    for mask in range(1 << len(fanins)):
        possible = []
        for zv in (0, 1):
            ok = True
            for cl in block:
                sat = False
                for lit in cl:
                    v = abs(lit)
                    if v == z:
                        val = zv
                    else:
                        idx = pos.get(v)
                        if idx is None:
                            ok = False
                            break
                        val = (mask >> idx) & 1
                    if (lit > 0 and val) or (lit < 0 and not val):
                        sat = True
                        break
                if not ok or not sat:
                    ok = False
                    break
            if ok:
                possible.append(zv)
        if len(possible) != 1:
            return None
        table.append(possible[0])
    return tuple(table)


def apply_gate(expr, hashes, z, block, eqdets):
    if z in eqdets:
        hashes[z] = expr.andor("and", [expr.lit(hashes, lit) for lit in eqdets[z]])
        return True
    if len(block) == 1 and len(block[0]) == 1:
        hashes[z] = expr.const1 if block[0][0] > 0 else expr.const0
        return True
    gate = decode_big_gate(block, z)
    if gate is not None:
        op, lits = gate
        args = [expr.lit(hashes, lit) for lit in lits]
        hashes[z] = expr.andor(op, args)
        return True
    fanins = sorted({abs(x) for c in block for x in c if abs(x) != z})
    tab = brute_func(block, z, fanins)
    if tab is not None:
        # Recognize 2-input XOR/XNOR to simplify aggressively.
        args = [hashes[v] for v in fanins]
        if len(fanins) == 2 and tab in ((0, 1, 1, 0), (1, 0, 0, 1)):
            h = expr.xor(args)
            hashes[z] = expr.neg(h) if tab[0] else h
        else:
            hashes[z] = expr.tt(tab, args)
        return True
    return False


def main():
    for arg in sys.argv[1:]:
        path = Path(arg)
        blocks = list(read_blocks(path))
        nvars = blocks[0][0] if blocks else 0
        expr = Expr()
        dsu = DSU(2 * nvars + 2)
        binset = set()
        for _nvars, _z, block in blocks:
            for c in block:
                if len(c) == 2 and c[0] != -c[1]:
                    binset.add(tuple(sorted(c)))
        for a, b in list(binset):
            if tuple(sorted((-a, -b))) in binset:
                dsu.union(lid(a), lid(-b))
                dsu.union(lid(-a), lid(b))

        rep_hash = {}
        hashes = [expr.const0] * (nvars + 1)
        for v in range(1, nvars + 1):
            rp = dsu.find(lid(v))
            rn = dsu.find(lid(-v))
            if rp == rn:
                hashes[v] = expr.var(v)
                continue
            base = min(rp, rn)
            h = rep_hash.setdefault(base, expr.var(lit_from_id(base)))
            hashes[v] = h if rp == base else expr.neg(h)
        eqdets = {}
        impls = []
        const_units = []
        decoded = collections.Counter()

        for _nvars, z, block in blocks:
            for c in block:
                if len(c) == 1:
                    const_units.append(c[0])
            pat = eqdet(block, z)
            if pat is not None and len(pat) in (8, 32):
                eqdets[z] = pat
                decoded[f"eq{len(pat)}"] += 1
            if implication_from_unit_gate(block, z) is not None:
                impls.append(implication_from_unit_gate(block, z))

        for lit in const_units:
            hashes[abs(lit)] = expr.const1 if lit > 0 else expr.const0

        for _nvars, z, block in blocks:
            if apply_gate(expr, hashes, z, block, eqdets):
                decoded["gate"] += 1
            else:
                decoded["unknown"] += 1

        rels = collections.defaultdict(dict)
        partial32 = collections.defaultdict(dict)
        for a, b in impls:
            if a < 0 and b > 0:
                prem, concl = -a, b
            elif b < 0 and a > 0:
                prem, concl = -b, a
            else:
                continue
            if prem in eqdets and concl in eqdets and len(eqdets[prem]) == 8 and len(eqdets[concl]) in (8, 32):
                pvars = tuple(abs(x) for x in eqdets[prem])
                cvars = tuple(abs(x) for x in eqdets[concl])
                pval = sum((1 if x > 0 else 0) << i for i, x in enumerate(eqdets[prem]))
                cval = sum((1 if x > 0 else 0) << i for i, x in enumerate(eqdets[concl]))
                rels[(pvars, cvars)][pval] = cval
            elif prem in eqdets and concl in eqdets and len(eqdets[prem]) == 32 and len(eqdets[concl]) == 8:
                pvars = tuple(abs(x) for x in eqdets[prem])
                cvars = tuple(abs(x) for x in eqdets[concl])
                pval = sum((1 if x > 0 else 0) << i for i, x in enumerate(eqdets[prem]))
                cval = sum((1 if x > 0 else 0) << i for i, x in enumerate(eqdets[concl]))
                partial32[(pvars, cvars)][pval] = cval

        full = [(k, v) for k, v in rels.items() if len(v) == 256]
        full.sort(key=lambda kv: max(kv[0][1]))
        table_ids = {}
        for (_pvars, _cvars), table in full:
            table_ids.setdefault(tuple(table[i] for i in range(256)), len(table_ids))
        tables_by_id = [None] * len(table_ids)
        for arr, tid in table_ids.items():
            tables_by_id[tid] = arr

        sbox_tid = None
        for tid, arr in enumerate(tables_by_id):
            if max(arr) < 256 and len(set(arr)) == 256:
                sbox_tid = tid
                break

        linear_from_sbox = {}
        if sbox_tid is not None:
            sarr = tables_by_id[sbox_tid]
            for tid, arr in enumerate(tables_by_id):
                lin = []
                maxbit = max(8, max(arr).bit_length())
                for bit in range(maxbit):
                    found = None
                    target = [(arr[x] >> bit) & 1 for x in range(256)]
                    for coeff in range(256):
                        vals = [((coeff & sarr[x]).bit_count() & 1) for x in range(256)]
                        if vals == target:
                            found = (0, coeff)
                            break
                        vals1 = [v ^ 1 for v in vals]
                        if vals1 == target:
                            found = (1, coeff)
                            break
                    lin.append(found)
                linear_from_sbox[tid] = lin

        word_sources = {}
        for (pvars, cvars), table in full:
            tid = table_ids[tuple(table[i] for i in range(256))]
            inp = tuple(hashes[v] for v in pvars)
            assigned_bits = []
            for bit, cv in enumerate(cvars):
                lin = linear_from_sbox.get(tid, [])
                if sbox_tid is not None and bit < len(lin) and lin[bit] is not None:
                    const, coeff = lin[bit]
                    args = []
                    for b in range(8):
                        if (coeff >> b) & 1:
                            sbitfunc = tuple((tables_by_id[sbox_tid][i] >> b) & 1 for i in range(256))
                            args.append(expr.sboxbit(sbox_tid, b, inp, sbitfunc))
                    if const:
                        args.append(expr.const1)
                    hashes[cv] = expr.xor(args)
                else:
                    bitfunc = tuple((table[i] >> bit) & 1 for i in range(256))
                    hashes[cv] = expr.sboxbit(tid, bit, inp, bitfunc)
                assigned_bits.append(hashes[cv])
            if len(cvars) == 32:
                word_sources[tuple(assigned_bits)] = (inp, tuple(table[i] for i in range(256)))

        composed32 = 0
        for (pvars, cvars), pmap in sorted(partial32.items(), key=lambda kv: max(kv[0][1])):
            src = word_sources.get(tuple(hashes[v] for v in pvars))
            if src is None:
                continue
            inp, table = src
            comp = []
            ok = True
            for x in range(256):
                y = pmap.get(table[x])
                if y is None:
                    ok = False
                    break
                comp.append(y)
            if not ok:
                continue
            tid = table_ids.setdefault(tuple(comp), len(table_ids))
            for bit, cv in enumerate(cvars):
                bitfunc = tuple((comp[i] >> bit) & 1 for i in range(256))
                hashes[cv] = expr.sboxbit(tid, bit, inp, bitfunc)
            composed32 += 1

        # One more deterministic pass after byte table summaries.
        for _nvars, z, block in blocks:
            apply_gate(expr, hashes, z, block, eqdets)

        conflicts = []
        unit_true = unit_false = 0
        for lit in const_units:
            h = expr.lit(hashes, lit)
            if h == expr.const0:
                conflicts.append(lit)
            elif h == expr.const1:
                unit_true += 1
            else:
                unit_false += 1
        print(path.name)
        print(f"  vars={nvars} exprs={len(expr.keys)} decoded={dict(decoded)}")
        print(f"  eqdets={len(eqdets)} impls={len(impls)} full_byte_relations={len(full)} composed32={composed32} tables={len(table_ids)}")
        print(f"  unit_conflicts={len(conflicts)} examples={conflicts[:20]} unit_literals_const_true={unit_true} nonconst={unit_false}")
        for v in range(max(1, nvars - 20), nvars + 1):
            key = expr.keys[hashes[v]]
            if key[0] == "const":
                desc = f"const{key[1]}"
            elif key[0] in ("and", "or", "xor"):
                desc = f"{key[0]}[{len(key[1])}]"
            elif key[0] == "tt":
                desc = f"tt[{len(key[2])}]"
            else:
                desc = str(key[:3])
            print(f"    v{v}: h{hashes[v]} {desc}")
        if os.environ.get("SIM_FINAL"):
            simulate_final(expr, hashes, blocks)
        if os.environ.get("PRINT_MITERS"):
            print_miters(expr, hashes, blocks)
        if os.environ.get("INSPECT_VARS"):
            wanted = [int(x) for x in os.environ["INSPECT_VARS"].replace(",", " ").split()]
            print_vars(expr, hashes, wanted)
        if os.environ.get("TRUTH_VARS"):
            wanted = [int(x) for x in os.environ["TRUTH_VARS"].replace(",", " ").split()]
            truth_vars(expr, hashes, wanted)


def splitmix64(x):
    x = (x + 0x9E3779B97F4A7C15) & ((1 << 64) - 1)
    z = x
    z = ((z ^ (z >> 30)) * 0xBF58476D1CE4E5B9) & ((1 << 64) - 1)
    z = ((z ^ (z >> 27)) * 0x94D049BB133111EB) & ((1 << 64) - 1)
    return z ^ (z >> 31)


def simulate_final(expr, hashes, blocks):
    mask = (1 << 64) - 1
    values = [0] * len(expr.keys)
    rand_cache = {}
    for i, key in enumerate(expr.keys):
        op = key[0]
        if op == "const":
            values[i] = mask if key[1] else 0
        elif op == "var":
            lit = key[1]
            base = rand_cache.setdefault(abs(lit), splitmix64(abs(lit) * 0x1234567 + 0xCAFE))
            values[i] = (~base & mask) if lit < 0 else base
        elif op == "not":
            values[i] = ~values[key[1]] & mask
        elif op == "and":
            v = mask
            for a in key[1]:
                v &= values[a]
            values[i] = v
        elif op == "or":
            v = 0
            for a in key[1]:
                v |= values[a]
            values[i] = v
        elif op == "xor":
            v = mask if key[2] else 0
            for a in key[1]:
                v ^= values[a]
            values[i] = v
        elif op == "tt":
            tab, args = key[1], key[2]
            out = 0
            for bit in range(64):
                idx = 0
                for j, a in enumerate(args):
                    idx |= ((values[a] >> bit) & 1) << j
                out |= (tab[idx] & 1) << bit
            values[i] = out
        elif op == "sboxbit":
            _tag, _tid, _bit, inp, bitfunc = key
            out = 0
            for bit in range(64):
                idx = 0
                for j, a in enumerate(inp):
                    idx |= ((values[a] >> bit) & 1) << j
                out |= (bitfunc[idx] & 1) << bit
            values[i] = out
        else:
            values[i] = splitmix64(i + 0xBAD)

    interesting = []
    for _nvars, _z, block in blocks:
        for c in block:
            if len(c) in (16, 56, 184, 257):
                cm = 0
                for lit in c:
                    hv = hashes[abs(lit)]
                    val = values[hv]
                    cm |= (~val & mask) if lit < 0 else val
                interesting.append((len(c), cm, c[:8]))
    for ln, cm, head in interesting[-20:]:
        print(f"    sim clause len={ln} true_pop={cm.bit_count()} head={head}")
        if ln <= 20:
            pops = []
            # Find the full clause again by matching the prefix.
            for _nvars, _z, block in blocks:
                for c in block:
                    if len(c) == ln and c[:len(head)] == head:
                        for lit in c:
                            hv = hashes[abs(lit)]
                            val = values[hv]
                            lv = (~val & mask) if lit < 0 else val
                            pops.append(lv.bit_count())
                        break
            if pops:
                print(f"      lit_pops={pops}")


def short_key(expr, h):
    key = expr.keys[h]
    if key[0] == "const":
        return f"const{key[1]}"
    if key[0] == "var":
        return f"var({key[1]})"
    if key[0] == "not":
        return f"not:{key[1]}"
    if key[0] in ("and", "or"):
        return f"{key[0]}[{len(key[1])}]"
    if key[0] == "xor":
        return f"xor[{len(key[1])}]^{key[2]}"
    if key[0] == "tt":
        return f"tt[{len(key[2])}]"
    if key[0] == "sboxbit":
        _tag, tid, bit, inp, _bitfunc = key
        return f"sboxbit(t{tid},b{bit},in={inp})"
    return str(key[:3])


def print_miters(expr, hashes, blocks):
    for _nvars, _z, block in blocks:
        for c in block:
            if len(c) not in (16, 56, 184, 257):
                continue
            cnt = collections.Counter()
            reps = {}
            for lit in c:
                h = expr.lit(hashes, lit)
                cnt[h] += 1
                reps.setdefault(h, lit)
            const0 = cnt.get(expr.const0, 0)
            const1 = cnt.get(expr.const1, 0)
            print(f"    miter len={len(c)} uniq={len(cnt)} const0={const0} const1={const1}")
            for h, n in cnt.most_common(24):
                print(f"      n={n:3d} lit={reps[h]:8d} h={h:8d} {short_key(expr, h)}")


def print_expr(expr, h, depth=0, limit=3, seen=None):
    if seen is None:
        seen = set()
    indent = "  " * depth
    key = expr.keys[h]
    print(f"{indent}h{h}: {short_key(expr, h)}")
    if depth >= limit or h in seen:
        return
    seen.add(h)
    op = key[0]
    children = []
    if op == "not":
        children = [key[1]]
    elif op in ("and", "or"):
        children = list(key[1])[:12]
    elif op == "xor":
        children = list(key[1])[:12]
    elif op == "tt":
        children = list(key[2])[:12]
    elif op == "sboxbit":
        children = list(key[3])[:12]
    for ch in children:
        print_expr(expr, ch, depth + 1, limit, seen)


def print_vars(expr, hashes, wanted):
    for v in wanted:
        h = hashes[abs(v)]
        if v < 0:
            h = expr.neg(h)
        print(f"    inspect {v}:")
        print_expr(expr, h)


def expr_support(expr, h, out):
    key = expr.keys[h]
    op = key[0]
    if op == "var":
        out.add(abs(key[1]))
    elif op == "not":
        expr_support(expr, key[1], out)
    elif op in ("and", "or"):
        for ch in key[1]:
            expr_support(expr, ch, out)
    elif op == "xor":
        for ch in key[1]:
            expr_support(expr, ch, out)
    elif op == "tt":
        for ch in key[2]:
            expr_support(expr, ch, out)
    elif op == "sboxbit":
        for ch in key[3]:
            expr_support(expr, ch, out)


def eval_expr(expr, h, vals, memo):
    if h in memo:
        return memo[h]
    key = expr.keys[h]
    op = key[0]
    if op == "const":
        ans = bool(key[1])
    elif op == "var":
        lit = key[1]
        ans = vals[abs(lit)]
        if lit < 0:
            ans = not ans
    elif op == "not":
        ans = not eval_expr(expr, key[1], vals, memo)
    elif op == "and":
        ans = all(eval_expr(expr, ch, vals, memo) for ch in key[1])
    elif op == "or":
        ans = any(eval_expr(expr, ch, vals, memo) for ch in key[1])
    elif op == "xor":
        acc = key[2]
        for ch in key[1]:
            acc ^= int(eval_expr(expr, ch, vals, memo))
        ans = bool(acc)
    elif op == "tt":
        tab, args = key[1], key[2]
        idx = 0
        for i, ch in enumerate(args):
            idx |= int(eval_expr(expr, ch, vals, memo)) << i
        ans = bool(tab[idx])
    elif op == "sboxbit":
        _tag, _tid, _bit, inp, bitfunc = key
        idx = 0
        for i, ch in enumerate(inp):
            idx |= int(eval_expr(expr, ch, vals, memo)) << i
        ans = bool(bitfunc[idx])
    else:
        ans = False
    memo[h] = ans
    return ans


def truth_vars(expr, hashes, wanted):
    for v in wanted:
        h = hashes[abs(v)]
        if v < 0:
            h = expr.neg(h)
        supp = set()
        expr_support(expr, h, supp)
        supp = sorted(supp)
        print(f"    truth {v}: support={len(supp)} {supp[:24]}")
        if len(supp) > 20:
            print("      skipped")
            continue
        ones = 0
        vals = {}
        for mask in range(1 << len(supp)):
            vals.clear()
            for i, var in enumerate(supp):
                vals[var] = bool((mask >> i) & 1)
            if eval_expr(expr, h, vals, {}):
                ones += 1
        print(f"      ones={ones}/{1 << len(supp)}")


if __name__ == "__main__":
    main()
