from __future__ import annotations
from typing import Any
import numpy as np
import pandas as pd
from sklearn.model_selection import GroupKFold


def build_group_folds(data: pd.DataFrame, config: dict[str, Any]):
    group_col = config["columns"]["group"]
    groups = data[group_col].astype(str).to_numpy()
    n_splits = int(config["cross_validation"]["n_splits"])
    seed = int(config["random_seed"])
    splitter = GroupKFold(n_splits=n_splits, shuffle=True, random_state=seed)
    folds = list(splitter.split(data, groups=groups))

    registry_rows = []
    seen_test = set()
    for fold_id, (train_idx, test_idx) in enumerate(folds, start=1):
        train_groups = set(groups[train_idx])
        test_groups = set(groups[test_idx])
        overlap = train_groups & test_groups
        if overlap:
            raise RuntimeError(f"Group leakage in fold {fold_id}: {sorted(overlap)}")
        for galaxy in sorted(test_groups):
            if galaxy in seen_test:
                raise RuntimeError(f"Galaxy {galaxy} appears in multiple test folds.")
            seen_test.add(galaxy)
            registry_rows.append({
                "galaxy_id": galaxy,
                "test_fold": fold_id,
                "radial_rows": int((groups == galaxy).sum()),
            })
    all_groups = set(groups)
    if seen_test != all_groups:
        raise RuntimeError("Fold registry does not cover every galaxy exactly once.")
    return folds, pd.DataFrame(registry_rows).sort_values(["test_fold", "galaxy_id"]).reset_index(drop=True)


def equal_group_weights(groups: np.ndarray) -> np.ndarray:
    groups = np.asarray(groups).astype(str)
    unique, counts = np.unique(groups, return_counts=True)
    count_map = dict(zip(unique, counts))
    weights = np.array([1.0 / count_map[g] for g in groups], dtype=float)
    return weights / weights.mean()
