#!/usr/bin/env python3
import glob
import os
import re
import subprocess
import sys
import time

files = sorted(glob.glob("/data/train/*.SATISFIABLE.cnf.xz"))
budget = float(sys.argv[1]) if len(sys.argv) > 1 else 2.0
limit = int(sys.argv[2]) if len(sys.argv) > 2 else len(files)

solved = 0
by = {}
for total, f in enumerate(files[:limit], 1):
    name = os.path.basename(f)
    m = re.search(r"-k(\d+)-", name)
    k = m.group(1) if m else "3"
    t0 = time.time()
    cnf = "/tmp/bench.cnf"
    with open(cnf, "wb") as out:
        subprocess.run(["xz", "-d", "-c", f], stdout=out, check=True)
    r = subprocess.run(
        ["/workspace/solver/solver", cnf, "/tmp/out.pbp"],
        env={**os.environ, "LYMPHOSAT_SOLVER_SECONDS": str(budget)},
        stdout=subprocess.DEVNULL,
        stderr=subprocess.DEVNULL,
    )
    dt = time.time() - t0
    ok = r.returncode == 10
    solved += ok
    by.setdefault(k, [0, 0, 0.0])
    by[k][0] += ok
    by[k][1] += 1
    by[k][2] += dt
    print(("SAT" if ok else "UNK"), f"{dt:.2f}", k, name, flush=True)

print("TOTAL", solved, min(limit, len(files)), by)
