#!/usr/bin/env python3
"""Fast-ish streaming CNF profile for the training split.

This intentionally does not implement a SAT solver.  It records structural
signals that are cheap to compute from xz-compressed DIMACS files.
"""

from __future__ import annotations

import json
import os
import subprocess
import sys
import time
from collections import Counter


def iter_tokens(path: str):
    proc = subprocess.Popen(["xz", "-d", "-c", path], stdout=subprocess.PIPE)
    assert proc.stdout is not None
    try:
        for raw in proc.stdout:
            if raw.startswith(b"c"):
                continue
            for tok in raw.split():
                yield tok
    finally:
        proc.stdout.close()
        proc.wait()


def profile(path: str, max_clauses: int | None = None):
    t0 = time.time()
    toks = iter_tokens(path)
    nvars = nclauses = None
    for tok in toks:
        if tok == b"p":
            fmt = next(toks)
            if fmt != b"cnf":
                raise ValueError(f"{path}: not cnf")
            nvars = int(next(toks))
            nclauses = int(next(toks))
            break
    if nvars is None:
        raise ValueError(f"{path}: no header")

    size_counts: Counter[int] = Counter()
    pos = neg = 0
    low_unit = 0
    units = []
    clause = []
    read = 0
    max_width = 0
    first_non_unit = None

    for tok in toks:
        lit = int(tok)
        if lit:
            clause.append(lit)
            continue
        read += 1
        k = len(clause)
        size_counts[k] += 1
        if k > max_width:
            max_width = k
        if k == 1:
            units.append(clause[0])
            if abs(clause[0]) < 1024:
                low_unit += 1
        elif first_non_unit is None:
            first_non_unit = read
        for lit in clause:
            if lit > 0:
                pos += 1
            else:
                neg += 1
        clause.clear()
        if max_clauses is not None and read >= max_clauses:
            break

    unit_set = set()
    contradictory_units = False
    dup_units = 0
    for lit in units:
        if -lit in unit_set:
            contradictory_units = True
        if lit in unit_set:
            dup_units += 1
        unit_set.add(lit)

    return {
        "file": os.path.basename(path),
        "nvars": nvars,
        "nclauses": nclauses,
        "read_clauses": read,
        "size_bytes": os.path.getsize(path),
        "sizes": dict(sorted(size_counts.items())),
        "units": len(units),
        "distinct_units": len(unit_set),
        "dup_units": dup_units,
        "contradictory_units": contradictory_units,
        "low_unit": low_unit,
        "first_non_unit": first_non_unit,
        "pos_lits": pos,
        "neg_lits": neg,
        "max_width": max_width,
        "seconds": round(time.time() - t0, 3),
    }


def main(argv: list[str]) -> int:
    max_clauses = None
    paths = []
    for arg in argv[1:]:
        if arg.startswith("--max-clauses="):
            max_clauses = int(arg.split("=", 1)[1])
        else:
            paths.append(arg)
    if not paths:
        paths = sorted(
            os.path.join("/data/train", p)
            for p in os.listdir("/data/train")
            if p.endswith(".cnf.xz")
        )
    for path in paths:
        try:
            print(json.dumps(profile(path, max_clauses=max_clauses), sort_keys=True), flush=True)
        except Exception as exc:
            print(json.dumps({"file": os.path.basename(path), "error": str(exc)}), flush=True)
    return 0


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