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

import argparse
import subprocess
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
from src.stacking_eval import run_stack_task


def main() -> None:
    set_thread_env(1)
    parser = argparse.ArgumentParser(description="Run PA3 individual and stacking tasks locally with joblib.")
    parser.add_argument("--config", default="config/experiment.yaml")
    parser.add_argument("--n-jobs", type=int, default=None)
    parser.add_argument("--skip-individual", action="store_true")
    parser.add_argument("--skip-stacking", action="store_true")
    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)
    n_jobs = args.n_jobs or int(cfg.get("n_jobs", 1))

    if not args.skip_individual:
        subprocess.run([sys.executable, str(ROOT / "scripts" / "00_make_tasks.py"), "--config", str(cfg_path)], check=True)
        tasks = pd.read_csv(dirs["output"] / "tasks.csv").to_dict("records")
        print(f"Running {len(tasks)} individual tasks with n_jobs={n_jobs}")
        Parallel(n_jobs=n_jobs, backend="loky")(delayed(run_task)(cfg, t, dirs) for t in tasks)

    if cfg.get("stacking", {}).get("enabled", False) and not args.skip_stacking:
        subprocess.run([sys.executable, str(ROOT / "scripts" / "00_make_stack_tasks.py"), "--config", str(cfg_path)], check=True)
        stasks = pd.read_csv(dirs["output"] / "stack_tasks.csv").to_dict("records")
        print(f"Running {len(stasks)} stacking tasks with n_jobs={n_jobs}")
        Parallel(n_jobs=n_jobs, backend="loky")(delayed(run_stack_task)(cfg, t, dirs) for t in stasks)

    print("Collecting summary metrics ...")
    subprocess.run([sys.executable, str(ROOT / "scripts" / "02_collect_summary.py"), "--config", str(cfg_path)], check=True)
    print("Generating HTML report ...")
    subprocess.run([sys.executable, str(ROOT / "scripts" / "03_analyze_results.py"), "--config", str(cfg_path)], check=True)


if __name__ == "__main__":
    main()