"""Prediction-space stratification utilities."""
import numpy as np
from .simplex import entropy, ilr
from sklearn.cluster import KMeans


def stratify_by_entropy(U: np.ndarray, n_bins: int = 5) -> np.ndarray:
    """Bin predictions by Shannon entropy."""
    H = entropy(U)
    return np.digitize(H, np.linspace(H.min(), H.max(), n_bins + 1)[1:-1])


def stratify_by_boundary(U: np.ndarray, n_bins: int = 5) -> np.ndarray:
    """Bin predictions by proximity to simplex boundary (min component)."""
    bprox = U.min(axis=-1)
    return np.digitize(bprox, np.linspace(bprox.min(), bprox.max(), n_bins + 1)[1:-1])


def stratify_by_kmeans(U: np.ndarray, n_clusters: int = 5, seed: int = 42) -> np.ndarray:
    """K-means clustering in ILR space."""
    Z = ilr(U)
    km = KMeans(n_clusters=n_clusters, random_state=seed, n_init=10)
    return km.fit_predict(Z)


def stratify_by_argmax_group(U: np.ndarray, split_index: int = 5) -> np.ndarray:
    """Two-group stratification based on whether argmax(U) is before split_index."""
    top_class = np.argmax(U, axis=1)
    return (top_class >= split_index).astype(int)


def _quantile_edges(values: np.ndarray, n_bins: int) -> np.ndarray:
    """Return stable interior quantile edges for binning a scalar score."""
    if n_bins <= 1:
        return np.array([], dtype=float)
    qs = np.linspace(0.0, 1.0, n_bins + 1)[1:-1]
    edges = np.quantile(values, qs)
    return np.unique(edges)


def _digitize_fixed(values: np.ndarray, n_bins: int) -> np.ndarray:
    """Digitize values using globally fixed quantile edges."""
    edges = _quantile_edges(values, n_bins)
    if edges.size == 0:
        return np.zeros(len(values), dtype=int)
    return np.digitize(values, edges)


def precompute_fixed_strata(
    U: np.ndarray,
    method: str,
    n_strata: int = 5,
    seed: int = 42,
) -> np.ndarray:
    """Precompute a fixed stratification on a full cached prediction matrix.

    This is intended for repeated cal/test splits of the same frozen task, where
    the stratification rule should remain constant across repetitions.

    Args:
        U: full prediction matrix with shape (n, K).
        method: one of {"entropy", "boundary", "dominant", "kmeans", "random"}.
        n_strata: target number of strata. For "dominant", this is treated as
            the number of grouped dominant-component bins; if it is at least K,
            each dominant component gets its own stratum.
        seed: random seed for methods that require stochastic initialization.

    Returns:
        Integer stratum labels of shape (n,).
    """
    method = method.lower()

    if method == "entropy":
        return _digitize_fixed(entropy(U), n_strata)

    if method == "boundary":
        return _digitize_fixed(U.min(axis=-1), n_strata)

    if method == "dominant":
        top = np.argmax(U, axis=1)
        k = U.shape[1]
        n_groups = min(max(int(n_strata), 1), k)
        if n_groups >= k:
            return top.astype(int)
        return np.floor(top * n_groups / k).astype(int)

    if method == "random":
        z = ilr(U)
        rng = np.random.default_rng(seed)
        direction = rng.normal(size=z.shape[1])
        direction /= np.linalg.norm(direction) + 1e-12
        score = z @ direction
        return _digitize_fixed(score, n_strata)

    if method == "kmeans":
        z = ilr(U)
        km = KMeans(n_clusters=n_strata, random_state=seed, n_init=10)
        return km.fit_predict(z)

    raise ValueError(f"Unknown fixed stratification method: {method}")