"""Uncertainty measures derived from ensemble member probabilities.

Definitions follow ai/ANALYSIS_PLAN.md section 2.

All measures take member probabilities ``probs`` of shape ``(M, S, C)`` -- M
ensemble members, S samples, C classes -- and return a per-sample uncertainty
``(S,)`` where **higher means more uncertain**. A ``single``-model configuration
passes ``M = 1``.

Key identity that keeps the two families consistent
---------------------------------------------------
The stored ``predictive_entropy`` of member *m* equals ``H(probs[m])``, because
``prob`` is already the MC-mean for a Bayesian member and the plain softmax for a
deterministic one. So the ensemble decomposition

    H(mean_m p_m)          = total       (mean_output_entropy)
    mean_m H(p_m)          = aleatoric   (mean_predictive_entropy)
    total - aleatoric      = epistemic   (ensemble_mutual_information)

holds for BOTH families, computed from ``prob`` alone. ``mean_mutual_information``
is different: it averages each member's *within-member* MC epistemic term, which
is identically zero for deterministic models.
"""

from dataclasses import dataclass
from typing import Callable, Dict, List, Optional

import numpy as np

from analysis.context import CONFIG_SINGLE
from evaluation import MODEL_KIND_BAYESIAN

#: Numerical floor for logarithms, matching bayesian_torch's convention.
EPS = 1e-15


# ---------------------------------------------------------------------------
# Primitive
# ---------------------------------------------------------------------------
def entropy(prob: np.ndarray) -> np.ndarray:
    """Shannon entropy along the last axis: ``-sum(p * log(p))``.

    Args:
        prob: Probabilities with classes on the last axis, e.g. ``(..., C)``.

    Returns:
        Entropy with the class axis removed. Natural log, so a uniform binary
        distribution gives ``ln 2 ~= 0.693``.
    """
    p = np.asarray(prob, dtype=float)
    return -np.sum(p * np.log(p + EPS), axis=-1)


def predicted_class(probs: np.ndarray) -> np.ndarray:
    """Ensemble prediction ``(S,)``: argmax of the mean member probability."""
    return np.asarray(probs, dtype=float).mean(axis=0).argmax(axis=1)


# ---------------------------------------------------------------------------
# Measures -- each returns (S,), higher = more uncertain
# ---------------------------------------------------------------------------
def mean_output_entropy(probs: np.ndarray, **_: object) -> np.ndarray:
    """Total uncertainty: entropy of the ensemble mean output, ``H(p_bar)``."""
    return entropy(np.asarray(probs, dtype=float).mean(axis=0))


def mean_predictive_entropy(probs: np.ndarray, **_: object) -> np.ndarray:
    """Aleatoric: mean of the members' own entropies, ``mean_m H(p_m)``.

    Equal to :func:`mean_output_entropy` when ``M == 1`` (hence skipped for the
    ``single`` configuration -- see :data:`MEASURES`).
    """
    return entropy(np.asarray(probs, dtype=float)).mean(axis=0)


def ensemble_std(probs: np.ndarray, **_: object) -> np.ndarray:
    """Member disagreement: ``std_m(p_m[y_hat])`` on the ensemble-predicted class.

    Measures how much the members disagree about the probability they assign to
    the class the ensemble ultimately predicts. Identically 0 for a single model;
    the applicability matrix excludes it there.
    """
    values = np.asarray(probs, dtype=float)
    y_hat = predicted_class(values)  # (S,)
    # Take each member's probability for the ensemble-predicted class -> (M, S).
    indices = np.broadcast_to(
        y_hat[None, :, None], (values.shape[0], values.shape[1], 1)
    )
    member_prob = np.take_along_axis(values, indices, axis=2)[:, :, 0]
    return member_prob.std(axis=0)


def ensemble_mutual_information(probs: np.ndarray, **_: object) -> np.ndarray:
    """Epistemic from disagreement: total minus aleatoric (Jensen-Shannon).

    ``H(p_bar) - mean_m H(p_m)`` -- non-negative by Jensen's inequality, and
    identically zero when ``M == 1``.
    """
    return mean_output_entropy(probs) - mean_predictive_entropy(probs)


def mean_mutual_information(
    probs: np.ndarray, member_mi: Optional[np.ndarray] = None, **_: object
) -> np.ndarray:
    """Within-member MC epistemic: ``mean_m MI_m`` from the stored statistic.

    Args:
        probs: Unused; kept for a uniform measure signature.
        member_mi: Stored ``mutual_information`` of shape ``(M, S)``.

    Bayesian only -- identically zero for deterministic members (K = 1).
    """
    if member_mi is None:
        raise ValueError(
            "mean_mutual_information needs the stored per-member "
            "`mutual_information` (M, S); pass member_mi=..."
        )
    return np.asarray(member_mi, dtype=float).mean(axis=0)


# ---------------------------------------------------------------------------
# Registry + applicability
# ---------------------------------------------------------------------------
MeasureFn = Callable[..., np.ndarray]


@dataclass(frozen=True)
class Measure:
    """An uncertainty measure and where it is scientifically meaningful."""

    name: str
    label: str
    fn: MeasureFn
    #: Needs more than one member (excludes the ``single`` configuration).
    requires_ensemble: bool = False
    #: Restricted to these families (empty = all).
    families: tuple[str, ...] = ()
    #: Needs the stored per-member ``mutual_information`` variable.
    needs_member_mi: bool = False
    #: Redundant with ``mean_output_entropy`` when M == 1.
    duplicate_when_single: bool = False

    def applies_to(self, family: str, config: str) -> bool:
        """Whether this measure is meaningful for a family x configuration."""
        if self.families and family not in self.families:
            return False
        if config == CONFIG_SINGLE and (
            self.requires_ensemble or self.duplicate_when_single
        ):
            return False
        return True


#: All uncertainty measures, in presentation order.
MEASURES: Dict[str, Measure] = {
    "mean_output_entropy": Measure(
        name="mean_output_entropy",
        label="Mean output",
        fn=mean_output_entropy,
    ),
    "mean_predictive_entropy": Measure(
        name="mean_predictive_entropy",
        label="Predictive uncertainty",
        fn=mean_predictive_entropy,
        duplicate_when_single=True,
    ),
    "ensemble_std": Measure(
        name="ensemble_std",
        label="Standard deviation",
        fn=ensemble_std,
        requires_ensemble=True,
    ),
    "ensemble_mutual_information": Measure(
        name="ensemble_mutual_information",
        label="Ensemble model uncertainty",
        fn=ensemble_mutual_information,
        requires_ensemble=True,
    ),
    "mean_mutual_information": Measure(
        name="mean_mutual_information",
        label="Model uncertainty",
        fn=mean_mutual_information,
        families=(MODEL_KIND_BAYESIAN,),
        needs_member_mi=True,
    ),
}


def measures_for(family: str, config: str) -> List[Measure]:
    """Measures that are meaningful for one family x configuration."""
    return [m for m in MEASURES.values() if m.applies_to(family, config)]


def compute(
    measure: Measure,
    probs: np.ndarray,
    member_mi: Optional[np.ndarray] = None,
) -> np.ndarray:
    """Evaluate ``measure`` on member probabilities, supplying what it needs.

    Args:
        measure: The measure to evaluate.
        probs: Member probabilities ``(M, S, C)``.
        member_mi: Stored per-member ``mutual_information`` ``(M, S)``; required
            only by measures with ``needs_member_mi``.

    Returns:
        Per-sample uncertainty ``(S,)``, higher = more uncertain.
    """
    if measure.needs_member_mi:
        return measure.fn(probs, member_mi=member_mi)
    return measure.fn(probs)