"""Check the (puzzle, instance) epoch enumeration on the real stage-1 counts.

Verifies the property the experiment rests on: after E epochs every instance of
every puzzle has been used exactly E times, and a batch holds distinct puzzles
(one instance of a given puzzle per batch).
"""
import numpy as np


def build_pairs(starts, counts, stage, n_puz):
    """Same expansion as SudokuDataset._build_instance_epoch_list."""
    c = np.asarray(counts[:n_puz, stage]).astype(np.int64)
    s = np.asarray(starts[:n_puz, stage]).astype(np.int64)
    total = int(c.sum())
    puzzle_ids = np.repeat(np.arange(n_puz, dtype=np.int64), c)
    offsets = np.arange(total, dtype=np.int64) - np.repeat(np.cumsum(c) - c, c)
    rows = np.repeat(s, c) + offsets
    return puzzle_ids, rows, c


def main():
    d = "/tmp/sudoku_superposition"
    starts = np.load(f"{d}/train_starts.npy")
    counts = np.load(f"{d}/train_counts.npy")
    stage, n_puz, epochs, bs = 0, 20000, 5, 64

    pids, rows, c = build_pairs(starts, counts, stage, n_puz)
    total = len(rows)
    print(f"stage {stage + 1}: {n_puz} puzzles -> {total} pairs, "
          f"{c.min()}-{c.max()} instances/puzzle (mean {c.mean():.2f})")

    # 1. rows are exactly each puzzle's contiguous assignment block, no repeats
    assert total == int(c.sum())
    assert len(np.unique(rows)) == total, "duplicate assignment rows"
    for p in (0, 1, 7, n_puz - 1):
        want = np.arange(starts[p, stage], starts[p, stage] + counts[p, stage])
        got = np.sort(rows[pids == p])
        assert np.array_equal(got, want), (p, got, want)
    print("rows match each puzzle's assignment block            OK")

    # 2. every instance seen exactly `epochs` times
    rng = np.random.RandomState(0)
    seen = np.zeros(total, dtype=np.int32)
    order = np.arange(total)
    batch_dup = 0
    n_batches = 0
    for _ in range(epochs):
        rng.shuffle(order)
        seen[order] += 1
        for b in range(0, total - bs, bs):
            sl = pids[order[b:b + bs]]
            batch_dup += bs - len(np.unique(sl))
            n_batches += 1
    assert seen.min() == seen.max() == epochs, (seen.min(), seen.max())
    print(f"every instance seen exactly {epochs}x                     OK")

    # 3. per-puzzle: all of its instances covered, each `epochs` times
    per_puzzle = np.bincount(pids, weights=seen, minlength=n_puz)
    assert np.array_equal(per_puzzle, epochs * c)
    print("every puzzle: all instances x epochs                 OK")

    # 4. batches hold distinct puzzles
    print(f"same-puzzle collisions in a batch: {batch_dup} over "
          f"{n_batches} batches ({batch_dup / (n_batches * bs) * 100:.3f}% "
          f"of slots)")

    # 5. averaged over the epochs, the target at a cell IS the candidate set
    A = np.load(f"{d}/train_assignments.npy", mmap_mode="r")
    p = 3
    blk = np.asarray(A[starts[p, stage]:starts[p, stage] + counts[p, stage]])
    multi = [c_ for c_ in range(81) if len(np.unique(blk[:, c_])) > 1]
    print(f"\npuzzle {p}: {counts[p, stage]} instances, "
          f"{len(multi)} cells whose digit varies across instances")
    for cell in multi[:5]:
        vals, cnt = np.unique(blk[:, cell], return_counts=True)
        emp = ", ".join(f"{v}:{n}/{len(blk)}" for v, n in zip(vals, cnt))
        print(f"  cell (r{cell // 9}, c{cell % 9}) digits {emp}")
    print("\nthis empirical spread over instances is exactly what the model "
          "must reproduce; it only sees it if all instances are visited")


if __name__ == "__main__":
    main()