#!/usr/bin/env python3
import ast
import bisect
import lzma
import sys
import time


def read_dimacs_xz(path):
    clauses = []
    comments = {}
    with lzma.open(path, "rt", errors="ignore") as fh:
        for line in fh:
            if line.startswith("c C = "):
                comments["C"] = ast.literal_eval(line[6:])
                continue
            if line.startswith("c d = "):
                comments["d"] = ast.literal_eval(line[6:])
                continue
            if line.startswith(("c", "p")):
                continue
            lits = [int(x) for x in line.split()]
            if lits and lits[-1] == 0:
                lits.pop()
            if lits:
                clauses.append(lits)
    return clauses, comments


def parse_runs(clauses):
    front = len(clauses[0])
    selectors = [front + 2 * i for i in range(1, front + 1)]
    runs = []
    i = 1 + 4 * front
    while i < len(clauses):
        found = False
        for s in selectors:
            j = i
            ys = []
            while j + 1 < len(clauses) and len(clauses[j]) == 2 and len(clauses[j + 1]) == 2:
                a, b = clauses[j]
                c, d = clauses[j + 1]
                if a == -s and c == s and d == -b and b > front:
                    ys.append(b)
                    j += 2
                else:
                    break
            if ys:
                runs.append({"start": i, "end": j, "s": s, "ys": ys})
                i = j
                found = True
                break
        if not found:
            i += 1
    return front, selectors, runs


def first_occurrences(clauses):
    maxv = max(abs(l) for cl in clauses for l in cl)
    first = [10**12] * (maxv + 1)
    for idx, cl in enumerate(clauses):
        for lit in cl:
            v = abs(lit)
            if first[v] > idx:
                first[v] = idx
    return first


def vars_in_range(clauses, start, end):
    out = set()
    for cl in clauses[start:end]:
        for lit in cl:
            out.add(abs(lit))
    return out


def first_clause_with_var(clauses, start, end, pred):
    for idx in range(start, end):
        if any(pred(abs(lit)) for lit in clauses[idx]):
            return idx
    return None


def first_clause_index(clauses, start, end, var):
    for idx in range(start, end):
        if any(abs(lit) == var for lit in clauses[idx]):
            return idx
    return None


def propagate(slice_clauses, assumptions):
    val = {}

    def assign(lit):
        v = abs(lit)
        bit = lit > 0
        if v in val:
            return val[v] == bit
        val[v] = bit
        return True

    for lit in assumptions:
        if not assign(lit):
            return None
    changed = True
    while changed:
        changed = False
        for cl in slice_clauses:
            sat = False
            unassigned = 0
            last = 0
            for lit in cl:
                bit = val.get(abs(lit))
                if bit is None:
                    unassigned += 1
                    last = lit
                elif bit == (lit > 0):
                    sat = True
                    break
            if sat:
                continue
            if unassigned == 0:
                return None
            if unassigned == 1:
                if not assign(last):
                    return None
                changed = True
    return val


def group_target(group_clauses, lhs, rhs, cin):
    vs = sorted({abs(lit) for cl in group_clauses for lit in cl})
    if len(vs) > 10:
        return None
    pos = {v: i for i, v in enumerate(vs)}
    values = set()
    for mask in range(1 << len(vs)):
        ok = True
        for cl in group_clauses:
            if not any(((mask >> pos[abs(lit)]) & 1) == (lit > 0) for lit in cl):
                ok = False
                break
        if ok:
            def get(v):
                if v == 0:
                    return 0
                return (mask >> pos[v]) & 1

            values.add(get(lhs) ^ get(rhs) ^ get(cin))
    if len(values) == 1:
        return next(iter(values))
    return None


def sat_masks(chunk, vs):
    pos = {v: i for i, v in enumerate(vs)}
    out = []
    for mask in range(1 << len(vs)):
        ok = True
        for cl in chunk:
            if not any(((mask >> pos[abs(lit)]) & 1) == (lit > 0) for lit in cl):
                ok = False
                break
        if ok:
            out.append(mask)
    return out


def parse_parity(chunk):
    vs = sorted({abs(lit) for cl in chunk for lit in cl})
    if not vs or len(vs) > 4:
        return None
    sat = sat_masks(chunk, vs)
    for parity in (0, 1):
        want = [mask for mask in range(1 << len(vs)) if (mask.bit_count() & 1) == parity]
        if sat == want:
            return {"vars": vs, "parity": parity}
    return None


def parse_carry(chunk):
    vs = sorted({abs(lit) for cl in chunk for lit in cl})
    if not vs or len(vs) > 4:
        return None
    sat = sat_masks(chunk, vs)
    pos = {v: i for i, v in enumerate(vs)}

    def bit(mask, var):
        return (mask >> pos[var]) & 1

    for out in vs:
        ins = [v for v in vs if v != out]
        want = [
            mask for mask in range(1 << len(vs))
            if bit(mask, out) == 0
        ]
        if sat == want:
            return {"vars": vs, "out": out}
        if len(ins) == 2:
            want = [
                mask for mask in range(1 << len(vs))
                if bit(mask, out) == (bit(mask, ins[0]) & bit(mask, ins[1]))
            ]
            if sat == want:
                return {"vars": vs, "out": out}
        if len(ins) == 3:
            want = [
                mask for mask in range(1 << len(vs))
                if bit(mask, out) == (sum(bit(mask, v) for v in ins) >= 2)
            ]
            if sat == want:
                return {"vars": vs, "out": out}
    if len(vs) == 2:
        a, b = vs
        want = [mask for mask in range(4) if not (bit(mask, a) & bit(mask, b))]
        if sat == want:
            return {"vars": vs, "out": 0}
    if len(vs) == 3:
        want = [mask for mask in range(8) if sum(bit(mask, v) for v in vs) < 2]
        if sat == want:
            return {"vars": vs, "out": 0}
    return None


def scan_adder(clauses, start, end):
    groups = []
    p = start
    while p < end:
        ps = None
        plen = 0
        for length in (8, 4, 2):
            if p + length <= end:
                cand = parse_parity(clauses[p:p + length])
                if cand is not None:
                    ps = cand
                    plen = length
                    break
        if ps is None:
            raise ValueError(f"could not parse sum relation at clause {p}")
        p += plen
        cr = None
        clen = 0
        for length in (6, 4, 3, 2, 1):
            if p + length <= end:
                cand = parse_carry(clauses[p:p + length])
                if cand is not None:
                    cr = cand
                    clen = length
                    break
        if cr is None:
            cr = {"vars": [], "out": 0}
            clen = 0
        p += clen
        groups.append({
            "sum_vars": set(ps["vars"]),
            "parity": ps["parity"],
            "carry_vars": set(cr["vars"]),
            "cout": cr["out"],
        })
    return groups


def recover(clauses):
    front, selectors, runs = parse_runs(clauses)
    selector_index = {s: i for i, s in enumerate(selectors)}
    row_runs_grid = []
    cur = [None] * front
    last_idx = -1
    for run in runs:
        idx = selector_index[run["s"]]
        if idx <= last_idx:
            row_runs_grid.append(cur)
            cur = [None] * front
        cur[idx] = run
        last_idx = idx
    row_runs_grid.append(cur)
    rows = len(row_runs_grid)
    first = first_occurrences(clauses)
    selset = set(selectors)
    matrix = []
    targets = []

    for r in range(rows):
        row_runs = row_runs_grid[r]
        first_run = next(run for run in row_runs if run is not None)
        next_first = next((run for run in row_runs_grid[r + 1] if run is not None), None) if r + 1 < rows else None
        row_end = next_first["start"] if next_first is not None else len(clauses)
        acc_map = {}
        row = [0] * front
        first_term_map = {}
        target_bits = {}

        for k in range(1, front):
            run = row_runs[k]
            if run is None:
                row[k] = 0
                continue
            next_run = next((row_runs[t] for t in range(k + 1, front) if row_runs[t] is not None), None)
            block_end = next_run["start"] if next_run is not None else row_end
            prev_vars = set(acc_map.values())
            future_for_start = vars_in_range(clauses, block_end, row_end) if k + 1 < front else set()
            add_start = first_clause_with_var(
                clauses,
                run["end"],
                block_end,
                lambda v: v in prev_vars or v in future_for_start,
            )
            if add_start is None:
                raise ValueError("could not find adder consuming previous accumulator")

            product_vars = {
                v for v in vars_in_range(clauses, run["start"], add_start)
                if run["start"] <= first[v] < add_start
            }
            groups = scan_adder(clauses, add_start, block_end)
            future_vars = vars_in_range(clauses, block_end, row_end) if k + 1 < front else set()
            term_map = {}
            next_acc = {}
            prev_map = {}
            for pos, group in enumerate(groups):
                gvars = group["sum_vars"] | group["carry_vars"]
                prev = sorted(
                    [v for v in gvars if first[v] < run["start"] and v not in selset],
                    key=lambda v: first[v],
                )
                if prev:
                    prev_map[pos] = prev[0]
                rhs = sorted(gvars & product_vars, key=lambda v: first[v])
                if rhs:
                    term_map[pos] = rhs[0]
                if k + 1 < front:
                    out = sorted(group["sum_vars"] & future_vars, key=lambda v: first[v])
                    if out:
                        next_acc[pos] = out[0]
                else:
                    target_bits[pos] = group["parity"]

            if k == 1:
                first_term_map = prev_map

            val = propagate(clauses[run["start"]:add_start], [run["s"]])
            if val is None:
                raise ValueError("term propagation conflict")
            num = 0
            for pos, var in term_map.items():
                if val.get(var, False):
                    num |= 1 << pos
            row[k] = num
            acc_map = next_acc

        if row_runs[0] is not None:
            next_run = next((row_runs[t] for t in range(1, front) if row_runs[t] is not None), None)
            first_end = next_run["start"] if next_run is not None else row_end
            val = propagate(clauses[row_runs[0]["start"]:first_end], [row_runs[0]["s"]])
            if val is None:
                raise ValueError("first term propagation conflict")
            first_num = 0
            for pos, var in first_term_map.items():
                if val.get(var, False):
                    first_num |= 1 << pos
            row[0] = first_num
        matrix.append(row)
        targets.append(sum(bit << pos for pos, bit in target_bits.items()))

    return matrix, targets


def mitm_solve(C, d):
    m = len(C)
    n = len(C[0])
    n1 = n // 2
    left = []
    for mask in range(1 << n1):
        vec = [0] * m
        mm = mask
        while mm:
            lsb = mm & -mm
            j = lsb.bit_length() - 1
            for i in range(m):
                vec[i] += C[i][j]
            mm -= lsb
        left.append((tuple(vec), mask))
    left.sort()
    for mask2 in range(1 << (n - n1)):
        vec = [0] * m
        mm = mask2
        while mm:
            lsb = mm & -mm
            j2 = lsb.bit_length() - 1
            j = n1 + j2
            for i in range(m):
                vec[i] += C[i][j]
            mm -= lsb
        need = tuple(d[i] - vec[i] for i in range(m))
        idx = bisect.bisect_left(left, (need, -1))
        if idx < len(left) and left[idx][0] == need:
            return left[idx][1] | (mask2 << n1)
    return None


def main():
    for path in sys.argv[1:]:
        t0 = time.time()
        clauses, comments = read_dimacs_xz(path)
        C, d = recover(clauses)
        print(path, "dims", len(C), len(C[0]), "recover", round(time.time() - t0, 3))
        if "C" in comments:
            mismatches = sum(
                1 for i in range(len(C)) for j in range(len(C[0]))
                if C[i][j] != comments["C"][i][j]
            )
            print("  comment mismatches", mismatches, "target mismatches",
                  sum(a != b for a, b in zip(d, comments.get("d", []))))
        sol = mitm_solve(C, d)
        print("  solution", sol, "weight", None if sol is None else sol.bit_count())


if __name__ == "__main__":
    main()
