"""ML validation utilities for machine learning research.

This module provides helper functions for validating ML constraints:
data leakage, class balance, feature scaling, NaN/Inf detection, etc.

Author: Generated by AgentBible
"""

from collections.abc import Sequence
from typing import Optional, Union

import numpy as np
from numpy.typing import NDArray


def check_no_data_leakage(
    train_indices: Union[Sequence[int], set[int], NDArray[np.integer]],
    test_indices: Union[Sequence[int], set[int], NDArray[np.integer]],
    name: Optional[str] = None,
) -> None:
    """Validate that train and test indices do not overlap.

    Args:
        train_indices: Indices used for training.
        test_indices: Indices used for testing.
        name: Optional name for error messages.

    Raises:
        ValueError: If there is overlap between train and test.

    Example:
        >>> check_no_data_leakage([0, 1, 2], [3, 4, 5])  # OK
        >>> check_no_data_leakage([0, 1, 2], [2, 3, 4])  # Raises ValueError
    """
    label = f" '{name}'" if name else ""
    train_set = set(train_indices)
    test_set = set(test_indices)

    overlap = train_set & test_set
    if overlap:
        raise ValueError(
            f"Data leakage{label}: {len(overlap)} indices appear in both "
            f"train and test sets. Overlapping indices: {sorted(overlap)[:10]}..."
        )


def check_class_balance(
    labels: NDArray[np.integer],
    max_ratio: float = 10.0,
    name: Optional[str] = None,
) -> None:
    """Validate that class distribution is not too imbalanced.

    Args:
        labels: Array of class labels.
        max_ratio: Maximum allowed ratio between largest and smallest class.
        name: Optional name for error messages.

    Raises:
        ValueError: If class imbalance exceeds max_ratio.
    """
    label = f" '{name}'" if name else ""

    unique, counts = np.unique(labels, return_counts=True)
    if len(unique) < 2:
        raise ValueError(f"Labels{label} contain only one class: {unique[0]}")

    max_count = np.max(counts)
    min_count = np.min(counts)

    if min_count == 0:
        raise ValueError(f"Labels{label} contain class with zero samples")

    ratio = max_count / min_count
    if ratio > max_ratio:
        raise ValueError(
            f"Class imbalance{label}: ratio {ratio:.1f} exceeds max {max_ratio}. "
            f"Class counts: {dict(zip(unique, counts))}"
        )


def check_feature_scaling(
    X: NDArray[np.floating],  # noqa: N803
    expected: str = "standard",
    tolerance: float = 0.1,
    name: Optional[str] = None,
) -> None:
    """Validate that features are properly scaled.

    Args:
        X: Feature matrix (n_samples, n_features).
        expected: Expected scaling - "standard" (mean=0, std=1),
                  "minmax" (min=0, max=1), or "none".
        tolerance: Tolerance for checking mean/std/min/max.
        name: Optional name for error messages.

    Raises:
        ValueError: If features don't match expected scaling.
    """
    label = f" '{name}'" if name else ""

    if expected == "standard":
        means = np.mean(X, axis=0)
        stds = np.std(X, axis=0)

        bad_means = np.abs(means) > tolerance
        bad_stds = np.abs(stds - 1.0) > tolerance

        if np.any(bad_means):
            raise ValueError(
                f"Features{label} not standardized: "
                f"{np.sum(bad_means)} features have |mean| > {tolerance}"
            )
        if np.any(bad_stds):
            raise ValueError(
                f"Features{label} not standardized: "
                f"{np.sum(bad_stds)} features have |std - 1| > {tolerance}"
            )

    elif expected == "minmax":
        mins = np.min(X, axis=0)
        maxs = np.max(X, axis=0)

        if np.any(mins < -tolerance) or np.any(maxs > 1 + tolerance):
            raise ValueError(
                f"Features{label} not min-max scaled to [0, 1]: "
                f"min={np.min(mins):.3f}, max={np.max(maxs):.3f}"
            )


def check_no_nan_inf(
    data: NDArray,
    name: Optional[str] = None,
) -> None:
    """Validate that data contains no NaN or Inf values.

    Args:
        data: Array to check.
        name: Optional name for error messages.

    Raises:
        ValueError: If data contains NaN or Inf.
    """
    label = f" '{name}'" if name else ""

    nan_count = np.sum(np.isnan(data))
    inf_count = np.sum(np.isinf(data))

    if nan_count > 0:
        raise ValueError(f"Data{label} contains {nan_count} NaN values")
    if inf_count > 0:
        raise ValueError(f"Data{label} contains {inf_count} Inf values")


def check_array_shape(
    data: NDArray,
    expected_ndim: Optional[int] = None,
    expected_shape: Optional[tuple] = None,
    min_samples: Optional[int] = None,
    name: Optional[str] = None,
) -> None:
    """Validate array shape constraints.

    Args:
        data: Array to check.
        expected_ndim: Expected number of dimensions.
        expected_shape: Expected shape (use -1 for any size).
        min_samples: Minimum number of samples (first dimension).
        name: Optional name for error messages.

    Raises:
        ValueError: If shape constraints are violated.
    """
    label = f" '{name}'" if name else ""

    if expected_ndim is not None and data.ndim != expected_ndim:
        raise ValueError(
            f"Array{label} has {data.ndim} dimensions, expected {expected_ndim}"
        )

    if expected_shape is not None:
        for i, (actual, expected) in enumerate(zip(data.shape, expected_shape)):
            if expected != -1 and actual != expected:
                raise ValueError(
                    f"Array{label} shape {data.shape} does not match "
                    f"expected {expected_shape} at dimension {i}"
                )

    if min_samples is not None and data.shape[0] < min_samples:
        raise ValueError(
            f"Array{label} has {data.shape[0]} samples, minimum is {min_samples}"
        )


def check_prediction_range(
    predictions: NDArray[np.floating],
    min_val: Optional[float] = None,
    max_val: Optional[float] = None,
    name: Optional[str] = None,
) -> None:
    """Validate that predictions fall within expected range.

    Args:
        predictions: Model predictions.
        min_val: Minimum allowed value (inclusive).
        max_val: Maximum allowed value (inclusive).
        name: Optional name for error messages.

    Raises:
        ValueError: If predictions are out of range.
    """
    label = f" '{name}'" if name else ""

    actual_min = np.min(predictions)
    actual_max = np.max(predictions)

    if min_val is not None and actual_min < min_val:
        raise ValueError(
            f"Predictions{label} below minimum: {actual_min:.4f} < {min_val}"
        )
    if max_val is not None and actual_max > max_val:
        raise ValueError(
            f"Predictions{label} above maximum: {actual_max:.4f} > {max_val}"
        )


def check_probability_predictions(
    probs: NDArray[np.floating],
    tolerance: float = 1e-6,
    name: Optional[str] = None,
) -> None:
    """Validate that probability predictions are valid.

    Checks:
    1. All values in [0, 1]
    2. Rows sum to 1 (for multi-class)

    Args:
        probs: Probability predictions (n_samples,) or (n_samples, n_classes).
        tolerance: Numerical tolerance.
        name: Optional name for error messages.

    Raises:
        ValueError: If probabilities are invalid.
    """
    label = f" '{name}'" if name else ""

    # Check range
    if np.any(probs < -tolerance) or np.any(probs > 1 + tolerance):
        raise ValueError(
            f"Probabilities{label} out of range [0, 1]: "
            f"min={np.min(probs):.4f}, max={np.max(probs):.4f}"
        )

    # Check sum to 1 for multi-class
    if probs.ndim == 2 and probs.shape[1] > 1:
        row_sums = np.sum(probs, axis=1)
        bad_rows = np.abs(row_sums - 1.0) > tolerance
        if np.any(bad_rows):
            raise ValueError(
                f"Probabilities{label} rows do not sum to 1: "
                f"{np.sum(bad_rows)} rows have sum != 1"
            )
