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


def parse(path):
    n = 0
    clauses = []
    cur = []
    with lzma.open(path, "rt", encoding="ascii", errors="ignore") if path.endswith(".xz") else open(path, "rt", encoding="ascii", errors="ignore") as f:
        for line in f:
            if not line or line[0] in "c\n":
                continue
            if line[0] == "p":
                n = int(line.split()[2])
                continue
            for tok in line.split():
                lit = int(tok)
                if lit == 0:
                    if cur:
                        clauses.append(tuple(cur))
                        cur = []
                else:
                    cur.append(lit)
    if cur:
        clauses.append(tuple(cur))
    return n, clauses


def recover(n, clauses):
    sizes = collections.Counter(len(c) for c in clauses if len(c) > 2)
    if not sizes:
        return None
    L = sizes.most_common(1)[0][0]
    if L % 2:
        return None
    b = L // 2
    count_by_set = collections.Counter()
    pair_counts = collections.Counter()
    for c in clauses:
        if len(c) == L:
            s = tuple(sorted({abs(x) for x in c}))
            if len(s) == L:
                count_by_set[s] += 1
                for a, z in itertools.combinations(s, 2):
                    pair_counts[(a, z)] += 1
    sets = list(count_by_set)
    adj = collections.defaultdict(set)
    # In this family, same-bit-vector pairs repeat with many partner supports.
    # Cross pairs repeat only as often as one binary constraint's forbidden table.
    vals = sorted(pair_counts.values())
    if not vals:
        return None
    # Pick a conservative cutoff: top co-occurrences, but never below twice the
    # median positive pair count.
    cutoff = max(vals[len(vals) // 2] * 2 + 1, vals[int(len(vals) * 0.90)])
    for (a, z), ct in pair_counts.items():
        if ct >= cutoff:
            adj[a].add(z)
            adj[z].add(a)
    groups_set = set()
    for v in range(1, n + 1):
        nbrs = sorted(adj.get(v, ()), key=lambda u: pair_counts[(min(u, v), max(u, v))], reverse=True)
        cand = tuple(sorted([v] + nbrs[: b - 1]))
        if len(cand) == b:
            groups_set.add(cand)
    changed = True
    while changed:
        changed = False
        for s in sets:
            ss = set(s)
            for comb in itertools.combinations(s, b):
                if comb in groups_set:
                    rest = tuple(sorted(ss - set(comb)))
                    if len(rest) == b and rest not in groups_set:
                        groups_set.add(rest)
                        changed = True
    # Some low-degree bit-vectors are not found by co-occurrence. Recover them
    # from intersections of supports sharing a variable: if two 2b-supports
    # share exactly b variables, that intersection is one complete bit-vector.
    by_var = collections.defaultdict(list)
    set_objs = [set(s) for s in sets]
    for i, s in enumerate(sets):
        for v in s:
            by_var[v].append(i)
    for ids in by_var.values():
        if len(ids) > 2000:
            # Avoid pathological dense cases; co-occurrence already handles
            # high-degree variables well.
            continue
        for i_pos in range(len(ids)):
            si = set_objs[ids[i_pos]]
            for j_pos in range(i_pos + 1, len(ids)):
                inter = si & set_objs[ids[j_pos]]
                if len(inter) == b:
                    groups_set.add(tuple(sorted(inter)))
    changed = True
    while changed:
        changed = False
        for s in sets:
            ss = set(s)
            for comb in itertools.combinations(s, b):
                if comb in groups_set:
                    rest = tuple(sorted(ss - set(comb)))
                    if len(rest) == b and rest not in groups_set:
                        groups_set.add(rest)
                        changed = True
    groups = sorted(groups_set)
    used = collections.Counter(v for g in groups for v in g)
    return L // 2, L, len(sets), count_by_set, groups, used, cutoff


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("paths", nargs="*")
    args = ap.parse_args()
    for p in args.paths or sorted(glob.glob("/data/train/*.cnf.xz")):
        n, clauses = parse(p)
        r = recover(n, clauses)
        if not r:
            continue
        b, L, nsets, count_by_set, groups, used, cutoff = r
        csets = collections.Counter(count_by_set.values())
        print(os.path.basename(p), "n", n, "b", b, "L", L, "unique_pair_sets", nsets)
        print("  repeated-varset clause counts", dict(sorted(csets.items())[:10]), "...")
        print("  cutoff", cutoff, "recovered groups", len(groups), "used_vars", len(used), "multiused", sum(1 for x in used.values() if x != 1))
        print("  sample groups", groups[:5])


if __name__ == "__main__":
    main()
