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


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


def lit_eval(lit, vals):
    v = vals[abs(lit)]
    return v if lit > 0 else not v


def clause_ok(clause, vals):
    return any(lit_eval(l, vals) for l in clause)


def classify_binary(z, clauses):
    defs = [c for c in clauses if len(c) > 1]
    if len(defs) != 3:
        return None
    small = sorted({abs(l) for c in defs for l in c if abs(l) != z})
    if len(small) != 2:
        return None
    a, b = small
    # Check whether defs determine z for every a,b assignment.
    outs = []
    for av, bv in itertools.product([False, True], repeat=2):
        possible = []
        for zv in [False, True]:
            vals = {a: av, b: bv, z: zv}
            if all(clause_ok(c, vals) for c in defs):
                possible.append(zv)
        if len(possible) != 1:
            return None
        outs.append(possible[0])
    # order: 00,01,10,11
    table = tuple(outs)
    names = {
        (False, False, False, True): "and",
        (False, True, True, True): "or",
        (True, True, True, False): "nand",
        (True, False, False, False): "nor",
        (False, True, True, False): "xor",
        (True, False, False, True): "xnor",
        (False, False, True, False): "a_and_not_b",
        (False, True, False, False): "not_a_and_b",
        (True, True, False, True): "b_or_not_a",
        (True, False, True, True): "a_or_not_b",
        (False, False, True, True): "a",
        (False, True, False, True): "b",
        (True, True, False, False): "not_a",
        (True, False, True, False): "not_b",
    }
    return names.get(table, f"table{''.join('1' if x else '0' for x in table)}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("path")
    ap.add_argument("--limit-clauses", type=int, default=0)
    ap.add_argument("--max-var", type=int, default=0)
    ap.add_argument("--show", type=int, default=40)
    args = ap.parse_args()

    groups = collections.defaultdict(list)
    nvars = nclauses = seen = 0
    with open_text(args.path) as f:
        for line in f:
            if not line or line[0] == "c":
                continue
            if line[0] == "p":
                p = line.split()
                nvars, nclauses = int(p[2]), int(p[3])
                continue
            vals = [int(x) for x in line.split()]
            if vals and vals[-1] == 0:
                vals.pop()
            if not vals:
                continue
            seen += 1
            z = max(abs(x) for x in vals)
            if not args.max_var or z <= args.max_var:
                groups[z].append(tuple(vals))
            if args.limit_clauses and seen >= args.limit_clauses:
                break
            if args.max_var and z > args.max_var + 1000:
                break

    counts = collections.Counter()
    arity_counts = collections.Counter()
    unit_polarity = collections.Counter()
    examples = []
    for z in sorted(groups):
        cs = groups[z]
        units = [c[0] for c in cs if len(c) == 1]
        if units:
            for u in units:
                unit_polarity["pos" if u > 0 else "neg"] += 1
        kind = classify_binary(z, cs)
        if kind:
            counts[kind] += 1
            arity_counts["binary"] += 1
        else:
            key = tuple(sorted(len(c) for c in cs))
            counts[f"shape{key}"] += 1
            arity_counts["other"] += 1
        if len(examples) < args.show:
            examples.append((z, kind or f"shape{tuple(sorted(len(c) for c in cs))}", cs[:6]))

    print(f"{os.path.basename(args.path)} vars={nvars} clauses={nclauses} scanned={seen} groups={len(groups)}")
    print("counts", counts)
    print("arity", arity_counts, "units", unit_polarity)
    for e in examples:
        print("example", e)


if __name__ == "__main__":
    main()
