from __future__ import annotations

import argparse
import json
import time
from dataclasses import asdict, fields, replace

import torch

from config import V2Config
from env import FastFloatingBeaconEnv
from mapgen import configure_motif_mapgen, install_motif_mapgen
from runner import BeaconController
from runtime import sync_device
from settings import ENV_DEFAULTS, MAPGEN_KWARGS


def sync(device: torch.device) -> None:
    sync_device(device)


@torch.no_grad()
def raw_env_sps(env: FastFloatingBeaconEnv, steps: int) -> float:
    actions = torch.zeros((env.n, env.action_size), device=env.device)
    actions[:, 0] = 1.0
    sync(env.device)
    start = time.perf_counter()
    for _ in range(int(steps)):
        env.step(actions)
    sync(env.device)
    return float(env.n * int(steps) / max(time.perf_counter() - start, 1e-9))


@torch.no_grad()
def policy_env_sps(env: FastFloatingBeaconEnv, model: BeaconController, steps: int) -> float:
    obs = env.observe()
    sync(env.device)
    start = time.perf_counter()
    for _ in range(int(steps)):
        action, _ = model.choose_action(obs, deterministic_controls=True, deterministic_buttons=True)
        obs = env.step(action).observation
    sync(env.device)
    return float(env.n * int(steps) / max(time.perf_counter() - start, 1e-9))


def main() -> None:
    parser = argparse.ArgumentParser(description="Benchmark the active v37 Torch parkour stack.")
    parser.add_argument("--device", default="cuda")
    parser.add_argument("--envs", type=int, default=65536)
    parser.add_argument("--steps", type=int, default=128)
    parser.add_argument("--repeats", type=int, default=3)
    parser.add_argument("--route-jumps", type=int, default=8)
    parser.add_argument("--distractors", type=int, default=6)
    parser.add_argument("--hidden", type=int, default=256)
    parser.add_argument("--depth", type=int, default=3)
    args = parser.parse_args()

    configure_motif_mapgen(**MAPGEN_KWARGS)
    install_motif_mapgen()
    config_fields = {field.name for field in fields(V2Config)}
    recipe_kwargs = {key: value for key, value in {**ENV_DEFAULTS, **MAPGEN_KWARGS}.items() if key in config_fields}
    cfg = replace(
        V2Config(),
        device=str(args.device),
        envs=int(args.envs),
        route_jumps=int(args.route_jumps),
        distractors=int(args.distractors),
        sensor_mode="topk",
        sensor_topk=16,
        sensor_token_range=9.5,
        **recipe_kwargs,
    )
    env = FastFloatingBeaconEnv(cfg)
    model = BeaconController(env.obs_size, hidden=int(args.hidden), depth=int(args.depth)).to(env.device).eval()
    raw = [raw_env_sps(env, int(args.steps)) for _ in range(max(1, int(args.repeats)))]
    policy = [policy_env_sps(env, model, int(args.steps)) for _ in range(max(1, int(args.repeats)))]
    result = {
        "device": str(env.device),
        "env": asdict(cfg),
        "raw_env_sps": sum(raw) / len(raw),
        "raw_env_sps_min": min(raw),
        "policy_env_sps": sum(policy) / len(policy),
        "policy_env_sps_min": min(policy),
    }
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()