#!/usr/bin/env python3
"""Profile the GF(3)-SAT CNF encoding used by the training split.

This script is analysis-only.  It reads xz-compressed DIMACS files, groups the
4-literal nogoods by one-hot ternary variable scope, and checks whether each
54-clause block is a single linear equation over GF(3).
"""

from __future__ import annotations

import collections
import glob
import itertools
import lzma
import os
import sys


def read_clauses(path: str):
    nvars = None
    clauses: list[list[int]] = []
    with lzma.open(path, "rt", encoding="ascii", errors="ignore") as f:
        for line in f:
            if not line.strip() or line.startswith("c"):
                continue
            if line.startswith("p"):
                parts = line.split()
                nvars = int(parts[2])
                continue
            clauses.append([int(x) for x in line.split() if x != "0"])
    if nvars is None:
        raise ValueError(f"{path}: missing p-line")
    return nvars, clauses


def fit_linear(forbidden):
    forbidden = set(forbidden)
    for coeffs in itertools.product((1, 2), repeat=4):
        if coeffs[0] != 1:
            continue
        for rhs in range(3):
            ok = True
            for vals in itertools.product(range(3), repeat=4):
                sat = sum(c * v for c, v in zip(coeffs, vals)) % 3 == rhs
                if (vals in forbidden) == sat:
                    ok = False
                    break
            if ok:
                return coeffs, rhs
    return None


def profile(path: str):
    nvars, clauses = read_clauses(path)
    sizes = collections.Counter(len(c) for c in clauses)
    groups = collections.defaultdict(list)
    for c in clauses:
        if len(c) != 4:
            continue
        scope_vals = sorted(((abs(lit) - 1) // 3, (abs(lit) - 1) % 3) for lit in c)
        scope = tuple(g for g, _ in scope_vals)
        vals = tuple(v for _, v in scope_vals)
        groups[scope].append(vals)

    fits = collections.Counter()
    bad = 0
    for vals in groups.values():
        found = fit_linear(vals) if len(vals) == 54 else None
        if found is None:
            bad += 1
        else:
            fits[found] += 1

    print(os.path.basename(path))
    print(f"  vars={nvars} clauses={len(clauses)} sizes={dict(sorted(sizes.items()))}")
    print(f"  one_hot_triples={sizes[3]} equation_blocks={len(groups)} bad_blocks={bad}")
    print(f"  fit_classes={len(fits)}")


def main(argv: list[str]) -> int:
    paths = argv[1:] or sorted(glob.glob("/data/train/*.cnf.xz"))
    for path in paths:
        profile(path)
    return 0


if __name__ == "__main__":
    raise SystemExit(main(sys.argv))
