#!/usr/bin/env python3
from __future__ import annotations

import argparse
from pathlib import Path
import sys

import pandas as pd
from joblib import Parallel, delayed

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from src.config_utils import prepare_output_dirs, read_config, set_thread_env
from src.train_eval import run_task


def main() -> None:
    set_thread_env(1)
    parser = argparse.ArgumentParser(description="Run all tasks locally with joblib parallelism.")
    parser.add_argument("--config", default="config/experiment.yaml")
    parser.add_argument("--n-jobs", type=int, default=None)
    args = parser.parse_args()

    cfg_path = Path(args.config).resolve()
    cfg = read_config(cfg_path)
    dirs = prepare_output_dirs(cfg, base_dir=cfg_path.parent.parent)
    task_file = dirs["output"] / "tasks.csv"
    if not task_file.exists():
        from subprocess import check_call
        check_call([sys.executable, str(ROOT / "scripts" / "00_make_tasks.py"), "--config", str(cfg_path)])
    tasks = pd.read_csv(task_file).to_dict("records")
    n_jobs = args.n_jobs or int(cfg.get("n_jobs", 1))
    print(f"Running {len(tasks)} tasks with n_jobs={n_jobs}")
    results = Parallel(n_jobs=n_jobs, backend="loky")(
        delayed(run_task)(cfg, task, dirs) for task in tasks
    )
    for res in results:
        print(res)


if __name__ == "__main__":
    main()