"""Run and resume the predeclared fold-local update-budget selection."""

from __future__ import annotations

import csv
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
CANDIDATES = {
    500: (ROOT / "configs/development_500.yaml", ROOT / "artifacts/development_update_500_v1"),
    1000: (ROOT / "configs/development_1000.yaml", ROOT / "artifacts/development_update_1000_v1"),
    2000: (ROOT / "configs/development_2000.yaml", ROOT / "artifacts/development_update_2000_v1"),
}
SOURCE_PREPARE = CANDIDATES[500][1] / "prepare"


def run(*arguments: str) -> None:
    command = [sys.executable, str(ROOT / "run_experiment.py"), *arguments]
    subprocess.run(command, cwd=ROOT, check=True)


def link_prepare(destination: Path) -> None:
    target = destination / "prepare"
    if target.exists():
        return
    destination.mkdir(parents=True, exist_ok=True)
    shutil.copytree(SOURCE_PREPARE, target, copy_function=os.link)


def main() -> None:
    if not (SOURCE_PREPARE / "complete.json").is_file():
        raise FileNotFoundError("Run the 500-update preparation before development selection")
    for _, destination in CANDIDATES.values():
        link_prepare(destination)

    for fold in range(3):
        for updates, (config, destination) in CANDIDATES.items():
            common = ["--config", str(config), "--artifact-dir", str(destination),
                      "--fold", str(fold), "--condition", "M50", "--seed", "42"]
            print(f"DEVELOPMENT fold={fold} updates={updates}", flush=True)
            run("train", *common, "--resume")
            run("extract", *common, "--role", "select")
            run("validate", *common)

    rows = []
    selected = {}
    for fold in range(3):
        fold_rows = []
        for updates, (_, destination) in CANDIDATES.items():
            path = destination / "runs" / f"fold_{fold}" / "M50" / "seed_42" / "validation.json"
            result = json.loads(path.read_text(encoding="utf-8"))
            row = {"fold": fold, "updates": updates, "select_AUROC": result["verification_AUROC"]}
            rows.append(row)
            fold_rows.append(row)
        best = max(row["select_AUROC"] for row in fold_rows)
        selected[fold] = min(row["updates"] for row in fold_rows if best - row["select_AUROC"] <= 0.001)

    output = ROOT / "artifacts" / "development_selection_v1"
    output.mkdir(parents=True, exist_ok=True)
    with (output / "results.csv").open("w", newline="", encoding="utf-8") as stream:
        writer = csv.DictWriter(stream, fieldnames=["fold", "updates", "select_AUROC"])
        writer.writeheader()
        writer.writerows(rows)
    result = {"selection_rule": "highest select AUROC; smallest budget within 0.001 of best",
              "selected_updates_by_fold": [selected[fold] for fold in range(3)], "results": rows,
              "roles_accessed": ["fit", "select"], "operate_or_test_accessed": False}
    (output / "selection.json").write_text(json.dumps(result, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(result, indent=2), flush=True)


if __name__ == "__main__":
    main()