#!/usr/bin/env python3
"""Structural checks and optional live evaluation for the 14-parameter LunarLander policy.

  python verify.py                      # structural checks only (torch)
  python verify.py --episodes 100       # + live Gymnasium evaluation (60 seeds x N episodes)

Reference numbers (60 seeds x 100 episodes): mean 264.06, 60/60 seeds >= 200, 0.35% crashes.
"""
from __future__ import annotations

import argparse
import os

import numpy as np

from model import TinyLunarPolicy, OBS_DIM, N_ACTIONS, N_PARAMS

# The 60 evaluation seeds (episode seed = seed * 1000 + episode_index).
EVAL_SEEDS = [
    40122, 42940, 51305, 55316, 60944, 61197, 93978, 96327, 103680, 157789,
    184578, 197504, 208678, 250957, 289579, 296198, 296953, 309957, 318961, 323741,
    349217, 350692, 373376, 382648, 383930, 385293, 393420, 400382, 414089, 450829,
    467030, 482346, 490224, 498162, 501685, 515927, 535096, 543079, 544042, 553850,
    555722, 571111, 615137, 634353, 642133, 660500, 665066, 684081, 732008, 743613,
    800434, 816622, 819515, 839392, 845498, 852069, 852799, 867535, 894908, 898516,
]


def structural_checks(path: str) -> TinyLunarPolicy:
    m = TinyLunarPolicy.from_safetensors(path)
    W = m.numpy_weights()
    print(f"parameters              : {m.n_params}   (expected {N_PARAMS})")
    print(f"materialised W shape    : {W.shape}   (expected ({OBS_DIM}, {N_ACTIONS}))")
    assert m.n_params == N_PARAMS

    rng = np.random.default_rng(0)
    O = rng.standard_normal((2000, OBS_DIM)).astype(np.float32)
    ref = np.argmax(O @ W, axis=1)
    got = np.asarray([m.act(o) for o in O[:500]])
    assert (ref[:500] == got).all(), "module disagrees with argmax(obs @ W)"
    print(f"argmax(obs @ W) == act  : {(ref[:500] == got).sum()}/500  OK")
    return m


def live_eval(m: TinyLunarPolicy, seeds, n_ep: int, env_id: str = "LunarLander-v3"):
    import gymnasium as gym
    env = gym.make(env_id)
    means, lands, crashes = [], 0, 0
    for s in seeds:
        rs = []
        for ep in range(n_ep):
            o, _ = env.reset(seed=int(s) * 1000 + ep)
            done, tot, last, term, trunc = False, 0.0, 0.0, False, False
            while not done:
                o, r, term, trunc, _ = env.step(m.act(o))
                tot += r
                last = r
                done = term or trunc
            rs.append(tot)
            lands += int(term and not trunc and last == 100.0)
            crashes += int(term and not trunc and last == -100.0)
        means.append(float(np.mean(rs)))
    means = np.asarray(means)
    n = len(seeds) * n_ep
    print(f"{len(seeds)} seeds x {n_ep} ep = {n} episodes")
    print(f"mean={means.mean():.3f}  std(seed means)={means.std():.3f}  "
          f"min={means.min():.2f}  max={means.max():.2f}")
    print(f"seeds with mean>=200: {int((means >= 200).sum())}/{len(means)}")
    print(f"landing rate={lands / n:.4f}   crash rate={crashes / n:.4f}")
    return means


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--weights", default=os.path.join(os.path.dirname(os.path.abspath(__file__)),
                                                     "model.safetensors"))
    ap.add_argument("--episodes", type=int, default=0, help="episodes per seed (0 = skip)")
    ap.add_argument("--seeds", type=int, default=60)
    ap.add_argument("--env-id", default="LunarLander-v3")
    a = ap.parse_args()
    m = structural_checks(a.weights)
    if a.episodes > 0:
        live_eval(m, EVAL_SEEDS[: a.seeds], a.episodes, a.env_id)


if __name__ == "__main__":
    main()