"""Impurity criteria for classification and regression."""

import numpy as np


def gini_impurity(counts: np.ndarray) -> float:
    """Gini impurity from class counts."""
    total = counts.sum()
    if total == 0:
        return 0.0
    proportions = counts / total
    return 1.0 - np.sum(proportions ** 2)


def entropy_impurity(counts: np.ndarray) -> float:
    """Shannon entropy from class counts."""
    total = counts.sum()
    if total == 0:
        return 0.0
    proportions = counts / total
    proportions = proportions[proportions > 0]
    return -np.sum(proportions * np.log2(proportions))


def mse_impurity(y: np.ndarray, sample_weight: np.ndarray = None) -> float:
    """Mean squared error impurity for regression."""
    if len(y) == 0:
        return 0.0
    if sample_weight is not None:
        w_total = sample_weight.sum()
        if w_total == 0:
            return 0.0
        mean = np.average(y, weights=sample_weight)
        return np.average((y - mean) ** 2, weights=sample_weight)
    return np.var(y)


def mae_impurity(y: np.ndarray, sample_weight: np.ndarray = None) -> float:
    """Mean absolute error impurity for regression."""
    if len(y) == 0:
        return 0.0
    if sample_weight is not None:
        w_total = sample_weight.sum()
        if w_total == 0:
            return 0.0
        median = _weighted_median(y, sample_weight)
        return np.average(np.abs(y - median), weights=sample_weight)
    median = np.median(y)
    return np.mean(np.abs(y - median))


def _weighted_median(y: np.ndarray, weights: np.ndarray) -> float:
    """Compute weighted median."""
    sorted_idx = np.argsort(y)
    sorted_y = y[sorted_idx]
    sorted_w = weights[sorted_idx]
    cum_w = np.cumsum(sorted_w)
    half = cum_w[-1] / 2.0
    idx = np.searchsorted(cum_w, half)
    return sorted_y[min(idx, len(sorted_y) - 1)]


def impurity_improvement(
    n_total: float,
    impurity_parent: float,
    n_left: float,
    impurity_left: float,
    n_right: float,
    impurity_right: float,
) -> float:
    """Weighted impurity improvement (gain) from a split."""
    if n_total == 0:
        return 0.0
    return impurity_parent - (
        (n_left / n_total) * impurity_left
        + (n_right / n_total) * impurity_right
    )


CLASSIFICATION_CRITERIA = {
    "gini": gini_impurity,
    "entropy": entropy_impurity,
    "log_loss": entropy_impurity,
}

REGRESSION_CRITERIA = {
    "squared_error": mse_impurity,
    "absolute_error": mae_impurity,
}
