#!/usr/bin/env python3
"""Run RS-fast / RS-full (spec §7) through the marco model wrapper; see src/lowbit/eval/rs.py."""
from __future__ import annotations

import argparse
import json

from lowbit.eval.rs import run_suite


def main() -> None:
    p = argparse.ArgumentParser()
    p.add_argument("--suite", required=True)
    p.add_argument("--out", required=True)
    p.add_argument("--model-args", required=True, help="lm-eval model_args for the chosen backend")
    p.add_argument("--batch-size", default="16", help="an int, or auto (vLLM schedules everything itself)")
    p.add_argument("--model", default="marco", choices=["marco", "vllm"])
    p.add_argument("--runs", nargs="*", help="run names from the suite (default: all)")
    p.add_argument("--tasks", nargs="*", help="tasks within those runs (default: all)")
    p.add_argument("--limit", type=int, default=None, help="items per task (dry runs only)")
    p.add_argument("--device", default="cuda")
    a = p.parse_args()
    bs = a.batch_size if a.batch_size == "auto" else int(a.batch_size)
    s = run_suite(a.suite, a.out, a.model_args, bs, a.limit, a.device, a.model, a.runs, a.tasks)
    print(json.dumps({k: s[k] for k in ("scores", "timings", "total_s")}, indent=1))


if __name__ == "__main__":
    main()