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


def open_maybe_xz(path):
    if path.endswith(".xz"):
        return lzma.open(path, "rt", encoding="ascii", errors="ignore")
    return open(path, "rt", encoding="ascii", errors="ignore")


def parse_dimacs(path):
    nvars = nclauses = None
    clauses = []
    cur = []
    with open_maybe_xz(path) as f:
        for line in f:
            if not line or line[0] in "c\n":
                continue
            if line[0] == "p":
                parts = line.split()
                nvars = int(parts[2])
                nclauses = int(parts[3])
                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 nvars, nclauses, clauses


def find_exact_one_groups(clauses, max_group=10000):
    neg_pairs = set()
    pos_clauses = []
    for c in clauses:
        if len(c) == 2 and c[0] < 0 and c[1] < 0:
            a, b = -c[0], -c[1]
            if a > b:
                a, b = b, a
            neg_pairs.add((a, b))
        elif len(c) >= 2 and all(x > 0 for x in c):
            if len(c) <= max_group:
                pos_clauses.append(tuple(sorted(c)))
    groups = []
    for c in pos_clauses:
        ok = True
        for a, b in itertools.combinations(c, 2):
            if (a, b) not in neg_pairs:
                ok = False
                break
        if ok:
            groups.append(c)
    return groups


def binary_components(nvars, clauses):
    parent = list(range(2 * nvars + 2))

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

    def find(x):
        while parent[x] != x:
            parent[x] = parent[parent[x]]
            x = parent[x]
        return x

    def union(a, b):
        ra, rb = find(a), find(b)
        if ra != rb:
            parent[rb] = ra

    for c in clauses:
        if len(c) == 2:
            a, b = c
            union(idx(a), idx(b))
            union(idx(-a), idx(-b))

    comps = collections.Counter(find(idx(x)) for v in range(1, nvars + 1) for x in (v, -v))
    return comps


def summarize(path, verbose=False):
    nvars, nclauses, clauses = parse_dimacs(path)
    sizes = collections.Counter(len(c) for c in clauses)
    pos = neg = 0
    units = []
    pos_all = neg_all = horn = dual_horn = 0
    for c in clauses:
        p = sum(1 for x in c if x > 0)
        q = len(c) - p
        pos += p
        neg += q
        if len(c) == 1:
            units.append(c[0])
        if p == len(c):
            pos_all += 1
        if q == len(c):
            neg_all += 1
        if p <= 1:
            horn += 1
        if q <= 1:
            dual_horn += 1
    deg = [0] * (nvars + 1)
    posdeg = [0] * (nvars + 1)
    negdeg = [0] * (nvars + 1)
    for c in clauses:
        for lit in c:
            v = abs(lit)
            deg[v] += 1
            if lit > 0:
                posdeg[v] += 1
            else:
                negdeg[v] += 1
    degs = sorted(deg[1:])
    groups = find_exact_one_groups(clauses)
    comps = binary_components(nvars, clauses)
    comp_sizes = sorted(comps.values(), reverse=True)
    print(os.path.basename(path))
    print("  vars", nvars, "clauses", len(clauses), "decl", nclauses, "lits", pos + neg)
    print("  sizes", " ".join(f"{k}:{v}" for k, v in sorted(sizes.items())[:20]))
    print("  polarity pos", pos, "neg", neg, "pos_frac", f"{pos/(pos+neg):.3f}")
    print("  allpos", pos_all, "allneg", neg_all, "horn", horn, "dual_horn", dual_horn, "units", len(units))
    print("  degree min/med/max", degs[0], degs[len(degs)//2], degs[-1])
    print("  bin comps", len(comp_sizes), "largest", comp_sizes[:8])
    if groups:
        gsz = collections.Counter(len(g) for g in groups)
        covered = len(set(v for g in groups for v in g))
        print("  exact1 groups", len(groups), "sizes", dict(sorted(gsz.items())), "covered_vars", covered)
        if verbose:
            for g in groups[:10]:
                print("    group", len(g), g[:30])
    if verbose:
        print("  units", units[:50])


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("paths", nargs="*")
    ap.add_argument("-v", "--verbose", action="store_true")
    args = ap.parse_args()
    paths = args.paths or sorted(glob.glob("/data/train/*.cnf.xz"))
    for p in paths:
        summarize(p, args.verbose)


if __name__ == "__main__":
    main()
