"""ESS (Effective Sample Size) monitoring with alerts.

Provides real-time monitoring of ESS across time steps with configurable
alert thresholds for detecting particle filter degeneracy.

Reference:
    Kong, A., Liu, J.S. & Wong, W.H. (1994). Sequential imputations and
    Bayesian missing data problems. JASA, 89(425), 278-288.
"""

from __future__ import annotations

from dataclasses import dataclass
from enum import Enum
from typing import Any

import numpy as np
from numpy.typing import NDArray


class AlertLevel(Enum):
    """Alert severity levels for ESS monitoring."""

    OK = "OK"
    WARNING = "WARNING"
    CRITICAL = "CRITICAL"


@dataclass
class ESSAlert:
    """An alert generated by the ESS monitor.

    Attributes:
        level: Severity level (OK, WARNING, CRITICAL).
        time_step: Time step at which the alert was generated.
        ess_value: ESS value that triggered the alert.
        n_particles: Number of particles at the time.
        message: Human-readable description of the alert.
    """

    level: AlertLevel
    time_step: int
    ess_value: float
    n_particles: int
    message: str


class ESSMonitor:
    """Real-time ESS monitoring with configurable alerts.

    Tracks ESS over time and generates alerts when ESS drops below
    configurable thresholds. Supports both normalized weights and
    ParticleCloud objects.

    Parameters:
        warning_ratio: ESS/N ratio below which WARNING is triggered (default 0.1).
        critical_ess: Absolute ESS value at or below which CRITICAL is triggered (default 1.0).

    Examples:
        >>> monitor = ESSMonitor()
        >>> # With uniform weights (healthy)
        >>> weights = np.ones(100) / 100
        >>> monitor.update_from_weights(weights, time_step=0)
        >>> assert monitor.is_healthy()
        >>> # With degenerate weights
        >>> weights = np.zeros(100)
        >>> weights[0] = 1.0
        >>> monitor.update_from_weights(weights, time_step=1)
        >>> assert not monitor.is_healthy()
    """

    def __init__(
        self,
        warning_ratio: float = 0.1,
        critical_ess: float = 1.0,
    ) -> None:
        self.warning_ratio = warning_ratio
        self.critical_ess = critical_ess

        self._ess_history: list[float] = []
        self._n_particles_history: list[int] = []
        self._time_steps: list[int] = []
        self._alerts: list[ESSAlert] = []

    @property
    def ess_min(self) -> float:
        """Minimum ESS observed across all time steps."""
        if not self._ess_history:
            return float("nan")
        return float(np.min(self._ess_history))

    @property
    def ess_mean(self) -> float:
        """Mean ESS across all time steps."""
        if not self._ess_history:
            return float("nan")
        return float(np.mean(self._ess_history))

    @property
    def ess_ratio_history(self) -> NDArray[np.float64]:
        """ESS/N ratio history as array."""
        if not self._ess_history:
            return np.array([], dtype=np.float64)
        ess = np.array(self._ess_history, dtype=np.float64)
        n = np.array(self._n_particles_history, dtype=np.float64)
        return ess / n

    @property
    def alerts(self) -> list[ESSAlert]:
        """List of all alerts generated."""
        return list(self._alerts)

    @staticmethod
    def compute_ess(weights: NDArray[np.float64]) -> float:
        """Compute ESS from normalized weights.

        ESS = 1 / sum(w_i^2)

        Parameters:
            weights: Normalized weights (must sum to 1).

        Returns:
            Effective sample size.
        """
        w = np.asarray(weights, dtype=np.float64)
        w = w / w.sum()  # ensure normalization
        return float(1.0 / np.sum(w**2))

    def update_from_weights(
        self,
        weights: NDArray[np.float64],
        time_step: int | None = None,
    ) -> AlertLevel:
        """Update monitor with new normalized weights.

        Parameters:
            weights: Normalized particle weights.
            time_step: Optional time step index. Auto-incremented if not provided.

        Returns:
            Alert level for this time step.
        """
        w = np.asarray(weights, dtype=np.float64)
        n_particles = len(w)
        ess = self.compute_ess(w)

        if time_step is None:
            time_step = len(self._ess_history)

        self._ess_history.append(ess)
        self._n_particles_history.append(n_particles)
        self._time_steps.append(time_step)

        # Determine alert level
        level = AlertLevel.OK
        message = f"ESS={ess:.1f}/{n_particles} ({ess / n_particles:.1%})"

        if ess <= self.critical_ess:
            level = AlertLevel.CRITICAL
            message = (
                f"CRITICAL: ESS={ess:.1f} (effective sample size collapsed to 1). "
                f"All weight on a single particle."
            )
        elif ess < self.warning_ratio * n_particles:
            level = AlertLevel.WARNING
            message = (
                f"WARNING: ESS={ess:.1f}/{n_particles} "
                f"({ess / n_particles:.1%} < {self.warning_ratio:.0%} threshold). "
                f"Consider resampling or increasing N."
            )

        if level != AlertLevel.OK:
            alert = ESSAlert(
                level=level,
                time_step=time_step,
                ess_value=ess,
                n_particles=n_particles,
                message=message,
            )
            self._alerts.append(alert)

        return level

    def update(self, cloud: Any, time_step: int | None = None) -> AlertLevel:
        """Update monitor from a ParticleCloud object.

        Parameters:
            cloud: Object with a `weights` attribute (normalized).
            time_step: Optional time step index.

        Returns:
            Alert level for this time step.
        """
        weights = np.asarray(cloud.weights, dtype=np.float64)
        return self.update_from_weights(weights, time_step=time_step)

    def is_healthy(self) -> bool:
        """Check if the filter is healthy (no CRITICAL alerts).

        Returns:
            True if no CRITICAL alerts have been generated.
        """
        return not any(a.level == AlertLevel.CRITICAL for a in self._alerts)

    def summary(self) -> dict[str, Any]:
        """Generate a summary of ESS monitoring.

        Returns:
            Dictionary with ESS statistics and alert counts.
        """
        n_warnings = sum(1 for a in self._alerts if a.level == AlertLevel.WARNING)
        n_critical = sum(1 for a in self._alerts if a.level == AlertLevel.CRITICAL)

        result: dict[str, Any] = {
            "n_time_steps": len(self._ess_history),
            "ess_min": self.ess_min,
            "ess_mean": self.ess_mean,
            "ess_std": float(np.std(self._ess_history)) if self._ess_history else float("nan"),
            "n_warnings": n_warnings,
            "n_critical": n_critical,
            "is_healthy": self.is_healthy(),
        }

        if self._ess_history:
            ratios = self.ess_ratio_history
            result["ess_ratio_min"] = float(np.min(ratios))
            result["ess_ratio_mean"] = float(np.mean(ratios))

        return result

    def reset(self) -> None:
        """Reset all monitoring state."""
        self._ess_history.clear()
        self._n_particles_history.clear()
        self._time_steps.clear()
        self._alerts.clear()
