"""
Detection Head Arena — single-domain runner using the heads/ registry.

Usage:
    python arena.py --head cofiber_threshold --steps 2000 --batch 1
    python arena.py --head all --steps 2000 --batch 1
    python arena.py --list
"""

import argparse
import sys

from heads import REGISTRY, ALL_NAMES, get_head


def main():
    parser = argparse.ArgumentParser(description="Detection Head Arena")
    parser.add_argument("--head", type=str, default="all",
                        help=f"Head name or 'all'. Available: {', '.join(ALL_NAMES)}")
    parser.add_argument("--list", action="store_true", help="List available heads")
    parser.add_argument("--steps", type=int, default=2000)
    parser.add_argument("--batch", type=int, default=1)
    args = parser.parse_args()

    if args.list:
        print(f"{len(ALL_NAMES)} heads available:")
        for name in ALL_NAMES:
            head = get_head(name)
            params = sum(p.numel() for p in head.parameters()) / 1e6
            print(f"  {name:<25} {params:.2f}M params")
        return

    if args.head == "all":
        heads = ALL_NAMES
    else:
        heads = [h.strip() for h in args.head.split(",")]

    for name in heads:
        head = get_head(name)
        params = sum(p.numel() for p in head.parameters()) / 1e6
        print(f"{name}: {params:.2f}M params, needs_intermediates={head.needs_intermediates}")


if __name__ == "__main__":
    main()