"""Is the large-sample q real, or an artifact of how it is sampled?

gen_bigq.py reuses the generator's own `propose` + `repair`. Neither is a
uniform sampler of the constrained set:

  aware=True  pre-places each confined digit into a random cell of its
              confinement set, so the confinements are satisfied by
              construction rather than by chance.
  repair      is a greedy hill-climber that forces digits into whichever cell
              reduces the violation count, so it can walk grids toward the one
              globally consistent assignment -- i.e. toward the solution.

Either could manufacture the q(true) = 0.63 at stage 4 / 0.99 at stage 10 that
the corrected gate now rests on. This compares three samplers on the same
cells:

  gen      aware + repair            (what gen_bigq.py did)
  aware    aware, no repair, reject violations
  plain    independent uniform draws, no repair, reject violations
             -- an exact rejection sampler for "uniform over S subject to the
                confinements", the honest definition of the stage-s target

If plain agrees with gen, the skew is a property of the constraints. If plain
comes out flat, the skew was the hill-climber and the gate result is wrong.
"""

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,
                                     digits_of, propose, repair, split_cells,
                                     _violations)

EMPTY = frozenset()
POP = np.array([bin(i).count("1") for i in range(512)], np.int16)


def sample(mask, mode, n_target, max_tries, seed):
    """Counts (81,9) plus (accepted, tried)."""
    rng = np.random.default_rng(seed)
    forced, empties = split_cells(mask)
    deps = dependencies_from_mask(mask, 3)
    cell_deps = build_dep_index(deps)
    counts = np.zeros((81, 9), np.int64)
    ok = 0
    tries = 0
    while ok < n_target and tries < max_tries:
        tries += 1
        aware = mode != "plain"
        inst = propose(mask, forced, empties, deps, EMPTY, rng, aware=aware)
        if mode == "gen":
            repair(inst, deps, forced, rng, 200, cell_deps)
        if _violations(inst, deps) != 0:
            continue
        ok += 1
        v = inst.astype(np.int64)
        m = (v >= 1) & (v <= 9)
        counts[np.arange(81)[m], v[m] - 1] += 1
    return counts, ok, tries


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--cand", required=True)
    ap.add_argument("--stages", default="4,6,8,10")
    ap.add_argument("--puzzles", type=int, default=40)
    ap.add_argument("--draws", type=int, default=150)
    ap.add_argument("--tries", type=int, default=40000)
    a = ap.parse_args()
    cm = np.load(a.cand, mmap_mode="r")
    for stage in [int(x) for x in a.stages.split(",")]:
        s = stage - 1
        acc = {m: dict(qt=[], h=[], lk=[], cov=[], ok=0, tries=0,
                       zero_true=0, n=0)
               for m in ("gen", "aware", "plain")}
        for idx in range(a.puzzles):
            mask = np.asarray(cm[idx, s]).astype(np.int64)
            true = np.array([int(b).bit_length()
                             for b in np.asarray(cm[idx, -1])])
            for mode in acc:
                cnt, ok, tries = sample(mask, mode, a.draws, a.tries,
                                        seed=7717 * idx + stage)
                A = acc[mode]
                A["ok"] += ok
                A["tries"] += tries
                if ok < 5:
                    continue
                for c in range(81):
                    k = int(POP[int(mask[c])])
                    if k < 2:
                        continue
                    tot = cnt[c].sum()
                    if tot <= 0:
                        continue
                    q = cnt[c] / tot
                    qq = q[q > 0]
                    A["qt"].append(q[true[c] - 1])
                    A["h"].append(-(qq * np.log(qq)).sum())
                    A["lk"].append(np.log(k))
                    A["cov"].append((q[np.array(
                        [(int(mask[c]) >> d) & 1 for d in range(9)], bool)]
                        > 0).mean())
                    A["zero_true"] += int(q[true[c] - 1] == 0)
                    A["n"] += 1
        print(f"\n===== stage {stage} =====")
        print(f"{'sampler':7s} {'accept%':>8s} {'cells':>7s} {'q(true)':>8s} "
              f"{'1/|S|':>7s} {'H(q)':>7s} {'log|S|':>7s} {'slack':>7s} "
              f"{'cov(q)':>7s} {'q(true)=0':>9s}")
        for mode in ("gen", "aware", "plain"):
            A = acc[mode]
            if not A["n"]:
                print(f"{mode:7s} {100.0 * A['ok'] / max(A['tries'], 1):8.2f} "
                      f"  no accepted samples")
                continue
            qt = np.mean(A["qt"])
            h = np.mean(A["h"])
            lk = np.mean(A["lk"])
            un = np.mean([1.0 / np.exp(x) for x in A["lk"]])
            print(f"{mode:7s} {100.0 * A['ok'] / max(A['tries'], 1):8.2f} "
                  f"{A['n']:7d} {qt:8.3f} {un:7.3f} {h:7.3f} {lk:7.3f} "
                  f"{lk - h:7.3f} {np.mean(A['cov']):7.3f} "
                  f"{A['zero_true'] / A['n']:9.3f}")


if __name__ == "__main__":
    main()