#!/usr/bin/env python3
import lzma
import sys


NETS = {
    13: [
        [(0, 11), (1, 7), (2, 4), (3, 5), (8, 9), (10, 12)],
        [(0, 2), (3, 6), (4, 12), (5, 7), (8, 10)],
        [(0, 8), (1, 3), (2, 5), (4, 9), (6, 11), (7, 12)],
        [(0, 1), (2, 10), (3, 8), (4, 6), (9, 11)],
        [(1, 3), (2, 4), (5, 10), (6, 8), (7, 9), (11, 12)],
        [(1, 2), (3, 4), (5, 8), (6, 9), (7, 10)],
        [(2, 3), (4, 7), (5, 6), (8, 11), (9, 10)],
        [(4, 5), (6, 7), (8, 9), (10, 11)],
        [(3, 4), (5, 6), (7, 8), (9, 10)],
    ],
    16: [
        [(0, 5), (1, 4), (2, 12), (3, 13), (6, 7), (8, 9), (10, 15), (11, 14)],
        [(0, 2), (1, 10), (3, 6), (4, 7), (5, 14), (8, 11), (9, 12), (13, 15)],
        [(0, 8), (1, 3), (2, 11), (4, 13), (5, 9), (6, 10), (7, 15), (12, 14)],
        [(0, 1), (2, 4), (3, 8), (5, 6), (7, 12), (9, 10), (11, 13), (14, 15)],
        [(1, 3), (2, 5), (4, 8), (6, 9), (7, 11), (10, 13), (12, 14)],
        [(1, 2), (3, 5), (4, 11), (6, 8), (7, 9), (10, 12), (13, 14)],
        [(2, 3), (4, 5), (6, 7), (8, 9), (10, 11), (12, 13)],
        [(4, 6), (5, 7), (8, 10), (9, 11)],
        [(3, 4), (5, 6), (7, 8), (9, 10), (11, 12)],
    ],
}

HINTS = {
    16: [
        [(0, 2), (1, 3), (4, 6), (5, 11), (7, 8), (9, 10), (12, 13), (14, 15)],
        [(0, 1), (2, 3), (4, 9), (5, 7), (6, 10), (8, 11), (12, 14), (13, 15)],
        [(1, 5), (2, 4), (3, 7), (6, 9), (8, 12), (10, 13), (11, 14)],
        [(1, 2), (3, 12), (4, 5), (6, 7), (8, 9), (10, 11), (13, 14)],
        [(2, 4), (3, 8), (5, 6), (7, 9), (10, 12), (11, 13)],
        [(3, 5), (6, 8), (7, 10), (9, 12)],
        [(1, 2), (3, 4), (5, 6), (7, 8), (9, 10), (11, 12)],
    ],
}


def pair_offset(n, a, b):
    if a > b:
        a, b = b, a
    off = 0
    for i in range(a):
        off += n - 1 - i
    return off + (b - a - 1)


def parse(path):
    opener = lzma.open if path.endswith(".xz") else open
    clauses = []
    nvars = nclauses = None
    with opener(path, "rt") as f:
        for line in f:
            if not line or line[0] == "c":
                continue
            if line[0] == "p":
                parts = line.split()
                nvars, nclauses = int(parts[2]), int(parts[3])
                continue
            lits = [int(x) for x in line.split()]
            if lits and lits[-1] == 0:
                lits.pop()
            if lits:
                clauses.append(lits)
    assert nclauses == len(clauses), (nclauses, len(clauses))
    return nvars, clauses


class Prop:
    def __init__(self, nvars, clauses):
        self.n = nvars
        self.clauses = clauses
        self.ass = [0] * (nvars + 1)
        self.w0 = [0] * len(clauses)
        self.w1 = [0] * len(clauses)
        self.watch = [[] for _ in range(2 * nvars)]
        for ci, c in enumerate(clauses):
            self.w0[ci] = 0
            self.w1[ci] = 0 if len(c) == 1 else 1
            self.watch[self.idx(c[self.w0[ci]])].append(ci)
            if self.w1[ci] != self.w0[ci]:
                self.watch[self.idx(c[self.w1[ci]])].append(ci)
        self.queue = []

    @staticmethod
    def lit_sign(lit):
        return 1 if lit > 0 else -1

    def idx(self, lit):
        v = abs(lit)
        return (v - 1) * 2 + (1 if lit < 0 else 0)

    def lit_val(self, lit):
        a = self.ass[abs(lit)]
        if a == 0:
            return 0
        return 1 if a == self.lit_sign(lit) else -1

    def set_lit(self, lit):
        v = abs(lit)
        s = self.lit_sign(lit)
        if self.ass[v] == s:
            return True
        if self.ass[v] == -s:
            return False
        self.ass[v] = s
        self.queue.append(lit)
        return True

    def propagate(self):
        while self.queue:
            lit = self.queue.pop()
            wi = self.idx(-lit)
            lst = self.watch[wi]
            out = 0
            scan = 0
            while scan < len(lst):
                ci = lst[scan]
                c = self.clauses[ci]
                if c[self.w0[ci]] == -lit:
                    this_attr = 0
                    other_pos = self.w1[ci]
                else:
                    this_attr = 1
                    other_pos = self.w0[ci]
                other_lit = c[other_pos]
                if self.lit_val(other_lit) == 1:
                    lst[out] = ci
                    out += 1
                    scan += 1
                    continue
                replacement = -1
                for pi, l in enumerate(c):
                    if pi == other_pos:
                        continue
                    if self.lit_val(l) != -1:
                        replacement = pi
                        break
                if replacement >= 0:
                    if this_attr == 0:
                        self.w0[ci] = replacement
                    else:
                        self.w1[ci] = replacement
                    self.watch[self.idx(c[replacement])].append(ci)
                    scan += 1
                    continue
                lst[out] = ci
                out += 1
                scan += 1
                if self.lit_val(other_lit) == -1:
                    del lst[out:]
                    return False, ci
                if not self.set_lit(other_lit):
                    del lst[out:]
                    return False, ci
            del lst[out:]
        return True, -1

    def count_status(self):
        sat = unsat = unresolved = 0
        for c in self.clauses:
            any_un = False
            ok = False
            for l in c:
                v = self.lit_val(l)
                if v == 1:
                    ok = True
                    break
                if v == 0:
                    any_un = True
            if ok:
                sat += 1
            elif any_un:
                unresolved += 1
            else:
                unsat += 1
        return sat, unresolved, unsat


def infer_shape(nvars, clauses):
    # The first layer starts with the lexicographic clique over variables
    # 1..n-1, i.e. comparators (0,1)..(0,n-1).
    best = None
    for t in range(2, 80):
        pos = 0
        good = True
        for i in range(1, t + 1):
            for j in range(i + 1, t + 1):
                if pos >= len(clauses) or clauses[pos] != [-i, -j]:
                    good = False
                    break
                pos += 1
            if not good:
                break
        if good:
            next_ok = pos < len(clauses) and clauses[pos] == [-1, -(t + 1)]
            if next_ok:
                best = t + 1
    n = best
    if n is None:
        return None, None
    negbin_prefix = 0
    for c in clauses:
        if len(c) == 2 and c[0] < 0 and c[1] < 0:
            negbin_prefix += 1
        else:
            break
    per_layer = n * ((n - 1) * (n - 2) // 2)
    blocks = negbin_prefix // per_layer if per_layer else 0
    return n, blocks


def main():
    if len(sys.argv) < 2:
        print("usage: probe_fixed_network.py formula.cnf[.xz]", file=sys.stderr)
        return 2
    path = sys.argv[1]
    nvars, clauses = parse(path)
    n, blocks = infer_shape(nvars, clauses)
    print("nvars", nvars, "clauses", len(clauses), "inferred_n", n, "blocks", blocks)
    if n not in NETS and n not in HINTS:
        return 0
    m = n * (n - 1) // 2
    d = blocks
    net = HINTS.get(n, NETS[n][-d:])
    true_vars = set()
    for layer, comps in enumerate(net):
        for a, b in comps:
            true_vars.add(layer * m + pair_offset(n, a, b) + 1)
    p = Prop(nvars, clauses)
    ok = True
    for v in range(1, d * m + 1):
        if not p.set_lit(v if v in true_vars else -v):
            ok = False
            break
    if ok:
        ok, conflict = p.propagate()
    print("after fixed network", "ok" if ok else f"conflict {conflict}", "assigned", sum(1 for x in p.ass if x), "of", nvars)
    print("clause_status", p.count_status())
    if ok:
        # Greedy completion, no backtracking: try false, then true if immediate conflict.
        # This mutates watches and is only a probe.
        for v in range(1, nvars + 1):
            if p.ass[v] == 0:
                if not p.set_lit(-v):
                    print("direct set conflict", v)
                    break
                ok, conflict = p.propagate()
                if not ok:
                    print("greedy conflict at", v, conflict)
                    break
        print("after greedy", "ok" if ok else f"conflict {conflict}", p.count_status())
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
