code/run_a1_end_to_end.py

#!/usr/bin/env python
"""Run the full A1 pipeline end-to-end.

Flow:
1) Bootstrap with feature extraction + alignment cache generation
2) Ridge fitting from cached regressors
3) Visualization of layer-wise and best-layer performance
"""

from __future__ import annotations

import argparse
import csv
import json
from datetime import datetime, timezone
import os
from pathlib import Path
import re
import subprocess
import sys
import threading
import time
from typing import Sequence


def _banner(msg: str) -> None:
    ts = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S UTC")
    width = 72
    print("\n" + "=" * width, flush=True)
    print(f"  [{ts}] {msg}", flush=True)
    print("=" * width, flush=True)


def _gpu_stats() -> str:
    try:
        out = subprocess.check_output(
            ["nvidia-smi",
             "--query-gpu=index,utilization.gpu,memory.used,memory.total,temperature.gpu",
             "--format=csv,noheader,nounits"],
            stderr=subprocess.DEVNULL,
            timeout=5,
        ).decode().strip()
        lines = []
        for row in out.splitlines():
            idx, util, mem_used, mem_total, temp = [x.strip() for x in row.split(",")]
            lines.append(f"GPU{idx}: {util}% util  {mem_used}/{mem_total} MiB  {temp}C")
        return "  ".join(lines)
    except Exception:
        return "(nvidia-smi unavailable)"


def _sys_stats() -> str:
    try:
        import psutil
        cpu = psutil.cpu_percent(interval=None)
        vm = psutil.virtual_memory()
        ram_used = vm.used // (1024 ** 3)
        ram_total = vm.total // (1024 ** 3)
        return f"CPU: {cpu:.0f}%  RAM: {ram_used}/{ram_total} GB"
    except Exception:
        return ""


def _start_utilization_monitor(interval_s: int = 60) -> threading.Event:
    """Print GPU/CPU/RAM stats every `interval_s` seconds in a background thread."""
    stop_evt = threading.Event()

    def _loop() -> None:
        while not stop_evt.wait(timeout=interval_s):
            ts = datetime.now(timezone.utc).strftime("%H:%M:%S UTC")
            gpu = _gpu_stats()
            sys_ = _sys_stats()
            parts = [p for p in [gpu, sys_] if p]
            print(f"[UTILIZATION {ts}] {' | '.join(parts)}", flush=True)

    t = threading.Thread(target=_loop, daemon=True)
    t.start()
    return stop_evt


def _build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Run end-to-end A1 pipeline")
    parser.add_argument(
        "--base-dir",
        type=str,
        default=".",
        help="Repository root (default: current working directory)",
    )
    parser.add_argument(
        "--bootstrap-output-dir",
        type=str,
        default=None,
        help="Output directory for bootstrap artifacts (default: <base-dir>/outputs/a1_bootstrap)",
    )
    parser.add_argument(
        "--model-profile",
        type=str,
        default="current",
        help="Model profile for bootstrap feature extraction",
    )
    parser.add_argument(
        "--model-slug",
        type=str,
        default="all",
        help=(
            "Model slug(s) for fit stage. Use 'all' to run every cached model from the selected "
            "profile (default), 'auto' to pick one cached slug, or a comma-separated list."
        ),
    )
    parser.add_argument(
        "--alpha",
        type=float,
        default=300.0,
        help="Ridge alpha used in fit stage",
    )
    parser.add_argument(
        "--protocols",
        type=str,
        default="C",
        help="Fit protocols to execute: C (bootstrap generates protocol_c_cross_subject_folds.csv automatically).",
    )
    parser.add_argument(
        "--metric",
        type=str,
        default="mean_corr",
        choices=[
            "mean_corr",
            "mean_r2",
            "mean_2v2_accuracy",
            "2v2",
            "2v2_accuracy",
            "two_v_two_accuracy",
        ],
        help="Metric for visualization",
    )
    parser.add_argument(
        "--allowed-runs",
        type=str,
        default="1,2,3,4",
        help="Run whitelist forwarded to bootstrap",
    )
    parser.add_argument(
        "--exclude-subjects",
        type=str,
        default="sub-03,sub-18",
        help="Excluded subjects forwarded to bootstrap",
    )
    parser.add_argument(
        "--feature-num-workers",
        type=str,
        default="auto",
        help=(
            "Number of GPU workers for feature extraction (forwarded to bootstrap). "
            "'auto' uses every visible CUDA device; recommended on 4xA10G."
        ),
    )
    parser.add_argument(
        "--num-fit-workers",
        type=str,
        default="auto",
        help=(
            "Number of CPU workers for layer-parallel fit (forwarded to run_a1_fit.py). "
            "'auto' uses min(n_layers, cpu_count // 2). Use 1 to force serial."
        ),
    )
    parser.add_argument(
        "--target-mask-mode",
        type=str,
        default="run_top10",
        choices=["run_top10", "run_top25", "core_roi"],
        help=(
            "Target mask family for run_a1_fit.py. 'run_top10' uses the Swati ISC top-10%% masks "
            "per canonical run, 'run_top25' keeps the 25%% mask option, and 'core_roi' keeps the "
            "legacy 7-ROI evaluation."
        ),
    )
    parser.add_argument(
        "--run-mask-dir",
        "--run-top10-mask-dir",
        "--run-top25-mask-dir",
        dest="run_mask_dir",
        type=str,
        default=None,
        help=(
            "Optional override for the Swati ISC run-conditioned mask directory passed to "
            "run_a1_fit.py. Can point to the output root or directly to isc_group/."
        ),
    )
    parser.add_argument(
        "--reuse-caches",
        action="store_true",
        help="Reuse existing feature/alignment caches (disables overwrite flags)",
    )
    parser.add_argument(
        "--skip-bootstrap",
        action="store_true",
        help="Skip bootstrap stage",
    )
    parser.add_argument(
        "--skip-fit",
        action="store_true",
        help="Skip fit stage",
    )
    parser.add_argument(
        "--skip-visualize",
        action="store_true",
        help="Skip visualize stage",
    )
    parser.add_argument(
        "--fit-output-dir",
        type=str,
        default=None,
        help="Fit output directory (default: <bootstrap-output-dir>/fit_results/<model-slug>)",
    )
    parser.add_argument(
        "--results-repo",
        type=str,
        default=None,
        help="Optional HF dataset repo for incremental result uploads during the run",
    )
    parser.add_argument(
        "--results-path",
        type=str,
        default="hf_jobs/latest",
        help="Base path inside --results-repo for incremental uploads",
    )
    parser.add_argument(
        "--results-token-env",
        type=str,
        default="HF_TOKEN",
        help="Environment variable containing the HF token used for incremental uploads",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Print commands without executing",
    )
    return parser


def _quote_for_log(args: Sequence[str]) -> str:
    quoted: list[str] = []
    for value in args:
        if any(ch.isspace() for ch in value):
            quoted.append(f'"{value}"')
        else:
            quoted.append(value)
    return " ".join(quoted)


def _run(cmd: list[str], dry_run: bool) -> None:
    print(f"[run] {_quote_for_log(cmd)}")
    if dry_run:
        return
    subprocess.run(cmd, check=True)


def _join_results_repo_path(*parts: str) -> str:
    tokens = [str(part).strip("/") for part in parts if str(part).strip("/")]
    return "/".join(tokens)


def _upload_results_folder(
    folder_path: Path,
    repo_id: str,
    path_in_repo: str,
    token_env: str,
    dry_run: bool,
    commit_message: str,
) -> None:
    if not folder_path.exists():
        raise FileNotFoundError(f"Cannot upload missing folder: {folder_path}")

    print(f"[upload] {folder_path} -> {repo_id}:{path_in_repo}", flush=True)
    if dry_run:
        return

    token = os.getenv(str(token_env).strip())
    if not token:
        raise ValueError(
            f"Results upload requested but env var {token_env!r} is not set"
        )

    from huggingface_hub import create_repo, upload_folder

    create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, token=token)
    upload_folder(
        folder_path=str(folder_path),
        repo_id=repo_id,
        repo_type="dataset",
        path_in_repo=path_in_repo,
        token=token,
        commit_message=commit_message,
    )


def _upload_results_file(
    file_path: Path,
    repo_id: str,
    path_in_repo: str,
    token_env: str,
    dry_run: bool,
    commit_message: str,
) -> None:
    if not file_path.exists():
        raise FileNotFoundError(f"Cannot upload missing file: {file_path}")

    print(f"[upload] {file_path} -> {repo_id}:{path_in_repo}", flush=True)
    if dry_run:
        return

    token = os.getenv(str(token_env).strip())
    if not token:
        raise ValueError(
            f"Results upload requested but env var {token_env!r} is not set"
        )

    from huggingface_hub import create_repo, upload_file

    create_repo(repo_id=repo_id, repo_type="dataset", exist_ok=True, token=token)
    upload_file(
        path_or_fileobj=str(file_path),
        repo_id=repo_id,
        repo_type="dataset",
        path_in_repo=path_in_repo,
        token=token,
        commit_message=commit_message,
    )


def _slugify_model_id(model_id: str) -> str:
    cleaned = re.sub(r"[^a-zA-Z0-9._-]+", "_", model_id.strip())
    return cleaned.strip("_") or "unknown_model"


def _read_available_model_slugs(bootstrap_output_dir: Path) -> list[str]:
    csv_summary_path = bootstrap_output_dir / "csv" / "alignment_regressor_summary.csv"
    summary_path = csv_summary_path if csv_summary_path.exists() else (bootstrap_output_dir / "alignment_regressor_summary.csv")
    if not summary_path.exists():
        return []

    slugs: set[str] = set()
    with summary_path.open("r", encoding="utf-8", newline="") as handle:
        reader = csv.DictReader(handle)
        for row in reader:
            value = str(row.get("model_slug", "")).strip()
            if value:
                slugs.add(value)
    return sorted(slugs)


def _resolve_fit_model_slug(
    requested_model_slug: str,
    bootstrap_output_dir: Path,
) -> str:
    available = _read_available_model_slugs(bootstrap_output_dir=bootstrap_output_dir)
    if not available:
        csv_summary_path = bootstrap_output_dir / "csv" / "alignment_regressor_summary.csv"
        summary_path = (
            csv_summary_path
            if csv_summary_path.exists()
            else (bootstrap_output_dir / "alignment_regressor_summary.csv")
        )
        raise FileNotFoundError(
            "Cannot resolve fit model slug because no cached regressor summary was found. "
            f"Expected file: {summary_path}"
        )

    requested = str(requested_model_slug).strip()
    if requested.lower() != "auto":
        if requested not in available:
            raise ValueError(
                f"Requested model_slug={requested} is unavailable in cached regressors. "
                f"Available slugs: {available}"
            )
        return requested

    bootstrap_summary_path = bootstrap_output_dir / "bootstrap_summary.json"
    if bootstrap_summary_path.exists():
        with bootstrap_summary_path.open("r", encoding="utf-8") as handle:
            bootstrap_summary = json.load(handle)
        preferred_ids = bootstrap_summary.get("models_locked", []) or []
        for model_id in preferred_ids:
            candidate = _slugify_model_id(str(model_id))
            if candidate in available:
                return candidate

    return available[0]


def main() -> None:
    args = _build_parser().parse_args()

    base_dir = Path(args.base_dir).resolve()
    code_dir = base_dir / "code"

    bootstrap_output_dir = (
        Path(args.bootstrap_output_dir).resolve()
        if args.bootstrap_output_dir
        else (base_dir / "outputs" / "a1_bootstrap")
    )

    bootstrap_script = code_dir / "run_a1_bootstrap.py"
    fit_script = code_dir / "run_a1_fit.py"
    visualize_script = code_dir / "run_a1_visualize.py"
    compare_script = code_dir / "run_a1_compare_models.py"

    for script in [bootstrap_script, fit_script, visualize_script]:
        if not script.exists():
            raise FileNotFoundError(f"Required script not found: {script}")

    py = sys.executable

    _banner("A1 end-to-end pipeline starting")
    print(f"  base_dir            : {base_dir}", flush=True)
    print(f"  bootstrap_output_dir: {bootstrap_output_dir}", flush=True)
    print(f"  model_profile       : {args.model_profile}", flush=True)
    print(f"  protocols           : {args.protocols}", flush=True)
    print(f"  alpha               : {args.alpha}", flush=True)
    print(f"  allowed_runs        : {args.allowed_runs}", flush=True)
    print(f"  exclude_subjects    : {args.exclude_subjects}", flush=True)
    print(f"  target_mask_mode    : {args.target_mask_mode}", flush=True)
    if args.run_mask_dir:
        print(f"  run_mask_dir        : {args.run_mask_dir}", flush=True)
    if args.results_repo:
        print(f"  results_repo        : {args.results_repo}", flush=True)
        print(f"  results_path        : {args.results_path}", flush=True)
    print(flush=True)
    print(f"  Initial GPU state:", flush=True)
    print(f"  {_gpu_stats()}", flush=True)
    print(f"  {_sys_stats()}", flush=True)

    _stop_monitor = _start_utilization_monitor(interval_s=60)

    bootstrap_cmd = [
        py,
        str(bootstrap_script),
        "--base-dir",
        str(base_dir),
        "--output-dir",
        str(bootstrap_output_dir),
        "--allowed-runs",
        str(args.allowed_runs),
        "--exclude-subjects",
        str(args.exclude_subjects),
        "--feature-num-workers",
        str(args.feature_num_workers),
        "--run-feature-extraction",
        "--run-alignment",
    ]
    bootstrap_cmd.extend(["--model-profile", str(args.model_profile)])

    if not args.reuse_caches:
        bootstrap_cmd.extend(["--feature-overwrite", "--alignment-overwrite"])

    if not args.skip_bootstrap:
        _banner("Bootstrap (feature extraction + alignment cache)")
        _run(bootstrap_cmd, dry_run=bool(args.dry_run))
        _banner("Bootstrap COMPLETE")
    else:
        print("[SKIP] Bootstrap stage skipped.", flush=True)

    requested_token = str(args.model_slug).strip()
    requested_lower = requested_token.lower()
    multi_mode = requested_lower == "all" or "," in requested_token

    resolved_model_slugs: list[str] = []
    if not args.skip_fit and not args.dry_run:
        if multi_mode:
            available = _read_available_model_slugs(bootstrap_output_dir=bootstrap_output_dir)
            if not available:
                raise FileNotFoundError(
                    "No cached regressors found; cannot resolve model slugs for multi-model fit."
                )
            if requested_lower == "all":
                resolved_model_slugs = list(available)
            else:
                requested_list = [s.strip() for s in requested_token.split(",") if s.strip()]
                missing = [s for s in requested_list if s not in available]
                if missing:
                    raise ValueError(
                        f"Requested model slugs {missing} not available. Available: {available}"
                    )
                resolved_model_slugs = requested_list
        else:
            resolved_model_slugs = [
                _resolve_fit_model_slug(
                    requested_model_slug=requested_token,
                    bootstrap_output_dir=bootstrap_output_dir,
                )
            ]
    else:
        resolved_model_slugs = [requested_token or "auto"]

    if args.fit_output_dir and len(resolved_model_slugs) > 1:
        raise ValueError(
            "--fit-output-dir cannot be combined with multi-model fitting; "
            "each slug needs its own output directory."
        )

    fit_dirs: list[Path] = []
    fit_cmd = visualize_cmd = []  # will be overwritten; kept for summary fallback

    for slug in resolved_model_slugs:
        slug_fit_dir = (
            Path(args.fit_output_dir).resolve()
            if (args.fit_output_dir and len(resolved_model_slugs) == 1)
            else (bootstrap_output_dir / "fit_results" / str(slug))
        )
        fit_dirs.append(slug_fit_dir)

        fit_cmd = [
            py, str(fit_script),
            "--bootstrap-output-dir", str(bootstrap_output_dir),
            "--model-slug", str(slug),
            "--alpha", str(args.alpha),
            "--protocols", str(args.protocols),
            "--output-dir", str(slug_fit_dir),
            "--num-fit-workers", str(args.num_fit_workers),
            "--target-mask-mode", str(args.target_mask_mode),
        ]
        if args.run_mask_dir:
            fit_cmd.extend(["--run-mask-dir", str(Path(args.run_mask_dir).resolve())])
        visualize_cmd = [
            py, str(visualize_script),
            "--fit-output-dir", str(slug_fit_dir),
            "--metric", str(args.metric),
        ]

        if not args.skip_fit:
            _banner(f"Fit (model_slug={slug}, alpha={args.alpha}, protocols={args.protocols})")
            _run(fit_cmd, dry_run=bool(args.dry_run))
            _banner(f"Fit COMPLETE for {slug}")
        else:
            print(f"[SKIP] Fit stage skipped for {slug}.", flush=True)

        if not args.skip_visualize:
            _banner(f"Visualize ({slug})")
            _run(visualize_cmd, dry_run=bool(args.dry_run))
            _banner(f"Visualize COMPLETE for {slug}")
        else:
            print(f"[SKIP] Visualize stage skipped for {slug}.", flush=True)

        if args.results_repo:
            _upload_results_folder(
                folder_path=slug_fit_dir,
                repo_id=str(args.results_repo),
                path_in_repo=_join_results_repo_path(str(args.results_path), "fit_results", str(slug)),
                token_env=str(args.results_token_env),
                dry_run=bool(args.dry_run),
                commit_message=f"Upload fit results for {slug}",
            )

    # Cross-model comparison when more than one slug was fit.
    compare_output_dir: Path | None = None
    if (
        len(resolved_model_slugs) > 1
        and not args.skip_visualize
        and compare_script.exists()
    ):
        compare_output_dir = (
            Path(args.compare_output_dir).resolve()
            if args.compare_output_dir
            else (bootstrap_output_dir / f"compare_{args.model_profile}")
        )
        compare_cmd = [
            py, str(compare_script),
            "--output-dir", str(compare_output_dir),
            "--metric", str(args.metric),
            "--title-suffix", f"profile={args.model_profile}",
        ]
        for slug, fdir in zip(resolved_model_slugs, fit_dirs):
            compare_cmd.extend(["--fit-dir", f"{slug}={fdir}"])
        _banner("Cross-model comparison plots")
        _run(compare_cmd, dry_run=bool(args.dry_run))
        _banner("Cross-model comparison COMPLETE")
        if args.results_repo:
            _upload_results_folder(
                folder_path=compare_output_dir,
                repo_id=str(args.results_repo),
                path_in_repo=_join_results_repo_path(str(args.results_path), compare_output_dir.name),
                token_env=str(args.results_token_env),
                dry_run=bool(args.dry_run),
                commit_message="Upload cross-model comparison outputs",
            )

    # ---- Summary ----
    resolved_model_slug = resolved_model_slugs[0]
    fit_output_dir = fit_dirs[0]

    summary = {
        "timestamp_utc": datetime.now(timezone.utc).isoformat(),
        "base_dir": str(base_dir),
        "bootstrap_output_dir": str(bootstrap_output_dir),
        "fit_output_dir": str(fit_output_dir),
        "fit_output_dirs": [str(d) for d in fit_dirs],
        "compare_output_dir": str(compare_output_dir) if compare_output_dir else None,
        "model_profile": str(args.model_profile),
        "model_slug_requested": str(args.model_slug),
        "model_slug_resolved": str(resolved_model_slug),
        "model_slugs_resolved": [str(s) for s in resolved_model_slugs],
        "alpha": float(args.alpha),
        "protocols": str(args.protocols),
        "metric": str(args.metric),
        "reuse_caches": bool(args.reuse_caches),
        "skip_bootstrap": bool(args.skip_bootstrap),
        "skip_fit": bool(args.skip_fit),
        "skip_visualize": bool(args.skip_visualize),
        "dry_run": bool(args.dry_run),
    }

    _stop_monitor.set()  # stop the utilization monitor thread

    fit_output_dir.mkdir(parents=True, exist_ok=True)
    summary_path = fit_output_dir / "end_to_end_summary.json"
    with summary_path.open("w", encoding="utf-8") as handle:
        json.dump(summary, handle, indent=2)

    if args.results_repo:
        _upload_results_file(
            file_path=summary_path,
            repo_id=str(args.results_repo),
            path_in_repo=_join_results_repo_path(str(args.results_path), "end_to_end_summary.json"),
            token_env=str(args.results_token_env),
            dry_run=bool(args.dry_run),
            commit_message="Upload end-to-end summary",
        )

    print("=" * 72)
    print("A1 end-to-end pipeline complete")
    print(f"Bootstrap output directory: {bootstrap_output_dir}")
    print(f"Fit output directory: {fit_output_dir}")
    print(f"Summary: {summary_path}")
    print("=" * 72)


if __name__ == "__main__":
    main()