gpu/single-gpu/home_benjamin_hdbasin/code/hdbasin/aggregate.py

from __future__ import annotations

import argparse
from pathlib import Path

import pandas as pd

from .experiments import write_outputs


def _read_csv_if_exists(path: Path) -> pd.DataFrame | None:
    if not path.exists():
        return None
    df = pd.read_csv(path)
    if path.name == "evaluations.csv":
        if "gpu_info" not in df.columns:
            df["gpu_info"] = ""
        if "gpu_name" not in df.columns:
            df["gpu_name"] = df["gpu_info"].fillna("")
        if "gpu_available" not in df.columns:
            gpu_name = df["gpu_name"].fillna("").astype(str)
            df["gpu_available"] = gpu_name.ne("") & ~gpu_name.str.lower().isin(["cpu", "cpu_sklearn"])
    df["source_run"] = path.parts[-3] if len(path.parts) >= 3 else path.parent.name
    return df


def _candidate_run_dirs(runs_root: Path) -> list[Path]:
    candidates: set[Path] = set()
    for parent in runs_root.iterdir():
        if not parent.is_dir():
            continue
        if (parent / "raw" / "evaluations.csv").exists() or (parent / "processed" / "real_data_status.csv").exists():
            candidates.add(parent)
        for child in parent.iterdir():
            if child.is_dir() and (
                (child / "raw" / "evaluations.csv").exists() or (child / "processed" / "real_data_status.csv").exists()
            ):
                candidates.add(child)
    return sorted(candidates)


def aggregate_runs(run_dirs: list[Path], outdir: Path) -> None:
    eval_frames = []
    status_frames = []
    pairwise_frames = []
    certificate_frames = []
    for run_dir in run_dirs:
        eval_df = _read_csv_if_exists(run_dir / "raw" / "evaluations.csv")
        if eval_df is not None:
            eval_frames.append(eval_df)
        status_df = _read_csv_if_exists(run_dir / "processed" / "real_data_status.csv")
        if status_df is not None:
            status_frames.append(status_df)
        pairwise_df = _read_csv_if_exists(run_dir / "processed" / "pairwise_score_validation.csv")
        if pairwise_df is not None:
            pairwise_frames.append(pairwise_df)
        certificate_df = _read_csv_if_exists(run_dir / "processed" / "certificate_validation.csv")
        if certificate_df is not None:
            certificate_frames.append(certificate_df)

    evaluations = pd.concat(eval_frames, ignore_index=True) if eval_frames else None
    real_status = pd.concat(status_frames, ignore_index=True) if status_frames else None
    pairwise = pd.concat(pairwise_frames, ignore_index=True) if pairwise_frames else None
    certificates = pd.concat(certificate_frames, ignore_index=True) if certificate_frames else None
    write_outputs(outdir, evaluations, real_status, pairwise, certificates)

    manifest = outdir / "processed" / "source_runs.csv"
    manifest.parent.mkdir(parents=True, exist_ok=True)
    pd.DataFrame({"source_run": [p.name for p in run_dirs], "path": [str(p) for p in run_dirs]}).to_csv(manifest, index=False)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Aggregate existing HD-BasinFlow run folders into one report.")
    parser.add_argument("--runs-root", type=Path, default=Path("runs"))
    parser.add_argument("--include", default="")
    parser.add_argument("--outdir", type=Path, default=Path("runs/combined_report"))
    args = parser.parse_args(argv)

    if args.include:
        run_dirs = [args.runs_root / name.strip() for name in args.include.split(",") if name.strip()]
    else:
        run_dirs = _candidate_run_dirs(args.runs_root)
    missing = [str(p) for p in run_dirs if not p.exists()]
    if missing:
        raise SystemExit(f"Missing run directories: {missing}")
    aggregate_runs(run_dirs, args.outdir)
    print(f"Aggregated {len(run_dirs)} runs into {args.outdir}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())