"""Driver for the exhaustive C oracle (brute_force.c): for a given n, k, finds
Z(P(n,k)) by increasing cardinality s, in parallel over the lowest vertex id
m0 (2n independent jobs per s), with orbit reduction under D_n (P3).

Also contains a pure-Python reference enumeration (no symmetry reduction) used
to cross-validate the C program on tiny cases (P(5,2) and P(11,5) at s <= 5).

Output: results/brute_force_runs.jsonl (one record per (n,k,s)) and the
value table results/brute_force_values.json.
"""
import itertools
import json
import os
import subprocess
import sys
import time
from concurrent.futures import ThreadPoolExecutor

import common
from closure_oracle import closure_fast
from graph_constructor import adjacency_masks, edge_list_sha256, edges_formula

BIN = os.path.join(common.ROOT, "code", "brute_force")


def build_binary():
    src = os.path.join(common.ROOT, "code", "brute_force.c")
    if not os.path.exists(BIN) or os.path.getmtime(BIN) < os.path.getmtime(src):
        subprocess.run(["gcc", "-O2", "-o", BIN, src], check=True)


def run_one(n, k, s, m0):
    out = subprocess.run([BIN, str(n), str(k), str(s), str(m0)], capture_output=True, text=True, check=True).stdout
    d = {}
    for tok in out.split():
        key, _, val = tok.partition("=")
        d[key] = val
    return d


def count_size(n, k, s, workers=16):
    """Enumerates all s-subsets (split over m0). Returns (enumerated, tested, forcing, witness or None, wall)."""
    t0 = time.time()
    with ThreadPoolExecutor(max_workers=workers) as ex:
        parts = list(ex.map(lambda m0: run_one(n, k, s, m0), range(2 * n)))
    enumerated = sum(int(p["enumerated"]) for p in parts)
    tested = sum(int(p["tested"]) for p in parts)
    forcing = sum(int(p["forcing"]) for p in parts)
    witness = next((p["witness"] for p in parts if "witness" in p), None)
    return enumerated, tested, forcing, witness, time.time() - t0


def python_reference(n, k, s):
    """No symmetry reduction: counts all forcing s-subsets (for tiny cases only)."""
    adj = adjacency_masks(n, k)
    N = 2 * n
    full = (1 << N) - 1
    cnt = 0
    for comb in itertools.combinations(range(N), s):
        m = 0
        for w in comb:
            m |= 1 << w
        if closure_fast(adj, m) == full:
            cnt += 1
    return cnt


def find_Z(n, k, s_start=1, s_max=12, time_limit=1800, workers=16):
    """Returns (Z or None, records). Each size s < Z is exhausted (forcing == 0)."""
    records = []
    t_total = time.time()
    for s in range(s_start, s_max + 1):
        if time.time() - t_total > time_limit:
            return None, records
        if s >= 12 and n >= 12:
            # Z <= 12 by the explicit construction P7 (replayed in verify_upper_bounds.py);
            # all sizes <= 11 have been exhausted above, so Z = 12 without enumerating 12-sets.
            rec = {"n": n, "k": k, "s": s, "enumerated_subsets": 0, "orbit_representatives_tested": 0,
                   "forcing_representatives": 1, "witness": ",".join("u%d" % i for i in range(12)),
                   "witness_source": "P7 construction (not enumerated)", "wall_seconds": 0.0,
                   "utc": common.utcnow(), "graph_sha256": edge_list_sha256(edges_formula(n, k))}
            records.append(rec)
            common.append_jsonl(os.path.join(common.RESULTS, "brute_force_runs.jsonl"), rec)
            return s, records
        enumerated, tested, forcing, witness, wall = count_size(n, k, s, workers)
        rec = {"n": n, "k": k, "s": s, "enumerated_subsets": enumerated, "orbit_representatives_tested": tested,
               "forcing_representatives": forcing, "witness": witness, "wall_seconds": round(wall, 3),
               "utc": common.utcnow(), "graph_sha256": edge_list_sha256(edges_formula(n, k))}
        records.append(rec)
        common.append_jsonl(os.path.join(common.RESULTS, "brute_force_runs.jsonl"), rec)
        print(json.dumps(rec), flush=True)
        if forcing:
            return s, records
    return None, records


if __name__ == "__main__":
    build_binary()
    # cross-validation of the C oracle against the Python reference on tiny cases
    checks = []
    for (n, k, s) in [(5, 2, 4), (5, 2, 5), (11, 5, 5)]:
        ref = python_reference(n, k, s)
        enumerated, tested, forcing, witness, wall = count_size(n, k, s)
        # orbit reduction: forcing representatives > 0 iff reference count > 0
        checks.append({"n": n, "k": k, "s": s, "python_forcing_subsets": ref, "c_forcing_representatives": forcing,
                       "agree": (ref > 0) == (forcing > 0)})
        print(checks[-1])
    assert all(c["agree"] for c in checks)
    with open(os.path.join(common.RESULTS, "brute_force_crosscheck.json"), "w") as f:
        json.dump(checks, f, indent=1)
    ns = [int(a) for a in sys.argv[1:]] or [11, 12, 13, 14]
    k = 5
    values = {}
    if os.path.exists(os.path.join(common.RESULTS, "brute_force_values.json")):
        values = json.load(open(os.path.join(common.RESULTS, "brute_force_values.json")))
    for n in ns:
        Z, recs = find_Z(n, k, s_start=max(1, (4 if n >= 11 else 1)), time_limit=float(os.environ.get("BF_TIME_LIMIT", "1800")))
        values[str(n)] = {"Z": Z, "sizes_exhausted": [r["s"] for r in recs if r["forcing_representatives"] == 0],
                          "witness": recs[-1]["witness"] if recs and recs[-1]["forcing_representatives"] else None,
                          "total_wall_seconds": round(sum(r["wall_seconds"] for r in recs), 3)}
        print("n=%d Z=%s" % (n, Z), flush=True)
        with open(os.path.join(common.RESULTS, "brute_force_values.json"), "w") as f:
            json.dump(values, f, indent=1)