"""Baselines on held-out TEST levels (seeds 1_000_000+) for SUPERBEAM; CPU only, emulator backend.

  .venv/bin/python scripts/gym_baselines_levels.py [--levels 20] [--difficulty 0.5]

Agents: random; scripted (hold right, jump when the frame shows a pit, a pipe/tall pipe (3 lit
pixels in a row at the level it must clear) or a walker within a lead; perception = the 64x32 frame
plus player x/y only); scripted+eps0.1.  Writes results/gym_baselines_levels.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)
sys.path.insert(0, os.path.join(ROOT, "scripts"))

from beamboot.superbeam_gym import GROUND_Y, PIPE_Y, SuperBeamEnv  # noqa: E402
from beamboot.superbeam_levels import TEST_SEEDS  # noqa: E402
from gym_baselines import RandomAgent, run_episode, obstacle_ahead  # noqa: E402

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


def wall_ahead(obs, lead):
    f, x, y = obs["frame"], obs["x"], obs["y"]
    row = 16 if y == GROUND_Y else 8
    for d in range(0, lead + 1):
        c = x + 8 + d
        if c + 2 < 64 and f[row, c] and f[row, c + 1] and f[row, c + 2]:
            return d
    return None


class Scripted:
    def __init__(self, pit_lead=4, enemy_lead=14, wall_lead=4, eps=0.0, seed=0):
        self.pl, self.el, self.wl, self.eps = pit_lead, enemy_lead, wall_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"] not in (GROUND_Y, PIPE_Y, 2):
            return 2
        pit, enemy = obstacle_ahead(obs, max(self.pl, self.el)) if obs["y"] == GROUND_Y else (None, None)
        wall = wall_ahead(obs, self.wl)
        if (pit is not None and pit <= self.pl) or (enemy is not None and enemy <= self.el) or wall is not None:
            return 4
        return 2


def evaluate(name, make_agent, seeds, args):
    eps = []
    for s in seeds:
        env = SuperBeamEnv(level=s, difficulty=args.difficulty, frame_skip=4, max_frames=600)
        r = run_episode(env, make_agent(s))
        r["coins"] = env.n_coins
        r["seed"] = s
        eps.append(r)
    return {"agent": name, "episodes": len(eps),
            "levels_won": sum(e["won"] for e in eps),
            "win_rate": sum(e["won"] for e in eps) / len(eps),
            "mean_progress": float(np.mean([e["progress"] for e in eps])),
            "mean_coins": float(np.mean([e["coins"] for e in eps])),
            "episodes_detail": eps}


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--levels", type=int, default=20)
    ap.add_argument("--difficulty", type=float, default=0.5)
    args = ap.parse_args()
    seeds = [TEST_SEEDS.start + i for i in range(args.levels)]
    rows = [evaluate("random", lambda s: RandomAgent(s), seeds, args),
            evaluate("scripted", lambda s: Scripted(), seeds, args),
            evaluate("scripted+eps0.1", lambda s: Scripted(eps=0.1, seed=s), seeds, args)]
    for r in rows:
        print("%-16s won %d/%d progress %.3f coins %.2f" % (r["agent"], r["levels_won"], r["episodes"],
                                                          r["mean_progress"], r["mean_coins"]))
    with open(OUT, "w") as f:
        json.dump({"levels": seeds, "difficulty": args.difficulty, "frame_skip": 4, "results": rows}, f, indent=1)


if __name__ == "__main__":
    main()