"""Score a saved checkpoint on the CE(q||p) <= log|S| promotion criterion.

Runs the same evaluation the trainer runs, on the same held-out split, but
offline against an arbitrary checkpoint directory, and reports the per-stage
value statistics with the CE criterion first.

The criterion, per cell with |S| >= 2 and headroom:

    CE(q||p)  <=  log|S|

where q is the empirical digit distribution of that puzzle's stored instances at
that stage (post-constraint, what the data teaches) and log|S| is the entropy of
the uniform distribution over the raw pre-constraint candidate set. Cells where
q is already uniform on S are skipped: there the bar equals the floor H(q), so by
Gibbs only p == q exactly could pass.

Usage:
  python wavecurriculum_run/eval_ceq_ckpt.py --ckpt_dir DIR [--stage 6]
                                             [--eval_epochs 40] [--label NAME]

DIR is a directory containing a checkpoint_<step> file, as unpacked from the
latest.tar / stageN_step.tar archives on the Hub.
"""
import argparse
import functools
import json
import os
import sys

os.environ.setdefault("SUDOKU_VERBOSE_EVAL", "0")

import numpy as np


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt_dir", required=True)
    ap.add_argument("--stage", type=int, default=0,
                    help="Curriculum stage to evaluate at; 0 = read it from "
                         "the checkpoint's sibling latest.json, else max_stage.")
    ap.add_argument("--eval_epochs", type=int, default=40)
    ap.add_argument("--label", default="")
    ap.add_argument("--out_json", default="")
    args = ap.parse_args()

    # The config is assembled from the environment exactly as the training jobs
    # assemble it, so the launcher must export the same SUDOKU_* variables. Only
    # the eval-side knobs are overridden here.
    os.environ["SUDOKU_EVAL_EPOCHS"] = str(args.eval_epochs)

    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    import jax
    import jax.numpy as jnp
    from jax import random
    from flax.training import checkpoints
    from flax import jax_utils

    from train import main as train_main
    from train import data, model, trainer, evaluater

    config = train_main.get_config()
    K = int(config.num_latent_slots)
    stage = args.stage or int(getattr(config, "curriculum_max_stage", 6))

    # Eval iterator only; the curriculum stage matters because the candidate
    # masks and instance pool are read per stage.
    curriculum = data.CurriculumState(
        stage=stage, max_stage=int(getattr(config, "curriculum_max_stage", 6)))
    eval_iter = data.create_iter(config, config.minibatch_size, train=False,
                                 curriculum=curriculum)

    model_config = model.TransformerConfig(
        dtype=config.dtype,
        vocab_size=config.vocab_size,
        seq_len=config.seq_len,
        num_heads=config.num_heads,
        num_layers=config.num_layers,
        emb_dim=config.emb_dim,
        qkv_dim=config.qkv_dim,
        mlp_dim=config.mlp_dim,
        dropout_rate=config.dropout_rate,
        attention_dropout_rate=config.attention_dropout_rate,
        deterministic=False,
        num_latent_slots=K,
        inject_latents=bool(int(getattr(config, "recurrent_latent", 1))),
    )

    rng = jax.random.PRNGKey(config.seed)
    rng, init_rng, _ = random.split(rng, num=3)
    rng, dropout_rng = jax.random.split(rng)
    net = model.TransformerLMHeadModel(model_config)
    dummy_latents = jnp.zeros(
        (config.minibatch_size, K, config.emb_dim), model_config.dtype)
    dummy_positions = jnp.zeros((config.minibatch_size, K), jnp.int32)
    dummy_active = jnp.zeros((config.minibatch_size, K), bool)
    _, initial_variables = jax.jit(net.init_with_output)(
        {"params": init_rng, "dropout": dropout_rng},
        jnp.ones((config.minibatch_size, config.seq_len), jnp.int32),
        dummy_latents, dummy_positions, dummy_active)
    state, _ = trainer.get_state(config, net, initial_variables)

    restored = checkpoints.restore_checkpoint(os.path.abspath(args.ckpt_dir),
                                              state)
    step = int(restored.step)
    if step == 0:
        sys.exit(f"no checkpoint restored from {args.ckpt_dir}")
    print(f"restored step {step} from {args.ckpt_dir}", flush=True)
    state = jax_utils.replicate(restored)

    p_eval_step = jax.pmap(
        functools.partial(evaluater.eval_step,
                          config=model_config.replace(deterministic=True)),
        axis_name="batch")

    em = evaluater.get_eval_metrics(state, eval_iter, p_eval_step, config)

    ceq = em.get("per_stage_val_ceq", {})
    hq = em.get("per_stage_val_hq", {})
    ok = em.get("per_stage_val_ceq_ok", {})
    gap = em.get("per_stage_val_ceq_gap", {})
    puz = em.get("per_stage_val_ceq_puz", {})
    mass = em.get("per_stage_val_mass_multi", {})
    spread = em.get("per_stage_val_spread", {})
    hb = em.get("per_stage_val_hbound", {})

    tag = args.label or os.path.basename(os.path.abspath(args.ckpt_dir))
    print(f"\n==== {tag}  step {step}  eval_stage {stage} "
          f"({args.eval_epochs} eval batches) ====")
    print("stage   H(q)      ceq   log|S|*    slack   ceq_ok   ceq_puz"
          "   mass*  spread*  hbound")
    rows = {}
    for s in sorted(ceq):
        if ceq[s] < 0:
            continue
        # log|S| on the counted cells, recovered from ceq + its slack.
        lm = ceq[s] + gap[s]
        rows[s] = {"h_q": hq[s], "ceq": ceq[s], "log_m": lm, "slack": gap[s],
                   "ceq_ok": ok[s], "ceq_puz": puz.get(s, -1.0),
                   "mass_multi": mass.get(s, -1.0),
                   "spread": spread.get(s, -1.0), "hbound": hb.get(s, -1.0)}
        flag = "" if gap[s] >= 0 else "  <-- over ceiling"
        print(f"{s:>5}  {hq[s]:6.4f}  {ceq[s]:7.4f}  {lm:7.4f}  {gap[s]:+7.4f}"
              f"   {ok[s]:6.3f}   {puz.get(s, -1):7.3f}"
              f"  {mass.get(s, -1):6.3f}  {spread.get(s, -1):7.3f}"
              f"  {hb.get(s, -1):6.3f}{flag}")

    if stage in rows:
        r = rows[stage]
        verdict = ("PROMOTE" if r["ceq_ok"] >= 0.90 else "HOLD")
        print(f"\nfrontier stage {stage}: ceq_ok = {r['ceq_ok']:.3f} "
              f"vs bar 0.90  ->  {verdict}")
        print(f"  mean CE(q||p) = {r['ceq']:.4f}, floor H(q) = {r['h_q']:.4f}, "
              f"ceiling log|S| = {r['log_m']:.4f}")

    if args.out_json:
        with open(args.out_json, "w") as f:
            json.dump({"label": tag, "step": step, "eval_stage": stage,
                       "eval_epochs": args.eval_epochs, "per_stage": rows}, f,
                      indent=1)
        print(f"wrote {args.out_json}")


if __name__ == "__main__":
    main()