#!/usr/bin/env python3
"""Validate grouped holdout outputs for galaxy-level leakage."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import pandas as pd


def parse_args() -> argparse.Namespace:
    root = Path(__file__).resolve().parents[1]
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--predictions", type=Path, default=root / "outputs/tables/heldout_predictions.csv")
    p.add_argument("--assignments", type=Path, default=root / "outputs/tables/fold_assignments.csv")
    return p.parse_args()


def main() -> int:
    args = parse_args()
    pred = pd.read_csv(args.predictions)
    assign = pd.read_csv(args.assignments)
    if pred.empty or assign.empty:
        print("FAIL: holdout outputs are empty", file=sys.stderr)
        return 2
    errors = []
    per_model = pred.groupby(["model", "galaxy_id"])["fold"].nunique()
    bad = per_model[per_model != 1]
    if not bad.empty:
        errors.append(f"Galaxy appears in multiple held-out folds within a model: {bad.to_dict()}")
    fold_map = assign.groupby("galaxy_id")["fold"].nunique()
    if (fold_map != 1).any():
        errors.append("Fold assignment table maps at least one galaxy to multiple folds")
    expected_rows = pred.groupby("model").size()
    if expected_rows.nunique() != 1:
        errors.append(f"Models do not have the same held-out row count: {expected_rows.to_dict()}")
    if errors:
        print("GROUPED HOLDOUT VALIDATION FAILED", file=sys.stderr)
        print("\n".join(errors), file=sys.stderr)
        return 2
    print(f"PASS: {assign['galaxy_id'].nunique()} galaxies are held out as indivisible groups")
    return 0


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