# SPDX-License-Identifier: BSD-3-Clause

"""Measure how many episodes the environment itself can supply.

Answers the question a training or eval run actually needs: is the environment
the bottleneck, or the model? Runs an agentic-shaped episode — reset, three
looks, optionally a pin, a guess — across a range of worker counts, both as
threads and as processes, and reports episodes and steps per second.

Every worker gets its own environment, because an episode is stateful; sharing
one interleaves resets and silently produces nonsense. A perfect guess is
asserted to score zero distance, so an interference bug fails loudly instead of
reporting excellent throughput.

Usage:
    python scripts/benchmark_throughput.py threads 32
    REVEAL=0 PIN=0 python scripts/benchmark_throughput.py threads 64
    python scripts/benchmark_throughput.py processes 32

Environment variables:
    REVEAL=0   skip the guess reveal map, as a training run would
    PIN=0      skip the pin, leaving view renders only
"""

import concurrent.futures as cf
import os
import pathlib
import statistics
import sys
import threading
import time

_LOCAL = threading.local()

_ROOT = pathlib.Path(__file__).resolve().parents[2]
sys.path.insert(0, str(_ROOT))
sys.path.insert(0, str(_ROOT.parent / "src"))

INDEX = str(_ROOT / "geoguesser_env" / "tasks" / "pano_v1.jsonl")
CACHE = str(_ROOT / "geoguesser_env" / "data" / "panos")


def episode(task_index: int) -> tuple[float, int]:
    """One agentic-shaped episode: reset, 3 looks, 1 pin, 1 guess."""
    from geoguesser_env.models import GuessAction, LookAction, PinAction, to_wire
    from geoguesser_env.server.geoguesser_environment import GeoGuesserEnvironment
    from geoguesser_env.server.render import minimap

    minimap.set_street_detail(False)  # training config: never touch Overpass
    # An episode is stateful, so every worker needs its own environment. Sharing
    # one through a module global interleaves resets between threads and silently
    # produces nonsense — which is exactly what the first version of this
    # benchmark measured.
    if not hasattr(_LOCAL, "env"):
        _LOCAL.env = GeoGuesserEnvironment(
            index_path=INDEX,
            cache_dir=CACHE,
            allow_fetch=False,
            view_size=640,
            reveal_map=os.environ.get("REVEAL", "1") == "1",
        )
    env = _LOCAL.env
    started = time.perf_counter()
    env.reset(task_index=task_index)
    steps = 1
    for heading in (0, 120, 240):
        env.step(to_wire(LookAction(heading_deg=heading, fov_deg=90)))
        steps += 1
    if os.environ.get("PIN", "1") == "1":
        env.step(to_wire(PinAction(lat=10.0, lon=10.0, span_deg=7.0)))
        steps += 1
    truth = env._task.truth
    result = env.step(to_wire(GuessAction(lat=truth[0], lon=truth[1])))
    steps += 1
    # A perfect guess must score ~1 minus costs. If workers interfered, this
    # fails instead of quietly reporting a great throughput number.
    assert result.distance_km is not None and result.distance_km < 0.01, (
        f"cross-worker interference: task {task_index} scored "
        f"{result.distance_km} km on a perfect guess"
    )
    return time.perf_counter() - started, steps


def run(pool_cls, workers: int, episodes: int) -> dict:
    tasks = [i % 100 for i in range(episodes)]
    wall0 = time.perf_counter()
    with pool_cls(max_workers=workers) as pool:
        results = list(pool.map(episode, tasks))
    wall = time.perf_counter() - wall0
    latencies = [r[0] for r in results]
    total_steps = sum(r[1] for r in results)
    return {
        "workers": workers,
        "episodes": episodes,
        "wall_s": wall,
        "episodes_per_s": episodes / wall,
        "steps_per_s": total_steps / wall,
        "median_episode_s": statistics.median(latencies),
    }


if __name__ == "__main__":
    kind = sys.argv[1]
    episodes = int(sys.argv[2]) if len(sys.argv) > 2 else 24
    pool = cf.ThreadPoolExecutor if kind == "threads" else cf.ProcessPoolExecutor
    print(f"{kind}, {os.cpu_count()} cpus, warm cache, street detail off")
    print(f"{'workers':>7} {'eps/s':>7} {'steps/s':>8} {'wall':>7} {'median ep':>10}")
    for workers in (1, 2, 4, 8, 16):
        r = run(pool, workers, episodes)
        print(
            f"{r['workers']:>7} {r['episodes_per_s']:>7.2f} {r['steps_per_s']:>8.1f} "
            f"{r['wall_s']:>6.1f}s {r['median_episode_s']:>9.2f}s"
        )