"""Exact-value driver (T5/T6). For each n:
  * upper bound: smallest b at which the fort-cover search finds a witness
    quickly (descending from min(12, published bound)); the witness is replayed
    by closure_oracle and stored in results/upper_bound_witnesses.jsonl with the
    full chronology;
  * SAT lower bound: fort-cover UNSAT at b = Z_upper - 1 within the time limit
    (static final CNF re-solved with proof logging, DRAT checked by drat-trim);
    and, independently, the round-based encoding (P4) at s = Z_upper - 1 within
    its own time limit;
  * brute-force lower bound is read from results/brute_force_values.json when
    present.
Every run is recorded in results/solver_runs.jsonl.

usage: python3 code/run_exact.py N [N ...]   (env: FORT_LIMIT, ROUND_LIMIT seconds)
"""
import json
import os
import sys
import time

import common
from closure_oracle import closure, replay
from graph_constructor import adjacency, edge_list_sha256, edges_formula, label
from sat_or_ilp_solver import fort_cover, round_sat, solve_with_proof

K = 5
FORT_LIMIT = float(os.environ.get("FORT_LIMIT", "900"))
ROUND_LIMIT = float(os.environ.get("ROUND_LIMIT", "600"))


def record(rec):
    rec["utc"] = common.utcnow()
    rec.update(common.versions())
    common.append_jsonl(os.path.join(common.RESULTS, "solver_runs.jsonl"), rec)
    print(json.dumps(rec)[:300], flush=True)


def upper_bound_search(n, start):
    """Descend b from start while a witness is found within a short limit."""
    adj = adjacency(n, K)
    best = None
    b = start
    while b >= 1:
        r = fort_cover(n, K, b, time_limit=FORT_LIMIT)
        record({"n": n, "k": K, "method": "fort_cover", "bound": b, "sat": r["sat"], "iterations": r["iterations"],
                "forts": len(r["forts"]), "wall_seconds": r["wall_seconds"], "timeout": r["sat"] is None})
        if r["sat"]:
            S = r["witness"]
            final, chron = closure(adj, S)
            assert len(final) == 2 * n
            ok, fin, bad = replay(adj, S, chron)
            assert ok and len(fin) == 2 * n
            best = {"n": n, "k": K, "size": len(S), "S": [label(w, n) for w in S],
                    "chronology": [[label(a, n), label(c, n)] for a, c in chron],
                    "graph_sha256": edge_list_sha256(edges_formula(n, K)), "source": "fort_cover", "utc": common.utcnow()}
            common.append_jsonl(os.path.join(common.RESULTS, "upper_bound_witnesses.jsonl"), best)
            b = len(S) - 1
            continue
        return best, b, r  # r is the UNSAT (or timeout) result at b
    return best, b, None


def lower_bound_certificate(n, b, r):
    """r: fort_cover result with sat False at bound b. Re-solve the static CNF with proof logging."""
    cert_dir = os.path.join(common.CERTS, "fortcover")
    os.makedirs(cert_dir, exist_ok=True)
    clauses = r["base_clauses"] + r["fort_clauses"]
    cnf_path = os.path.join(cert_dir, "P%d_5_atmost%d.cnf" % (n, b))
    proof_path = os.path.join(cert_dir, "P%d_5_atmost%d.drat" % (n, b))
    forts_path = os.path.join(cert_dir, "P%d_5_atmost%d.forts.json" % (n, b))
    with open(forts_path, "w") as f:
        json.dump({"n": n, "k": K, "bound": b, "svar": r["svar"], "normalised": r["normalised"], "symmetric": r["symmetric"],
                   "forts": r["forts"], "forts_labels": [[label(w, n) for w in F] for F in r["forts"]]}, f)
    out = solve_with_proof(clauses, r["nv"], "cadical195", cnf_path, proof_path)
    out.update({"n": n, "k": K, "bound": b, "method": "fort_cover_static_proof", "forts_path": os.path.relpath(forts_path, common.ROOT),
                "forts_sha256": common.sha256_file(forts_path), "num_forts": len(r["forts"])})
    record(out)
    common.append_jsonl(os.path.join(common.RESULTS, "lower_bound_certificates.jsonl"), out)
    return out


if __name__ == "__main__":
    ns = [int(a) for a in sys.argv[1:]]
    published = {11: 6, 12: 8, 13: 10, 14: 11, 16: 9}
    for n in ns:
        start = min(12, published.get(n, 12))
        t0 = time.time()
        best, b, r = upper_bound_search(n, start)
        Zu = best["size"] if best else start
        lower = None
        if r is not None and r["sat"] is False:
            cert = lower_bound_certificate(n, b, r)
            lower = b + 1
        # second path: round-based encoding at s = Zu - 1
        # SKIP_ROUND=1 (reproduce.sh default mode) omits the round-based path, whose runs at n >= 18 only time out
        rr = round_sat(n, K, Zu - 1, time_limit=ROUND_LIMIT) if (n <= 30 and os.environ.get("SKIP_ROUND") != "1") else None
        if rr is not None:
            record({"n": n, "k": K, "method": "round_sat", "solver": rr["solver"], "bound": Zu - 1, "sat": rr["sat"], "T": rr["T"],
                    "num_vars": rr["num_vars"], "num_clauses": rr["num_clauses"], "wall_seconds": rr["solve_seconds"], "timeout": rr["timeout"]})
        print("n=%d: upper=%d (witness size), fort-cover lower=%s, round_sat(%d)=%s, total %.1fs" % (
            n, Zu, lower, Zu - 1, rr["sat"] if rr else None, time.time() - t0), flush=True)