#!/usr/bin/env python3
import argparse
import collections
import lzma
import os
import sys


def read_dimacs(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] in "c%\n":
                continue
            if line.startswith("p"):
                p = line.split()
                nvars = int(p[2])
                nclauses = int(p[3])
                continue
            lits = [int(x) for x in line.split()]
            if lits and lits[-1] == 0:
                lits.pop()
            if lits:
                clauses.append(lits)
    return nvars, nclauses, clauses


def unit_run(units):
    vals = sorted(abs(u) for u in units)
    best = []
    cur = []
    for v in vals:
        if not cur or v == cur[-1] + 1:
            cur.append(v)
        else:
            if len(cur) > len(best):
                best = cur
            cur = [v]
    if len(cur) > len(best):
        best = cur
    return best


def infer_layout(clauses):
    first = next(c for c in clauses if c)
    units = [c[0] for c in clauses if len(c) == 1]
    run = unit_run(units)
    if len(run) not in (7, 8, 10):
        raise ValueError(f"could not infer counter width from units: run={run}")
    width = len(run)
    first_min = min(abs(x) for x in first)
    block_start = first_min if width == 8 else first_min - 2
    base = block_start - 1
    fixed_block = (run[0] - block_start) // width
    fixed_bits = [None] * width
    for lit in units:
        v = abs(lit)
        if run[0] <= v <= run[-1]:
            fixed_bits[v - run[0]] = 1 if lit > 0 else 0
    low_units = [u for u in units if abs(u) < block_start]
    return {
        "width": width,
        "block_start": block_start,
        "base": base,
        "fixed_block": fixed_block,
        "fixed_bits": fixed_bits,
        "low_units": low_units,
    }


def block_of(v, layout):
    return (v - layout["block_start"]) // layout["width"]


def off_of(v, layout):
    return (v - layout["block_start"]) % layout["width"]


def infer_edge_from_guard(guarded, layout):
    width = layout["width"]
    for i in range(max(0, len(guarded) - 3)):
        chunk = guarded[i : i + 4]
        hs = sorted({abs(l) for c in chunk for l in c if abs(l) >= layout["block_start"]})
        if len(hs) != 3:
            continue
        by_block = collections.defaultdict(list)
        for v in hs:
            by_block[block_of(v, layout)].append(off_of(v, layout))
        for b, offs in by_block.items():
            if sorted(offs) == [0, 1]:
                others = [bb for bb in by_block if bb != b]
                if len(others) == 1:
                    return b, others[0]
    # Width 8/105 starts with a 3-cycle factor; its first relation only has
    # offsets 0/1 but the source/destination roles are reversed from width 7.
    if width == 8 and len(guarded) >= 4:
        chunk = guarded[:4]
        hs = sorted({abs(l) for c in chunk for l in c if abs(l) >= layout["block_start"]})
        if len(hs) == 3:
            by_block = collections.defaultdict(list)
            for v in hs:
                by_block[block_of(v, layout)].append(off_of(v, layout))
            for b, offs in by_block.items():
                if sorted(offs) == [0, 1]:
                    others = [bb for bb in by_block if bb != b]
                    if len(others) == 1:
                        # In the 3-cycle factor d0 = s1, d1 = s0 xor s1.
                        return b, others[0]
    return None


def extract(path):
    nvars, nclauses, clauses = read_dimacs(path)
    layout = infer_layout(clauses)
    base = layout["base"]
    guarded = collections.defaultdict(list)
    for c in clauses:
        lows = [-l for l in c if l < 0 and -l <= base]
        if len(lows) == 1 and any(abs(l) >= layout["block_start"] for l in c):
            e = lows[0]
            if len(guarded[e]) < 40:
                guarded[e].append(c)
    edges = {}
    for e in range(1, base + 1):
        r = infer_edge_from_guard(guarded[e], layout)
        if r is not None:
            edges[e] = r
    # Special forced reverse arcs are deliberately not encoded as increments.
    # They are paired with the preceding/following normal directed edge.
    for e in range(1, base + 1, 2):
        if e in edges and e + 1 not in edges:
            s, d = edges[e]
            edges[e + 1] = (d, s)
        elif e + 1 in edges and e not in edges:
            s, d = edges[e + 1]
            edges[e] = (d, s)
    groups = []
    for c in clauses:
        if c and all(l > 0 for l in c) and any(l <= base for l in c) and all(l <= base for l in c):
            groups.append(c)
    max_block = -1
    for s, d in edges.values():
        max_block = max(max_block, s, d)
    n = max_block + 1
    undirected = set()
    for e in range(1, base + 1, 2):
        if e in edges and e + 1 in edges:
            a, b = edges[e]
            if a != b:
                undirected.add((min(a, b), max(a, b)))
    return nvars, nclauses, clauses, layout, edges, groups, n, sorted(undirected)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("cnf")
    args = ap.parse_args()
    nvars, nclauses, clauses, layout, edges, groups, n, undirected = extract(args.cnf)
    deg = [0] * n
    for a, b in undirected:
        deg[a] += 1
        deg[b] += 1
    print(os.path.basename(args.cnf))
    print("vars", nvars, "clauses", nclauses)
    print("layout", layout)
    print("blocks", n, "directed", len(edges), "undirected", len(undirected), "groups", len(groups))
    print("degree_hist", dict(sorted(collections.Counter(deg).items())))
    print("low_units", layout["low_units"])
    if layout["low_units"]:
        for u in layout["low_units"]:
            print("unit_low", u, "edge", edges.get(abs(u)))


if __name__ == "__main__":
    main()
