"""Page data for the "Beam writes the weights" section of the Space.

Reads results/beam_weights.json (written by scripts/beam_weights.py), replays the best network that
Beam wrote on the 20 held-out test levels with the reference emulator, and writes
static-space/weights-data.js with:
  - the weights, the network shape, the feature and action names;
  - per test level: the exact ROM (same bytes the in-weights model runs), the result, and the action
    sequence the network chose, so the browser can rerun the network live and check it picks the
    same actions;
  - the per-round training history (wins / progress per round of feedback);
  - Beam's visible answer for the best round.

Nothing is invented here: if the recorded run has no usable network, this script refuses to write.

  .venv/bin/python scripts/build_weights_page_data.py
"""
import base64
import json
import os
import sys

import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
sys.path.insert(0, os.path.join(ROOT, "scripts"))

from beamboot.superbeam_gym import SuperBeamEnv  # noqa: E402
import beam_weights as bw  # noqa: E402

SRC = os.path.join(ROOT, "results", "beam_weights.json")
OUT = os.path.join(ROOT, "static-space", "weights-data.js")


def replay(net, seed):
    env = SuperBeamEnv(level=seed, difficulty=0.5, frame_skip=4, max_frames=600)
    env.reset()
    term = trunc = False
    info, actions = {}, []
    while not (term or trunc):
        a = bw.act(net, bw.features(env))
        actions.append(a)
        _, _, term, trunc, info = env.step(a)
    return {"seed": seed, "rom_b64": base64.b64encode(bytes(env.backend.rom)).decode(),
            "won": bool(info.get("won")), "dead": bool(info.get("dead")),
            "progress": round(float(info.get("best_progress", 0)), 3),
            "coins": int(info.get("coins", 0)), "coins_total": int(info.get("coins_total", 0)),
            "frames": int(info.get("frames", 0)), "actions": "".join(map(str, actions))}


def main():
    res = json.load(open(SRC))
    if "best_weights" not in res:
        raise SystemExit("no best_weights in %s: the run has not produced a usable network" % SRC)
    net = {k: np.array(v, dtype=float) for k, v in res["best_weights"].items()}
    levels = [replay(net, s) for s in bw.TEST_SEEDS]
    history = [{"round": r["round"], "won": r["train"]["won"], "levels": r["train"]["levels"],
                "progress": r["train"]["mean_progress"]} if "train" in r else
               {"round": r["round"], "error": r.get("error")} for r in res["rounds"]]
    best = next(r for r in res["rounds"] if r["round"] == res["best_round"])
    data = {"model": res["model"], "shape": [bw.N_IN, bw.N_HID, bw.N_OUT],
            "features": [[n, d] for n, d in bw.FEATURES], "actions": bw.ACTIONS,
            "weights": res["best_weights"], "best_round": res["best_round"],
            "train": {"won": res["best_train"]["won"], "levels": res["best_train"]["levels"],
                      "progress": res["best_train"]["mean_progress"]},
            "test": {"won": sum(l["won"] for l in levels), "levels": len(levels),
                     "progress": round(float(np.mean([l["progress"] for l in levels])), 3)},
            "history": history, "answer": best.get("reply", ""), "frame_skip": 4, "max_frames": 600,
            "levels": levels}
    if data["test"]["won"] != res["test"]["won"]:
        raise SystemExit("replay disagrees with the recorded test result (%d vs %d)" % (data["test"]["won"], res["test"]["won"]))
    with open(OUT, "w") as f:
        f.write("window.WS = " + json.dumps(data, separators=(",", ":")) + ";\n")
    print("wrote %s: test won %d/20, mean progress %.0f%%" % (OUT, data["test"]["won"], 100 * data["test"]["progress"]))


if __name__ == "__main__":
    main()