"""Weight-space programming with Beam: Beam-1.0-501B writes the weights of a small neural network
that plays SUPERBEAM, with no training (no data, no loss, no gradients), in the spirit of
https://rikiyat.github.io/blog/weight-space-programming/ .

The network: 8 features -> 16 ReLU -> 6 action logits, action = argmax. Features are read from the
game's registers/RAM each turn (see `features`). Beam gets the game rules, the feature definitions and
the network shape, and replies with the weights as JSON. Each candidate is scored on TRAINING levels
(seeds 0..TRAIN_LEVELS-1); Beam gets per-level results as feedback and may revise. The candidate with
the best training score is then evaluated once on the 20 held-out TEST levels used by the leaderboard
(seeds 1000000..1000019, difficulty 0.5, one attempt per level).

  modal run scripts/beam_weights.py [--rounds 4]
  -> results/beam_weights.json (all rounds: Beam's reply, weights, scores) and the test result.

The Hugging Face token stays in the Modal secret `huggingface-secret` and is never printed.
"""
import json
import os
import re
import sys
import time
import urllib.error
import urllib.request

import modal
import numpy as np

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
OUT = os.path.join(ROOT, "results", "beam_weights.json")
MODEL = "reflection/Beam-1.0-501B-A23B:baseten"
TRAIN_LEVELS = 10
TEST_SEEDS = list(range(1000000, 1000020))
N_IN, N_HID, N_OUT = 8, 16, 6
ACTIONS = ["noop", "left", "right", "jump", "right+jump", "left+jump"]
FEATURES = [
    ("standing", "1 if you stand on ground or on a pipe this frame, else 0 (in the air)"),
    ("height", "how high your feet are: 0 on ground level, 8 on a pipe top, 16 on a tall pipe top (px)"),
    ("gap_dist", "px from your front (right) edge to the start of the next gap in the ground ahead; 0 if the gap is already under your front edge; 32 if no gap within 32 px"),
    ("gap_width", "width of that gap in px (8 or 16); 0 if none"),
    ("wall_dist", "px from your front edge to the next block whose top is higher than your feet (a pipe seen from the ground, or a tall pipe seen from a pipe); 32 if none within 32 px"),
    ("wall_height", "how many px higher that block's top is than your feet (8 or 16); 0 if none"),
    ("walker_dist", "px from your front edge to the nearest walker ahead of you (0 if touching); 32 if none within 32 px"),
    ("vy", "vertical speed in px per frame: negative while rising, positive while falling, 0 when standing"),
]

app = modal.App("beamboot-beam-weights")


@app.function(image=modal.Image.debian_slim(python_version="3.12").pip_install("numpy"),
              secrets=[modal.Secret.from_name("huggingface-secret")], timeout=600)
def beam_chat(messages, model=MODEL, max_tokens=32000):
    token = (os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN")
             or os.environ.get("HUGGINGFACE_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN"))
    if not token:
        raise RuntimeError("Hugging Face token not found in the Modal secret")
    payload = {"model": model, "messages": messages, "max_tokens": max_tokens}
    req = urllib.request.Request("https://router.huggingface.co/v1/chat/completions",
                                 data=json.dumps(payload).encode(),
                                 headers={"Authorization": "Bearer " + token, "Content-Type": "application/json"})
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=580) as response:
            out = json.load(response)
    except urllib.error.HTTPError as error:
        raise RuntimeError("HF router HTTP %d: %s" % (error.code, error.read(800).decode(errors="replace"))) from None
    out["_latency_s"] = round(time.time() - t0, 1)
    return out


# ------------------------------------------------------------------ the game side (local, CPU)
def features(env):
    """The 8 features, read from the emulator's registers and RAM."""
    regs = env.backend.regs()
    x, y, cam = regs[0xA], regs[0xB], regs[0xD]
    wx = cam * 8 + x
    front = wx + 8
    standing = 1.0 if regs[0xE] == y and y in (18, 10, 2) else 0.0
    feet = y + 6                                         # first pixel row below the player
    height = float(max(0, 24 - feet))
    lv = env.live_level()
    ncol = len(lv)

    def top(c):                                          # top pixel row of column c (None = gap)
        if c >= ncol:
            return 24
        t = lv[c] & 3
        return None if t == 0 else {1: 24, 2: 16, 3: 8}[t]

    gap_dist, gap_width = 32.0, 0.0
    wall_dist, wall_height = 32.0, 0.0
    c0 = max(0, wx // 8)
    for c in range(c0, min(ncol, c0 + 7)):
        left = c * 8
        if left + 8 <= front - 8:                        # column entirely behind the player
            continue
        d = max(0, left - front)
        if d > 32:
            break
        t = top(c)
        if t is None and gap_dist == 32.0:
            w = 0
            while c + w < ncol and top(c + w) is None:
                w += 1
            gap_dist, gap_width = float(d), float(8 * w)
        if t is not None and t < feet and wall_dist == 32.0 and left >= front - 1:
            wall_dist, wall_height = float(d), float(feet - t)
    walker_dist = 32.0
    et = env.sym["ETAB"]
    for e in range(2):
        alive, col, off = (env.backend.ram(et + 4 * e + i) for i in range(3))
        if not alive:
            continue
        d = col * 8 + off - front
        if -8 <= d <= 32:
            walker_dist = min(walker_dist, float(max(0, d)))
    vy = 0.0 if standing else float(regs[0xC] - 8)
    return np.array([standing, height, gap_dist, gap_width, wall_dist, wall_height, walker_dist, vy])


def act(net, f):
    h = np.maximum(0.0, net["W1"] @ f + net["b1"])
    return int(np.argmax(net["W2"] @ h + net["b2"]))


def play(net, seed, max_frames=600):
    from beamboot.superbeam_gym import SuperBeamEnv
    env = SuperBeamEnv(level=seed, difficulty=0.5, frame_skip=4, max_frames=max_frames)
    env.reset()
    term = trunc = False
    info, last_f = {}, None
    while not (term or trunc):
        last_f = features(env)
        _, _, term, trunc, info = env.step(act(net, last_f))
    cause = None
    if info.get("dead"):
        cause = ("fell into a gap" if last_f[2] <= 8 else "hit a walker" if last_f[6] <= 8 else "died")
    elif not info.get("won"):
        cause = "ran out of time (stuck: e.g. walking into a pipe wall, or standing still)"
    return {"seed": seed, "won": bool(info.get("won")), "progress": round(float(info.get("best_progress", 0)), 3),
            "coins": int(info.get("coins", 0)), "cause": cause,
            "where": None if last_f is None else dict(zip([n for n, _ in FEATURES], [round(float(v), 1) for v in last_f]))}


def score(net, seeds):
    res = [play(net, s) for s in seeds]
    return {"won": sum(r["won"] for r in res), "levels": len(res),
            "mean_progress": round(float(np.mean([r["progress"] for r in res])), 3), "detail": res}


def parse_net(text):
    blocks = re.findall(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.S) or re.findall(r"(\{\s*\"W1\".*\})", text, re.S)
    if not blocks:
        raise ValueError("no JSON object with W1/b1/W2/b2 found")
    d = json.loads(blocks[-1])
    net = {k: np.array(d[k], dtype=float) for k in ("W1", "b1", "W2", "b2")}
    shapes = {"W1": (N_HID, N_IN), "b1": (N_HID,), "W2": (N_OUT, N_HID), "b2": (N_OUT,)}
    for k, s in shapes.items():
        if net[k].shape != s:
            raise ValueError("%s has shape %s, expected %s" % (k, net[k].shape, s))
    return net


SYSTEM = """You are a weight-space programmer. You write the weights of a small neural network directly, by reasoning about what each unit should compute. There is no training, no data and no gradient descent: the numbers you write are the final network."""


def task_prompt():
    sys.path.insert(0, os.path.join(ROOT, "scripts"))
    from beam_play import SYSTEM as RULES                 # the same game rules Beam is given when it plays
    feats = "\n".join("  x[%d] %s: %s" % (i, n, d) for i, (n, d) in enumerate(FEATURES))
    acts = "\n".join("  %d %s" % (i, a) for i, a in enumerate(ACTIONS))
    rules = RULES.split("Actions:")[0].strip()
    return f"""Write the weights of a neural network that plays SUPERBEAM well.

THE GAME (these are the rules a player is given):
{rules}

THE NETWORK. Every 4 frames it reads 8 numbers x[0..7] and picks one action:
  h = relu(W1 @ x + b1)        # W1: 16 x 8, b1: 16
  logits = W2 @ h + b2         # W2: 6 x 16, b2: 6
  action = argmax(logits)
Inputs (read from the game's memory; distances in pixels):
{feats}
Actions (index: name):
{acts}

Design it like a program: give each hidden unit a job (for example "a gap is close", "a wall is close
and I stand on the ground", "a walker is close", "I am in the air"), set its weights and bias so it
switches on only in that situation (ReLU), then wire the units to the action logits so that the right
action wins. Keep "right" as the default so the player keeps moving. Unused units can be all zeros.

Reply with your reasoning, then ONE JSON object in a ```json block with keys "W1" (16 lists of 8 numbers),
"b1" (16 numbers), "W2" (6 lists of 16 numbers), "b2" (6 numbers). Numbers only, no expressions."""


def feedback(sc):
    lines = ["Your network on %d training levels: won %d, mean progress %.0f%%." % (sc["levels"], sc["won"], 100 * sc["mean_progress"])]
    for r in sc["detail"]:
        if r["won"]:
            lines.append("  level %d: won" % r["seed"])
        else:
            lines.append("  level %d: %s at %.0f%% of the level; inputs at that moment: %s"
                         % (r["seed"], r["cause"], 100 * r["progress"], json.dumps(r["where"])))
    lines.append("Improve the weights to win more levels. Reply with reasoning, then the full JSON (all four arrays).")
    return "\n".join(lines)


@app.local_entrypoint()
def main(rounds: int = 4):
    train = list(range(TRAIN_LEVELS))
    messages = [{"role": "system", "content": SYSTEM}, {"role": "user", "content": task_prompt()}]
    log, best = [], None
    for r in range(rounds):
        out = beam_chat.remote(messages)
        msg = out["choices"][0]["message"]
        text = msg.get("content") or ""
        reasoning = msg.get("reasoning_content") or msg.get("reasoning") or ""
        finish = out["choices"][0].get("finish_reason")
        entry = {"round": r + 1, "latency_s": out.get("_latency_s"), "finish_reason": finish,
                 "usage": out.get("usage"), "reply": text, "reasoning_chars": len(reasoning)}
        print("round %d: finish=%s content=%d chars reasoning=%d chars" % (r + 1, finish, len(text), len(reasoning)), flush=True)
        try:
            try:
                net = parse_net(text)
            except ValueError:                           # JSON sometimes lands only in the reasoning
                net = parse_net(reasoning)
                entry["parsed_from"] = "reasoning_content"
            sc = score(net, train)
            entry.update(weights={k: v.tolist() for k, v in net.items()}, train=sc)
            print("round %d: train won %d/%d, mean progress %.0f%% (%.0fs)" % (r + 1, sc["won"], sc["levels"], 100 * sc["mean_progress"], out.get("_latency_s", 0)), flush=True)
            if best is None or (sc["won"], sc["mean_progress"]) > (best[1]["won"], best[1]["mean_progress"]):
                best = (net, sc, r + 1)
            follow = feedback(sc)
        except Exception as e:                           # unparseable reply: ask again
            entry["error"] = str(e)
            print("round %d: %s" % (r + 1, e), flush=True)
            follow = "I could not use that: %s. Reply with the full JSON object in a ```json block." % e
        log.append(entry)
        json.dump({"model": MODEL, "rounds": log}, open(OUT, "w"), indent=1)
        if best and best[1]["won"] == best[1]["levels"]:
            break
        turn = {"role": "assistant", "content": text}
        if reasoning:
            turn["reasoning_content"] = reasoning
        messages += [turn, {"role": "user", "content": follow}]
    result = {"model": MODEL, "network": "8 -> 16 ReLU -> 6, weights written by Beam, no training",
              "rounds": log}
    if best:
        test = score(best[0], TEST_SEEDS)
        result.update(best_round=best[2], best_train=best[1], test=test,
                      best_weights={k: v.tolist() for k, v in best[0].items()})
        print("TEST (20 held-out levels): won %d/20, mean progress %.0f%% (round %d weights)"
              % (test["won"], 100 * test["mean_progress"], best[2]), flush=True)
    json.dump(result, open(OUT, "w"), indent=1)
    print("wrote", OUT)