import argparse
import os
import yaml
import json
from dataclasses import asdict
from typing import Dict

from lkb.benchmark import BenchmarkConfig
from lkb.sweep import sweep_one_algo


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--sweep", type=str, required=True, help="YAML defining sweep_name, base_config, pilot overrides, and grids")
    args = parser.parse_args()

    with open(args.sweep, "r") as f:
        s = yaml.safe_load(f)

    sweep_name = s["sweep_name"]
    base_config_path = s["base_config"]

    with open(base_config_path, "r") as f:
        d = yaml.safe_load(f)

    R_pilot = s.get("R_pilot", 5)
    T_pilot = s.get("T_pilot", None)

    base_cfg = BenchmarkConfig(**d)
    base_cfg.R = R_pilot
    if T_pilot is not None:
        base_cfg.T = int(T_pilot)

    grids: Dict[str, Dict[str, list]] = s["grids"]

    outdir = os.path.join("out", "sweeps", sweep_name)
    os.makedirs(outdir, exist_ok=True)

    recommended = {"algo_params": {}}
    all_rows = []

    for algo_name, grid in grids.items():
        print(f"[SWEEP] Tuning {algo_name} on base={base_config_path} with grid={grid}")
        best, rows = sweep_one_algo(base_cfg, algo_name, grid)
        recommended["algo_params"][algo_name] = best
        all_rows.extend(rows)

    with open(os.path.join(outdir, "recommended.yaml"), "w") as f:
        yaml.safe_dump(recommended, f, sort_keys=False)

    with open(os.path.join(outdir, "results.json"), "w") as f:
        json.dump(all_rows, f, indent=2)

    print(f"[SWEEP] Done. Wrote recommended params to {os.path.join(outdir, 'recommended.yaml')}")
    print(f"[SWEEP] Full results in {os.path.join(outdir, 'results.json')}")


if __name__ == "__main__":
    main()
