"""Baselines for the SUPERBEAM gym env (beamboot/superbeam_gym.py).

  python scripts/gym_baselines.py [--seeds 20] [--frame-skip 4]

Agents
  random        uniform over the 6 actions
  scripted      hold right; jump (right+jump) when the frame shows a pit in the
                ground row or a walker within `lead` px ahead.  Perception uses
                only the 64x32 frame plus the player's screen x/y.
  scripted+eps  scripted with 10% uniformly random actions (seeded)

Writes results/gym_baselines.json.
"""
import argparse
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)

from beamboot.superbeam_gym import GROUND_Y, SuperBeamEnv  # noqa: E402

OUT = os.path.join(ROOT, "results", "gym_baselines.json")


def obstacle_ahead(obs, lead):
    """(pit_dist, enemy_dist) in px from the player's right edge, or None."""
    f, x, y = obs["frame"], obs["x"], obs["y"]
    right = x + 8
    pit = enemy = None
    for d in range(0, lead + 1):
        c = right + d
        if c >= 64:
            break
        if pit is None and not f[24, c]:
            pit = d
        # walker body row (sprite row 3 = 0xFF) sits at y = 21; flag pole is 1 px wide
        if enemy is None and f[GROUND_Y + 3, c] and f[GROUND_Y + 3, min(c + 1, 63)] \
                and f[GROUND_Y + 3, min(c + 2, 63)]:
            enemy = d
    return pit, enemy


class ScriptedAgent:
    def __init__(self, pit_lead=4, enemy_lead=14, eps=0.0, seed=0):
        self.pit_lead, self.enemy_lead, self.eps = pit_lead, enemy_lead, eps
        self.rng = np.random.default_rng(seed)

    def act(self, obs):
        if self.eps and self.rng.random() < self.eps:
            return int(self.rng.integers(6))
        if obs["y"] != GROUND_Y:
            return 2
        pit, enemy = obstacle_ahead(obs, max(self.pit_lead, self.enemy_lead))
        if (pit is not None and pit <= self.pit_lead) or \
                (enemy is not None and enemy <= self.enemy_lead):
            return 4
        return 2


class RandomAgent:
    def __init__(self, seed=0):
        self.rng = np.random.default_rng(seed)

    def act(self, obs):
        return int(self.rng.integers(6))


def run_episode(env, agent):
    obs, info = env.reset()
    steps = 0
    while True:
        obs, r, term, trunc, info = env.step(agent.act(obs))
        steps += 1
        if term or trunc:
            return {"won": info["won"], "dead": info["dead"], "progress": info["best_progress"],
                    "frames": info["frames"], "steps": steps, "reward": env.total_reward}


def evaluate(name, make_agent, seeds, frame_skip, max_frames):
    env = SuperBeamEnv(frame_skip=frame_skip, max_frames=max_frames)
    eps = [run_episode(env, make_agent(s)) for s in seeds]
    return {"agent": name, "episodes": len(eps),
            "win_rate": sum(e["won"] for e in eps) / len(eps),
            "death_rate": sum(e["dead"] for e in eps) / len(eps),
            "mean_progress": float(np.mean([e["progress"] for e in eps])),
            "max_progress": float(max(e["progress"] for e in eps)),
            "mean_steps": float(np.mean([e["steps"] for e in eps])),
            "episodes_detail": eps}


def tune(frame_skip, max_frames):
    """Grid over leads for the deterministic scripted agent; returns winning (pit, enemy) leads."""
    env = SuperBeamEnv(frame_skip=frame_skip, max_frames=max_frames)
    wins = []
    for pl in range(0, 24, 2):
        for el in range(0, 40, 2):
            if run_episode(env, ScriptedAgent(pl, el))["won"]:
                wins.append((pl, el))
    return wins


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--seeds", type=int, default=20)
    ap.add_argument("--frame-skip", type=int, default=4)
    ap.add_argument("--max-frames", type=int, default=600)
    args = ap.parse_args(argv)
    seeds = list(range(args.seeds))
    fs, mf = args.frame_skip, args.max_frames

    wins = tune(fs, mf)
    pl, el = wins[len(wins) // 2] if wins else (4, 14)
    print("frame_skip=%d: scripted lead grid wins %d/%d cells; using pit_lead=%d enemy_lead=%d"
          % (fs, len(wins), 12 * 20, pl, el))
    rows = [
        evaluate("random", lambda s: RandomAgent(s), seeds, fs, mf),
        evaluate("scripted", lambda s: ScriptedAgent(pl, el), seeds[:1], fs, mf),
        evaluate("scripted+eps0.1", lambda s: ScriptedAgent(pl, el, eps=0.1, seed=s), seeds, fs, mf),
    ]
    print("%-16s %4s %8s %8s %10s %10s" % ("agent", "eps", "win", "death", "mean_prog", "max_prog"))
    for r in rows:
        print("%-16s %4d %8.2f %8.2f %10.3f %10.3f" % (r["agent"], r["episodes"], r["win_rate"],
                                                    r["death_rate"], r["mean_progress"], r["max_progress"]))
    os.makedirs(os.path.dirname(OUT), exist_ok=True)
    with open(OUT, "w") as f:
        json.dump({"frame_skip": fs, "max_frames": mf, "scripted_leads": [pl, el],
                   "lead_grid_wins": wins, "results": rows}, f, indent=1)
    print("wrote", os.path.relpath(OUT, ROOT))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())