"""
🏓 Mr. Pong — an interactive demo for fromziro/MrPong.

Two ways to meet the agent:
  1. Play against it live in the browser. The 28K-parameter policy is exported
     from the checkpoint and evaluated client-side at 40 physics steps/second,
     so the paddle reacts with zero network latency.
  2. Watch / benchmark the real PyTorch checkpoint playing full matches against
     the baseline opponents from the model card, rendered to video on the server.
"""

import base64
import html as html_lib
import math
import os
import random
import tempfile
import time
from pathlib import Path

import gradio as gr
import numpy as np
import torch
from PIL import Image, ImageDraw, ImageFont
from transformers import AutoConfig, AutoModel

from pong_engine import (
    OPPONENTS,
    PhysicsConfig,
    StandalonePongEnv,
)

MODEL_ID = "fromziro/MrPong"
HERE = Path(__file__).parent

# ----------------------------------------------------------------------------
# Model (28,484 params — pure CPU, a forward pass costs microseconds)
# ----------------------------------------------------------------------------
config = AutoConfig.from_pretrained(MODEL_ID, trust_remote_code=True)
model = AutoModel.from_pretrained(MODEL_ID, trust_remote_code=True)
model.eval()
OBS_DIM = int(getattr(config, "obs_dim", 12))
N_PARAMS = sum(p.numel() for p in model.parameters())
print(f"[OK] {MODEL_ID} loaded — {N_PARAMS:,} params, obs_dim={OBS_DIM}, "
      f"hidden_dims={config.hidden_dims}")


@torch.no_grad()
def policy_action(obs: np.ndarray, deterministic: bool = True) -> int:
    """Greedy action from the Mr. Pong policy (0 stay, 1 up, 2 down)."""
    x = obs[:OBS_DIM] if len(obs) >= OBS_DIM else np.pad(obs, (0, OBS_DIM - len(obs)))
    return int(model.act(x, deterministic=deterministic))


# ----------------------------------------------------------------------------
# Export the checkpoint for the in-browser copy of the policy
# ----------------------------------------------------------------------------
def _export_weights_b64() -> str:
    sd = model.state_dict()
    order = ["trunk.0.weight", "trunk.0.bias", "trunk.2.weight", "trunk.2.bias",
             "actor.weight", "actor.bias", "critic.weight", "critic.bias"]
    flat = np.concatenate([sd[k].detach().cpu().numpy().astype("<f4").ravel()
                           for k in order])
    assert flat.size == N_PARAMS, (flat.size, N_PARAMS)
    return base64.b64encode(flat.tobytes()).decode("ascii")


GAME_HTML = (HERE / "game_template.html").read_text().replace(
    "__WEIGHTS_B64__", _export_weights_b64()
)
GAME_IFRAME = (
    '<iframe title="Play Mr. Pong" srcdoc="{}" '
    'style="width:100%;height:700px;border:0;border-radius:12px;overflow:hidden" '
    'scrolling="no"></iframe>'
).format(html_lib.escape(GAME_HTML, quote=True))


# ----------------------------------------------------------------------------
# Server-side match rendering (real PyTorch checkpoint)
# ----------------------------------------------------------------------------
SC = 0.8                      # render scale: 800x500 table -> 640x400 video
VW, VH = int(800 * SC), int(500 * SC)
FPS = 40                      # one frame per physics sub-step == real time
MAX_FRAMES = 1100             # hard cap on video length (~27 s)
POINT_SUBSTEP_CAP = 420       # a single point is called a draw after this
MAX_POINTS = 12               # safety net against endless draw sequences


def _font(size: int):
    for p in ("/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
              "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf"):
        if os.path.exists(p):
            return ImageFont.truetype(p, size)
    try:
        return ImageFont.load_default(size=size)
    except TypeError:
        return ImageFont.load_default()


F_BIG, F_MED, F_SMALL = _font(34), _font(15), _font(13)


def _render_frame(env, score_a, score_b, opp_name, trail, banner=""):
    img = Image.new("RGB", (VW, VH), (15, 22, 38))
    d = ImageDraw.Draw(img)

    for y in range(0, VH, 26):
        d.line([(VW // 2, y), (VW // 2, min(VH, y + 12))], fill=(38, 48, 70), width=2)
    d.rectangle([0, 0, VW - 1, VH - 1], outline=(28, 36, 54), width=2)

    for i, (tx, ty) in enumerate(trail):
        a = (i + 1) / (len(trail) + 1)
        c = int(60 + 150 * a)
        r = 8 * SC * (0.35 + 0.6 * a)
        d.ellipse([tx * SC - r, ty * SC - r, tx * SC + r, ty * SC + r], fill=(c, c, c))

    pw = 14 * SC
    eh = env.ego_paddle_h * SC / 2
    oh = env.opp_paddle_h * SC / 2
    d.rounded_rectangle([2, env.ego_y * SC - eh, 2 + pw, env.ego_y * SC + eh],
                        radius=5, fill=(79, 140, 255))          # Mr. Pong (blue)
    d.rounded_rectangle([VW - pw - 2, env.opp_y * SC - oh, VW - 2, env.opp_y * SC + oh],
                        radius=5, fill=(249, 115, 22))          # baseline (orange)

    r = 8 * SC
    d.ellipse([env.ball_x * SC - r, env.ball_y * SC - r,
               env.ball_x * SC + r, env.ball_y * SC + r], fill=(255, 255, 255))

    d.text((VW // 2 - 46, 16), str(score_a), font=F_BIG, fill=(79, 140, 255), anchor="ma")
    d.text((VW // 2 + 46, 16), str(score_b), font=F_BIG, fill=(249, 115, 22), anchor="ma")
    d.text((20, 20), "MR. PONG", font=F_MED, fill=(120, 150, 200))
    d.text((VW - 20, 20), opp_name.upper(), font=F_MED, fill=(190, 130, 80), anchor="ra")

    speed = math.hypot(env.ball_vx, env.ball_vy)
    d.text((20, VH - 26),
           f"rally {env.rally_count} hits   ball {speed:4.1f} px/f", font=F_SMALL,
           fill=(120, 135, 160))

    if banner:
        d.rectangle([0, VH // 2 - 34, VW, VH // 2 + 34], fill=(11, 15, 25))
        d.text((VW // 2, VH // 2), banner, font=F_BIG, fill=(249, 115, 22), anchor="mm")
    return np.asarray(img)


def _play_point(env, opponent, frames, opp_name, score_a, score_b, serve_dir, record):
    """Run one point sub-step by sub-step; returns the winner string."""
    env.reset(serve_direction=serve_dir)
    trail = []
    ai_act = 0
    opp_act = 0
    for i in range(POINT_SUBSTEP_CAP):
        # Both policies are re-queried every `frame_skip` sub-steps, exactly as
        # StandalonePongEnv.step() does in the author's simulate() loop.
        if i % env.phys.frame_skip == 0:
            ai_act = policy_action(env.get_ego_observation())
            opp_act = opponent.act(env)
        trail.append((env.ball_x, env.ball_y))
        if len(trail) > 8:
            trail.pop(0)
        done, info = env.physics_substep(ego_action=ai_act, opp_action=opp_act)
        if record and len(frames) < MAX_FRAMES:
            frames.append(_render_frame(env, score_a, score_b, opp_name, trail))
        if done:
            return info.get("winner")
    return "draw"


def watch_match(opponent_name: str, points_to_win: int = 3,
                randomize_seed: bool = True, seed: int = 0):
    """Simulate and render a Mr. Pong match against a baseline opponent.

    Args:
        opponent_name: which scripted baseline to face (from the model card).
        points_to_win: number of points needed to take the match.
        randomize_seed: draw a fresh random seed for this match.
        seed: RNG seed used when randomize_seed is off.

    Returns:
        An MP4 of the match, a markdown scoreline, and the seed actually used.
    """
    import imageio.v2 as imageio

    if randomize_seed:
        seed = random.randint(0, 2**31 - 1)
    seed = int(seed)
    rng = random.Random(seed)
    random.seed(seed)
    np.random.seed(seed % (2**32))

    points_to_win = int(points_to_win)
    opp_cls = OPPONENTS[opponent_name]
    opponent = opp_cls(rng=rng)

    env = StandalonePongEnv(PhysicsConfig(), seed=seed)
    frames = []
    score_a = score_b = 0
    serve_dir = 1
    rallies = []
    t0 = time.perf_counter()

    while (score_a < points_to_win and score_b < points_to_win
           and len(rallies) < MAX_POINTS):
        winner = _play_point(env, opponent, frames, opponent_name,
                             score_a, score_b, serve_dir,
                             record=len(frames) < MAX_FRAMES)
        rallies.append(env.rally_count)
        if winner == "ego":
            score_a += 1
            serve_dir = 1
        elif winner == "opponent":
            score_b += 1
            serve_dir = -1
        else:
            serve_dir = -serve_dir
        if len(frames) >= MAX_FRAMES:
            break

    banner = ("MR. PONG WINS %d-%d" % (score_a, score_b) if score_a > score_b
              else ("%s WINS %d-%d" % (opponent_name.upper(), score_b, score_a)
                    if score_b > score_a else "TIME LIMIT"))
    for _ in range(FPS * 2):
        if len(frames) < MAX_FRAMES + FPS * 2:
            frames.append(_render_frame(env, score_a, score_b, opponent_name,
                                        [], banner=banner))

    sim_s = time.perf_counter() - t0
    out = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
    with imageio.get_writer(out, fps=FPS, codec="libx264", quality=7,
                            macro_block_size=1, pixelformat="yuv420p") as w:
        for f in frames:
            w.append_data(f)

    stats = (
        f"### Mr. Pong **{score_a} – {score_b}** {opponent_name}\n"
        f"| | |\n|---|---|\n"
        f"| Points played | {len(rallies)} |\n"
        f"| Longest rally | {max(rallies)} hits |\n"
        f"| Average rally | {np.mean(rallies):.1f} hits |\n"
        f"| Video length | {len(frames) / FPS:.1f} s ({len(frames)} frames) |\n"
        f"| Simulation time | {sim_s:.2f} s |\n"
        f"| Seed | `{seed}` |\n"
    )
    return out, stats, seed


def benchmark(opponent_name: str, num_matches: int = 25, seed: int = 0):
    """Replay the model card's evaluation: N single-point matches, no rendering.

    Args:
        opponent_name: which scripted baseline to face.
        num_matches: how many matches to simulate.
        seed: RNG seed for reproducibility.

    Returns:
        A markdown summary of the win / draw / loss record.
    """
    seed = int(seed)
    rng = random.Random(seed)
    random.seed(seed)
    num_matches = int(num_matches)

    opponent = OPPONENTS[opponent_name](rng=rng)
    env = StandalonePongEnv(PhysicsConfig(), seed=seed)

    wins = draws = losses = 0
    rallies = []
    t0 = time.perf_counter()

    for i in range(1, num_matches + 1):
        obs = env.reset(serve_direction=1 if i % 2 == 0 else -1)
        done = False
        info = {}
        while not done:
            ego_act = policy_action(obs)
            opp_act = opponent.act(env)
            obs, done, info = env.step(ego_action=ego_act, opp_action=opp_act)
        rallies.append(env.rally_count)
        w = info.get("winner")
        if w == "ego":
            wins += 1
        elif w == "opponent":
            losses += 1
        else:
            draws += 1

    el = time.perf_counter() - t0
    n = max(1, num_matches)
    return (
        f"### Mr. Pong vs **{opponent_name}** — {num_matches} matches\n\n"
        f"| Metric | Value |\n|---|---|\n"
        f"| Win rate | **{wins / n * 100:.1f}%** |\n"
        f"| Draw rate | {draws / n * 100:.1f}% |\n"
        f"| Loss rate | {losses / n * 100:.1f}% |\n"
        f"| Record (W/D/L) | {wins}W / {draws}D / {losses}L |\n"
        f"| Average rally | {np.mean(rallies):.1f} hits (max {max(rallies)}) |\n"
        f"| Wall time | {el:.2f} s ({num_matches / max(el, 1e-6):.1f} matches/s) |\n\n"
        f"*A draw means the rally hit the {env.phys.max_rally_steps}-step limit "
        f"without either side scoring.*"
    )


# ----------------------------------------------------------------------------
# UI
# ----------------------------------------------------------------------------
CSS = """
#col-container { max-width: 1000px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

INTRO = f"""# 🏓 Mr. Pong

[**fromziro/MrPong**](https://huggingface.co/{MODEL_ID}) is a {N_PARAMS:,}-parameter
actor–critic MLP trained with PPO and self-play for 10M steps to play 2D table tennis.
It wins ~87% of its matches against the strongest scripted baseline it was trained on.

**Play it yourself below** — the policy runs in your browser, so the paddle answers
in real time. Or watch the real PyTorch checkpoint take on the model-card opponents.
"""

with gr.Blocks(theme=gr.themes.Citrus(), css=CSS, title="Mr. Pong") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(INTRO)

        with gr.Tabs():
            with gr.Tab("🎮 Play Mr. Pong"):
                gr.HTML(GAME_IFRAME)
                gr.Markdown(
                    "Click the court to start. Hold **W** / **↑** and **S** / **↓** "
                    "(or switch to mouse / touch control). First to 5 points.  \n"
                    "*V(s) is the critic's value estimate from Mr. Pong's point of "
                    "view — it climbs when the agent thinks it is about to win the point.*"
                )

            with gr.Tab("🎬 Watch a match"):
                with gr.Row():
                    opponent = gr.Dropdown(
                        list(OPPONENTS.keys()), value="Realistic Hard",
                        label="Opponent", scale=3)
                    watch_btn = gr.Button("Play match", variant="primary", scale=1)
                video_out = gr.Video(label="Match replay", autoplay=True,
                                     height=430, format="mp4")
                stats_out = gr.Markdown()
                with gr.Accordion("Advanced settings", open=False):
                    points = gr.Slider(1, 5, value=3, step=1, label="Points to win")
                    rand_seed = gr.Checkbox(value=True, label="Randomize seed")
                    seed_num = gr.Number(value=0, precision=0, label="Seed")
                gr.Examples(
                    examples=[["Realistic Hard"], ["Impossible Hard"],
                              ["Medium Logic"], ["Random Agent"]],
                    inputs=[opponent],
                    outputs=[video_out, stats_out, seed_num],
                    fn=watch_match,
                    cache_examples=True,
                    cache_mode="lazy",
                    label="Opponents from the model card",
                )
                watch_btn.click(
                    watch_match,
                    inputs=[opponent, points, rand_seed, seed_num],
                    outputs=[video_out, stats_out, seed_num],
                    api_name="watch_match",
                )

            with gr.Tab("📊 Benchmark"):
                with gr.Row():
                    bench_opp = gr.Dropdown(
                        list(OPPONENTS.keys()), value="Realistic Hard",
                        label="Opponent", scale=3)
                    bench_btn = gr.Button("Run benchmark", variant="primary", scale=1)
                bench_out = gr.Markdown()
                with gr.Accordion("Advanced settings", open=False):
                    n_matches = gr.Slider(5, 200, value=25, step=5, label="Matches")
                    bench_seed = gr.Number(value=0, precision=0, label="Seed")
                gr.Markdown(
                    "Reference numbers reported by the author over 1000 matches:  \n"
                    "Easy 99.7% · Medium 99.0% · Realistic Hard 86.6% · "
                    "Random 100% · Impossible Hard 0% W / 98.4% D."
                )
                bench_btn.click(
                    benchmark,
                    inputs=[bench_opp, n_matches, bench_seed],
                    outputs=[bench_out],
                    api_name="benchmark",
                )

        gr.Markdown(
            "Model, physics environment and baseline opponents by "
            "[FromZero](https://huggingface.co/fromziro) (Apache-2.0). "
            "The in-browser policy is the exact checkpoint re-evaluated in JavaScript; "
            "the *Watch* and *Benchmark* tabs run the PyTorch model itself."
        )

if __name__ == "__main__":
    demo.launch(mcp_server=True)