"""Benchmark LiDAR perception improvements.

Measures success rate, mean steps, and tracker quality metrics for
demo collection across multiple seeds. Compares configurations.

Usage::

    python -m tools.benchmark_lidar --n-sheep 5 --seeds 15
    HERDING_WORLD=field_round python -m tools.benchmark_lidar --n-sheep 5
"""

from __future__ import annotations

import argparse
import time
from collections import Counter

from training.bc.collect import collect_one
from herding.control.universal import compute_action


def run_benchmark(n_sheep: int, n_seeds: int, max_steps: int = 100000,
                  drive_mode: str = "differential"):
    results = []
    t0 = time.time()
    for seed in range(n_seeds):
        obs, actions, success, steps = collect_one(
            n_sheep, seed, max_steps, 5, compute_action,
            frame_stack=1, privileged=False, drive_mode=drive_mode,
        )
        results.append({
            "seed": seed,
            "success": success,
            "steps": steps,
            "logged": len(obs),
        })
        tag = "+" if success else "x"
        print(f"  [{tag}] seed={seed:>2d}  steps={steps:>6d}")
    elapsed = time.time() - t0

    successes = [r for r in results if r["success"]]
    failures = [r for r in results if not r["success"]]
    n_ok = len(successes)
    rate = 100.0 * n_ok / len(results)

    mean_steps_ok = (sum(r["steps"] for r in successes) / n_ok) if n_ok else 0
    mean_steps_all = sum(r["steps"] for r in results) / len(results)

    print(f"\n  Results: {n_ok}/{len(results)} success ({rate:.0f}%)")
    print(f"  Mean steps (success): {mean_steps_ok:>8.0f}")
    print(f"  Mean steps (all):     {mean_steps_all:>8.0f}")
    print(f"  Elapsed: {elapsed:.0f}s")
    return {
        "n_sheep": n_sheep,
        "n_seeds": n_seeds,
        "success_rate": rate,
        "n_success": n_ok,
        "mean_steps_success": mean_steps_ok,
        "mean_steps_all": mean_steps_all,
        "elapsed_s": elapsed,
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--n-sheep", type=int, default=5)
    parser.add_argument("--seeds", type=int, default=15)
    parser.add_argument("--max-steps", type=int, default=100000)
    parser.add_argument("--drive-mode", default="differential",
                        choices=["differential", "mecanum"])
    args = parser.parse_args()

    from herding.world.geometry import FIELD_SHAPE
    print(f"[bench] world={FIELD_SHAPE}  n_sheep={args.n_sheep}  "
          f"seeds={args.seeds}  drive={args.drive_mode}")
    print()
    result = run_benchmark(args.n_sheep, args.seeds, args.max_steps,
                           args.drive_mode)
    print()
    print("[bench] summary:", result)


if __name__ == "__main__":
    main()