"""Where does the entropy ceiling H(p) <= log|S| sit, per stage?

Two things are printed.

1. The ceiling itself, log|S| averaged over the |S| >= 2 cells of the eval set,
   read straight from the candidate masks. This is the number the model's
   entropy has to stay under, and it shrinks as the curriculum deepens.

2. Whether a model with a given (mass*, spread*) can satisfy the ceiling at
   all. Writing p as mass mu on S with conditional q, and leak 1-mu spread as r
   off S,

       H(p) = mu*H(q) + (1-mu)*H(r) + Hbin(mu),   H(q) = spread * log|S|

   so the margin is

       log|S| - H(p) = log|S| * (1 - mu*spread) - Hbin(mu) - (1-mu)*H(r).

   H(r) is the only term not pinned by the reported metrics, and it is bounded:
   0 if the leak sits on one digit, log(9-|S|) if it spreads over all illegal
   digits. That brackets the margin, which is enough to tell whether the gate is
   reachable without running the model.

The point this makes: Uniform(S) sits exactly ON the ceiling, so at spread* = 1
any leak at all breaks it. The ceiling and a high spread* gate are in tension,
not independent.

Run: python wavecurriculum_run/probe_entropy_ceiling.py
"""
import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
CAND = os.path.join(os.path.dirname(HERE), "datasets_multicandidate_s12",
                    "test_cand_masks.npy")

# (run, stage, mass*, spread*) as last reported by each run's heartbeat.
REPORTED = [
    ("w12_multistage", 3, 0.9512, 0.9227),
]


def hbin(mu):
    mu = np.clip(mu, 1e-12, 1 - 1e-12)
    return -mu * np.log(mu) - (1 - mu) * np.log(1 - mu)


def main():
    if not os.path.exists(CAND):
        sys.exit(f"candidate masks not found at {CAND}")
    # (n_puzzles, n_stages, 81) bitmasks; bit d-1 set iff digit d is a candidate.
    masks = np.load(CAND, mmap_mode="r")
    print(f"{CAND}\nshape {masks.shape}\n")
    sub = np.asarray(masks[:2000])
    n_stages = sub.shape[1]

    sizes = np.zeros(sub.shape, dtype=np.int16)
    for d in range(9):
        sizes += ((sub >> d) & 1).astype(np.int16)

    print("stage  cells|S|>=2   mean|S|  mean log|S|   %|S|=2  %|S|>=5")
    ceil = {}
    for s in range(n_stages):
        sz = sizes[:, s, :].ravel()
        multi = sz[sz >= 2]
        if multi.size == 0:
            print(f"{s + 1:>5}  {'no |S|>=2 cells':>38}")
            continue
        ceil[s + 1] = float(np.log(multi).mean())
        print(f"{s + 1:>5}  {multi.size:>10}  {multi.mean():>8.3f}  "
              f"{ceil[s + 1]:>10.4f}  {100 * (multi == 2).mean():>7.1f}  "
              f"{100 * (multi >= 5).mean():>7.1f}")

    print("\nIs the ceiling reachable at the (mass*, spread*) each run reports?")
    print("margin = log|S|(1 - mass*spread) - Hbin(mass) - (1-mass)H(leak)")
    print("best  = leak on ONE digit (H(r)=0); worst = leak spread over all "
          "illegal digits\n")
    for run, stage, mass, spread in REPORTED:
        lm = ceil.get(stage)
        if lm is None:
            continue
        sz = sizes[:, stage - 1, :].ravel()
        multi = sz[sz >= 2]
        best = lm * (1 - mass * spread) - hbin(mass)
        # worst case pairs each cell's own |S| with a maximally diffuse leak
        h_r = np.log(np.maximum(9 - multi, 1)).mean()
        worst = best - (1 - mass) * h_r
        print(f"{run} stage {stage}: mass*={mass:.4f} spread*={spread:.4f} "
              f"log|S|={lm:.4f}")
        print(f"    H(p) is between {lm - best:.4f} and {lm - worst:.4f} nats")
        print(f"    margin best {best:+.4f}   worst {worst:+.4f}   "
              f"-> ceiling {'HOLDS' if worst >= 0 else ('borderline' if best >= 0 else 'VIOLATED')}")
        need = (hbin(mass)) / max(1 - mass * spread, 1e-9)
        print(f"    to hold in the best case this stage needs log|S| >= "
              f"{need:.4f}, i.e. |S| >= {np.exp(need):.2f}\n")

    print("At spread*=1 the margin is log|S|(1-mass) - Hbin(mass) - ..., which "
          "is < 0 for every mass < 1:")
    print("  mass   log|S|=log2  log4  log6   (spread*=1, leak on one digit)")
    for mass in (1.0, 0.99, 0.95, 0.90):
        row = [f"{np.log(m) * (1 - mass) - hbin(mass):+.4f}" for m in (2, 4, 6)]
        print(f"  {mass:.2f}   " + "  ".join(f"{v:>10}" for v in row))

    # Which threshold triples are self-consistent? A gate set is useless if no
    # distribution can clear it, so check the worst-case margin at the corner of
    # the accept region (mass exactly at bar, spread exactly at bar) for every
    # stage. spread must be gated LOW, purely as a collapse guard, because the
    # ceiling already forbids the excess entropy that a high spread bar demands.
    print("\nWorst-case ceiling margin at the corner of each candidate gate set")
    print("(mass* bar, spread* bar); needs to be >= 0 at every stage to be "
          "satisfiable:\n")
    trials = [(0.95, 0.88), (0.95, 0.75), (0.95, 0.60), (0.98, 0.60),
              (0.95, 0.50)]
    hdr = "  ".join(f"m{int(100 * a)}/s{int(100 * b)}" for a, b in trials)
    print(f"stage  log|S|   {hdr}")
    for s in sorted(ceil):
        sz = sizes[:, s - 1, :].ravel()
        multi = sz[sz >= 2]
        h_r = np.log(np.maximum(9 - multi, 1)).mean()
        cells = []
        for mass, spread in trials:
            mg = (ceil[s] * (1 - mass * spread) - hbin(mass)
                  - (1 - mass) * h_r)
            cells.append(f"{mg:+.3f}" if mg >= 0 else f"[{mg:+.3f}]")
        print(f"{s:>5}  {ceil[s]:.4f}   " + "  ".join(f"{c:>9}" for c in cells))
    print("\n[bracketed] = infeasible: no model can satisfy both that spread "
          "bar and the ceiling at that stage.")


if __name__ == "__main__":
    main()