import argparse
import os
import yaml
import json

from lkb.benchmark import BenchmarkConfig
from lkb.tpe import tpe_discrete


def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--tpe", type=str, required=True, help="YAML with sweep_name, base_config, algo, search_space, and tpe settings")
    args = parser.parse_args()

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

    sweep_name = s["sweep_name"]
    base_config_path = s["base_config"]
    algo = s["algo"]
    space = s["search_space"]
    tpe_cfg = s.get("tpe", {})
    R_pilot = s.get("R_pilot", 5)
    T_pilot = s.get("T_pilot", None)

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

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

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

    best, trials = tpe_discrete(
        cfg,
        algo,
        space,
        n_init=tpe_cfg.get("n_init", 8),
        n_iter=tpe_cfg.get("n_iter", 32),
        gamma=tpe_cfg.get("gamma", 0.2),
        seed=tpe_cfg.get("seed", 0),
    )

    rec = {"algo_params": {algo: best}}
    with open(os.path.join(outdir, "recommended.yaml"), "w") as f:
        yaml.safe_dump(rec, f, sort_keys=False)

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

    print(f"[TPE] Done. Recommended params for {algo} written to {os.path.join(outdir, 'recommended.yaml')}")
    print(f"[TPE] Trials written to {os.path.join(outdir, 'trials.json')}")


if __name__ == "__main__":
    main()
