"""Large-sample instance pool, so q_s is a distribution estimate and not a
coverage checklist.

`instances_for_stage` stops as soon as every (cell, candidate) pair has been
seen once (`while uncovered and ...`, plus require_new rejecting redundant
draws), which yields ~5 grids per puzzle and drops to 1.4 by stage 11. The
resulting per-cell histogram is far too small to estimate an entropy: measured
H(q_hat) climbs from 0.70 at n=3 to 0.99 at n=7 against log|S| = 1.10 and is
still rising, so the headroom log|S| - H(q_hat) that the promotion gate spends
is mostly finite-sample bias.

This samples the same generator in its large-n limit. With `uncovered` empty,
`propose` falls through to an independent uniform draw per cell (see its
docstring), still preceded by the confined-disjunction pass and followed by
`repair`, so accepted grids are distributed as "uniform over S subject to the
local confinements" -- exactly what the pool converges to once coverage is
met, and what the paper means by the stage-s superposition.

Writes per-(puzzle, stage) digit counts plus a handful of whole grids, the
latter to be used as teacher-forcing prefixes so the model is never shown the
unique solution while its candidate distribution is being scored.
"""

import argparse
import os
import sys

import numpy as np

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from superposition_instances import (build_dep_index, dependencies_from_mask,
                                     propose, repair, split_cells,
                                     _violations)

EMPTY = frozenset()


def sample_stage(mask, n_draws, max_confine=3, seed=0, max_repair=200,
                 keep_grids=8):
    """n_draws accepted grids for one stage's candidate mask.

    Returns (counts (81,9) int32, grids (keep_grids,81) int8, n_accepted).
    Duplicates are kept: this is a frequency estimate, so throwing repeats
    away would distort exactly the quantity being measured.
    """
    rng = np.random.default_rng(seed)
    forced, empties = split_cells(mask)
    deps = dependencies_from_mask(mask, max_confine)
    cell_deps = build_dep_index(deps)
    counts = np.zeros((81, 9), dtype=np.int32)
    grids = []
    n_ok = 0
    if not empties:
        only = np.zeros(81, dtype=np.int8)
        for c, d in forced.items():
            only[c] = d
        for _ in range(max(keep_grids, 1)):
            grids.append(only.copy())
        counts[np.arange(81), np.maximum(only.astype(np.int64) - 1, 0)] += 1
        counts[only == 0] = 0
        return counts, np.array(grids[:keep_grids], dtype=np.int8), 1
    # Budget the attempts: rejection rates are a few percent at shallow stages
    # but climb as the confinements multiply, so allow generous headroom and
    # report what was actually accepted.
    for a in range(n_draws * 12):
        if n_ok >= n_draws:
            break
        inst = propose(mask, forced, empties, deps, EMPTY, rng, aware=True)
        if max_repair:
            repair(inst, deps, forced, rng, max_repair, cell_deps)
        if _violations(inst, deps) != 0:
            continue
        n_ok += 1
        v = inst.astype(np.int64)
        ok = (v >= 1) & (v <= 9)
        counts[np.arange(81)[ok], v[ok] - 1] += 1
        if len(grids) < keep_grids:
            grids.append(inst.copy())
    while len(grids) < keep_grids and grids:
        grids.append(grids[len(grids) % max(n_ok, 1)].copy())
    return counts, np.array(grids, dtype=np.int8), n_ok


def run_puzzle(args):
    idx, masks, n_draws, keep = args
    S = masks.shape[0]
    counts = np.zeros((S, 81, 9), dtype=np.int32)
    grids = np.zeros((S, keep, 81), dtype=np.int8)
    acc = np.zeros(S, dtype=np.int32)
    for s in range(S):
        c, g, n = sample_stage(masks[s], n_draws, seed=1000 * idx + s,
                               keep_grids=keep)
        counts[s] = c
        if len(g):
            grids[s, :len(g)] = g[:keep]
        acc[s] = n
    return idx, counts, grids, acc


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cand", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--puzzles", type=int, default=400)
    ap.add_argument("--draws", type=int, default=200)
    ap.add_argument("--keep", type=int, default=8)
    ap.add_argument("--workers", type=int, default=16)
    a = ap.parse_args()

    cm = np.load(a.cand, mmap_mode="r")
    P = min(a.puzzles, cm.shape[0])
    S = cm.shape[1]
    print(f"[bigq] {P} puzzles x {S} stages x {a.draws} draws", flush=True)
    counts = np.zeros((P, S, 81, 9), dtype=np.int32)
    grids = np.zeros((P, S, a.keep, 81), dtype=np.int8)
    acc = np.zeros((P, S), dtype=np.int32)

    jobs = ((i, np.asarray(cm[i]).astype(np.int64), a.draws, a.keep)
            for i in range(P))
    if a.workers > 1:
        import multiprocessing as mp
        with mp.Pool(a.workers) as pool:
            for k, (i, c, g, n) in enumerate(
                    pool.imap_unordered(run_puzzle, jobs, chunksize=4)):
                counts[i], grids[i], acc[i] = c, g, n
                if (k + 1) % 50 == 0:
                    print(f"[bigq] {k + 1}/{P} puzzles", flush=True)
    else:
        for k, j in enumerate(jobs):
            i, c, g, n = run_puzzle(j)
            counts[i], grids[i], acc[i] = c, g, n

    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    np.savez_compressed(a.out, counts=counts, grids=grids, accepted=acc)
    print("[bigq] accepted draws per stage (mean):",
          " ".join(f"{s + 1}:{acc[:, s].mean():.0f}" for s in range(S)),
          flush=True)
    print(f"[bigq] wrote {a.out}", flush=True)


if __name__ == "__main__":
    main()