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

budget = float(sys.argv[1]) if len(sys.argv) > 1 else 5.0
limit = int(sys.argv[2]) if len(sys.argv) > 2 else 100
want = sys.argv[3] if len(sys.argv) > 3 else "sat"

with open("/workspace/baselines/kissat-sc2024.json") as f:
    rows = {r["name"]: r for r in json.load(f)["results"]}

files = []
for p in glob.glob("/data/train/*.cnf.xz"):
    row = rows.get(os.path.basename(p))
    if row and row.get("expected_result") == want:
        files.append(p)

random.Random(123).shuffle(files)
files = files[:limit]

solved = 0
by = {}
for i, p in enumerate(files, 1):
    name = os.path.basename(p)
    row = rows[name]
    m = re.search(r"-k(\d+)-", name)
    k = m.group(1) if m else "3"
    with open("/tmp/bench.cnf", "wb") as out:
        subprocess.run(["xz", "-d", "-c", p], stdout=out, check=True)
    t0 = time.time()
    r = subprocess.run(
        ["/workspace/solver/solver", "/tmp/bench.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])
    by[k][0] += ok
    by[k][1] += 1
    print(("SAT" if ok else "UNK"), f"{dt:.2f}", k, row["status"], name, flush=True)

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