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


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


def extract_structure(clauses):
    long_neg = [i for i, c in enumerate(clauses) if len(c) > 10 and all(l < 0 for l in c)]
    assert len(long_neg) == 1
    selectors = sorted(abs(l) for l in clauses[long_neg[0]])
    idx = long_neg[0] + 1
    init_ranges = []
    y_groups = []
    for sel in selectors:
        start = idx
        direct = []
        while idx + 1 < len(clauses):
            c1, c2 = clauses[idx], clauses[idx + 1]
            if len(c1) == len(c2) == 2 and c1[0] == -sel and c2[0] == sel and c1[1] == -c2[1] and c1[1] > 0:
                direct.append(c1[1])
                idx += 2
            else:
                break
        outs = []
        for v in direct:
            c1, c2 = clauses[idx], clauses[idx + 1]
            assert len(c1) == len(c2) == 2 and c1 == (v, -c2[1]) and c2 == (-v, c2[1])
            # The actual pattern is (v, -out), (-v, out).
            out = -c1[1]
            assert c2 == (-v, out)
            outs.append(out)
            idx += 2
        init_ranges.append((start, idx))
        y_groups.append(outs)

    init_end = idx
    add_starts = []
    for outs in y_groups:
        s = set(outs)
        pos = None
        for j in range(init_end, len(clauses)):
            if any(abs(l) in s for l in clauses[j]):
                pos = j
                break
        add_starts.append(pos)
    assert add_starts[0] == init_end
    return selectors, init_ranges, y_groups, init_end, add_starts


class Propagator:
    def __init__(self, nvars, clauses, clause_ids):
        self.nvars = nvars
        self.clauses = [clauses[i] for i in clause_ids]
        self.occ = [[] for _ in range(2 * (nvars + 1))]
        for ci, c in enumerate(self.clauses):
            for lit in c:
                self.occ[self.lit_index(lit)].append(ci)
        self.assign = [0] * (nvars + 1)
        self.trail = []
        self.queue = collections.deque()

    @staticmethod
    def lit_index(lit):
        v = abs(lit)
        return 2 * v + (1 if lit < 0 else 0)

    def lit_value(self, lit):
        a = self.assign[abs(lit)]
        if a == 0:
            return 0
        val = a > 0
        return 1 if (val == (lit > 0)) else -1

    def set_lit(self, lit):
        v = abs(lit)
        val = 1 if lit > 0 else -1
        if self.assign[v] == val:
            return True
        if self.assign[v] == -val:
            return False
        self.assign[v] = val
        self.trail.append(v)
        self.queue.append(lit)
        return True

    def propagate(self):
        while self.queue:
            lit = self.queue.popleft()
            for ci in self.occ[self.lit_index(-lit)]:
                c = self.clauses[ci]
                sat = False
                open_lit = 0
                open_count = 0
                for l in c:
                    val = self.lit_value(l)
                    if val == 1:
                        sat = True
                        break
                    if val == 0:
                        open_count += 1
                        open_lit = l
                        if open_count > 1:
                            break
                if sat or open_count > 1:
                    continue
                if open_count == 0:
                    return False
                if not self.set_lit(open_lit):
                    return False
        return True

    def run(self, assumptions):
        mark = len(self.trail)
        oldq = collections.deque(self.queue)
        ok = True
        for lit in assumptions:
            if not self.set_lit(lit):
                ok = False
                break
        if ok:
            ok = self.propagate()
        result = self.assign[:] if ok else None
        while len(self.trail) > mark:
            self.assign[self.trail.pop()] = 0
        self.queue = oldq
        return result


def solve_probe(path, split=None, max_right=None):
    nvars, clauses = read_formula(path)
    selectors, init_ranges, y_groups, init_end, add_starts = extract_structure(clauses)
    n = len(selectors)
    if split is None:
        split = n // 2
    print("selectors", n, "init_end", init_end, "add_starts", add_starts[:5], "...", add_starts[-3:])
    prefix_ids = [0]
    for a, b in init_ranges[:split]:
        prefix_ids.extend(range(a, b))
    prefix_ids.extend(range(add_starts[0], add_starts[split] if split < n else len(clauses)))
    suffix_ids = [0]
    for a, b in init_ranges[split:]:
        suffix_ids.extend(range(a, b))
    suffix_ids.extend(range(add_starts[split], len(clauses)))
    pset, sset = set(prefix_ids), set(suffix_ids)
    pvars = set(abs(l) for i in pset for l in clauses[i])
    svars = set(abs(l) for i in sset for l in clauses[i])
    boundary = sorted(pvars & svars - set(selectors))
    print("clauses prefix/suffix", len(prefix_ids), len(suffix_ids), "boundary", len(boundary), boundary[:10], boundary[-10:])
    P = Propagator(nvars, clauses, prefix_ids)
    S = Propagator(nvars, clauses, suffix_ids)

    def key(assign):
        vals = []
        missing = 0
        for v in boundary:
            a = assign[v]
            if a == 0:
                missing += 1
            vals.append("1" if a > 0 else "0" if a < 0 else "?")
        return "".join(vals), missing

    left = {}
    misscnt = collections.Counter()
    for mask in range(1 << split):
        assumps = []
        for i in range(split):
            lit = selectors[i] if (mask >> i) & 1 else -selectors[i]
            assumps.append(lit)
        a = P.run(assumps)
        if a is None:
            continue
        k, miss = key(a)
        misscnt[miss] += 1
        if miss == 0:
            left.setdefault(k, mask)
    print("left states", len(left), "missing hist", misscnt.most_common(8))
    rbits = n - split
    checked = 0
    for rmask in range(1 << rbits):
        if max_right is not None and checked >= max_right:
            break
        checked += 1
        assumps = []
        for j in range(rbits):
            i = split + j
            lit = selectors[i] if (rmask >> j) & 1 else -selectors[i]
            assumps.append(lit)
        a = S.run(assumps)
        if a is None:
            continue
        k, miss = key(a)
        if miss == 0 and k in left:
            mask = left[k] | (rmask << split)
            print("candidate", mask, bin(mask))
            return mask
    print("no candidate")
    return None


if __name__ == "__main__":
    for arg in sys.argv[1:]:
        print("==", Path(arg).name)
        solve_probe(arg)
