"""
Epiplexity: Epistemic Health Monitoring via Bayesian Surprise
==============================================================

Biological Analogy:
- Trophic Factors: Growth signals that maintain neuronal health. Without novel
  stimulation, neurons atrophy. Similarly, agents need "informational nutrition"
  to maintain healthy reasoning.
- Free Energy Principle: Living systems minimize surprise while maintaining
  viability. An agent with no surprise is either perfectly adapted (unlikely)
  or epistemically stagnant (loop pathology).

The key insight: If an agent's outputs stabilize (low embedding distance) while
its perplexity remains high (model is uncertain), it's in a pathological loop.
Healthy agents either:
1. Converge (low perplexity) → task complete
2. Explore (high embedding novelty) → making progress

Epiplexity measures the *absence* of this healthy pattern.

References:
- Friston, K. (2010). The free-energy principle: a unified brain theory?
- Article Section 5.1: Oncology - Infinite Loops as Epistemic Starvation
"""
from dataclasses import dataclass, field
from typing import Any, Protocol, Sequence, runtime_checkable
from enum import Enum
from collections import deque
import math
import hashlib
import operator


class HealthStatus(Enum):
    """Epistemic health status."""
    HEALTHY = "healthy"           # Normal operation
    CONVERGING = "converging"     # Low epiplexity, likely task completion
    EXPLORING = "exploring"       # High novelty, making progress
    STAGNANT = "stagnant"         # Low epiplexity (no surprise), possible loop
    CRITICAL = "critical"         # Sustained low epiplexity, intervention needed


@dataclass
class EpiplexityResult:
    """Result of a single epiplexity measurement."""

    # Raw components
    embedding_novelty: float  # ½(1 - cos(e_t, e_{t-1})), range [0, 1]
    normalized_perplexity: float  # σ(H), range [0, 1]

    # Combined score
    epiplexity: float  # Ê_t, range [0, 1]

    # Windowed integral
    epiplexic_integral: float  # E_w = mean of window

    # Interpretation
    status: HealthStatus
    window_size: int
    threshold: float

    @property
    def is_healthy(self) -> bool:
        """Check if agent is in healthy epistemic state."""
        return self.status in (HealthStatus.HEALTHY, HealthStatus.CONVERGING, HealthStatus.EXPLORING)


@dataclass
class EpiplexityState:
    """Tracks epistemic state over time."""

    # Current embedding
    current_embedding: list[float] | None = None
    previous_embedding: list[float] | None = None

    # History for windowed analysis
    epiplexity_history: deque[float] = field(default_factory=lambda: deque(maxlen=100))

    # Distance-based tracking (for DistanceProvider path)
    previous_item: Any | None = None

    # Statistics
    total_measurements: int = 0
    stagnant_episodes: int = 0
    max_consecutive_stagnant: int = 0
    current_consecutive_stagnant: int = 0


@runtime_checkable
class DistanceProvider(Protocol):
    """Protocol for distance-based novelty measurement (non-embedding).

    Enables EpiplexityMonitor to operate on arbitrary objects (e.g.,
    CandidateConfig) rather than requiring text embeddings.  Tests whether
    epistemic health monitoring generalizes across abstraction scales.
    """

    def distance(self, a: Any, b: Any) -> float:
        """Return distance in [0, 1] between two objects."""
        ...


class EmbeddingProvider(Protocol):
    """Protocol for embedding providers."""

    def embed(self, text: str) -> list[float]:
        """Get embedding for text."""
        ...


class MockEmbeddingProvider:
    """
    Mock embedding provider for testing.

    Uses hash-based pseudo-embeddings that are deterministic
    but provide reasonable similarity structure.
    """

    def __init__(self, dim: int = 128):
        self.dim = dim

    def embed(self, text: str) -> list[float]:
        """Generate deterministic pseudo-embedding from text hash."""
        # Create deterministic seed from text
        text_hash = hashlib.sha256(text.encode()).digest()

        # Generate embedding components from hash
        embedding = []
        for i in range(self.dim):
            # Use different parts of hash for each dimension
            byte_idx = i % len(text_hash)
            val = (text_hash[byte_idx] + i * 17) % 256
            # Normalize to [-1, 1]
            embedding.append((val / 127.5) - 1.0)

        # Normalize to unit vector
        # Optimization: use C-backed math.hypot instead of math.sqrt(sum(x * x for x in vector))
        magnitude = math.hypot(*embedding) if embedding else 0.0
        if magnitude > 0:
            embedding = [x / magnitude for x in embedding]

        return embedding


@dataclass
class EpiplexityMonitor:
    """
    Monitor for detecting epistemic stagnation.

    Implements the operational Epiplexity approximation from the
    oncology / epistemic-starvation section:

        Ê_t = α·½(1 - cos(e_t, e_{t-1})) + (1-α)·σ(H(m_t|m_{<t}))

    Where:
    - e_t is the embedding of message t
    - cos is cosine similarity
    - σ(H) = 1 - e^{-H/H_0} is exponential saturation to [0,1]
    - H is conditional perplexity (approximated)
    - α is the mixing parameter

    The Epiplexic Integral for windowed detection:

        E_w = (1/w) Σ Ê_t  for t in window

    Low Epiplexity = low Bayesian surprise = stagnation.
    An agent is flagged when E_w < δ for sustained periods.

    Example:
        >>> monitor = EpiplexityMonitor(
        ...     embedding_provider=MockEmbeddingProvider(),
        ...     alpha=0.5,
        ...     window_size=5,
        ...     threshold=0.2,
        ... )
        >>> result = monitor.measure("First message")
        >>> result.status
        HealthStatus.HEALTHY
        >>> # Repeating same message creates stagnation
        >>> for _ in range(10):
        ...     result = monitor.measure("Repeating message")
        >>> result.status  # Will eventually become STAGNANT
    """

    embedding_provider: EmbeddingProvider | None = None
    alpha: float = 0.5  # Mixing parameter (0 = perplexity only, 1 = embedding only)
    window_size: int = 10  # Window for integral calculation
    threshold: float = 0.2  # δ threshold: E_w < δ indicates stagnation
    critical_duration: int = 5  # Consecutive stagnant measurements before CRITICAL

    # Distance-based novelty (alternative to embedding path)
    distance_provider: DistanceProvider | None = None

    # Internal state
    state: EpiplexityState = field(default_factory=EpiplexityState)

    # Perplexity normalization parameter (H_0 in the paper)
    perplexity_h0: float = 2.0  # Baseline perplexity for exponential saturation

    def measure(
        self,
        message: str = "",
        perplexity: float | None = None,
        *,
        item: Any | None = None,
    ) -> EpiplexityResult:
        """
        Measure epiplexity for a new message or item.

        When ``distance_provider`` is set and ``item`` is provided, novelty
        is computed via the distance provider instead of embedding cosine
        similarity.  This enables monitoring of non-text objects (e.g.,
        CandidateConfig) for epistemic health.

        Args:
            message: The agent's output message (embedding path)
            perplexity: Optional measured perplexity. If not provided,
                       approximated from embedding stability.
            item: Object to measure distance-based novelty for (distance path)

        Returns:
            EpiplexityResult with current health status
        """
        # --- Distance-provider path (non-embedding) ---
        if self.distance_provider is not None and item is not None:
            if self.state.previous_item is not None:
                embedding_novelty = self.distance_provider.distance(
                    self.state.previous_item, item,
                )
            else:
                embedding_novelty = 1.0  # First item — maximum novelty

            # Without real perplexity, approximate from novelty trend
            normalized_perplexity = self._approximate_perplexity_from_novelty(
                embedding_novelty,
            )

            epiplexity = (
                self.alpha * embedding_novelty
                + (1 - self.alpha) * normalized_perplexity
            )

            self.state.previous_item = item
            self.state.epiplexity_history.append(epiplexity)
            self.state.total_measurements += 1

            window = list(self.state.epiplexity_history)[-self.window_size:]
            epiplexic_integral = sum(window) / len(window) if window else 0.0

            status = self._determine_status(
                embedding_novelty, normalized_perplexity, epiplexic_integral,
            )
            self._track_stagnation(status)

            return EpiplexityResult(
                embedding_novelty=embedding_novelty,
                normalized_perplexity=normalized_perplexity,
                epiplexity=epiplexity,
                epiplexic_integral=epiplexic_integral,
                status=status,
                window_size=len(window),
                threshold=self.threshold,
            )

        # --- Embedding path (original behavior) ---
        if self.embedding_provider is None:
            raise ValueError(
                "EpiplexityMonitor requires either embedding_provider or "
                "distance_provider (with item=) to be set"
            )
        embedding = self.embedding_provider.embed(message)

        # Calculate embedding novelty: 1 - cos(e_t, e_{t-1})
        if self.state.previous_embedding is not None:
            cosine_sim = self._cosine_similarity(embedding, self.state.previous_embedding)
            # Normalize from [-1, 1] to [0, 2], then to [0, 1]
            embedding_novelty = (1.0 - cosine_sim) / 2.0
        else:
            # First message - maximum novelty
            embedding_novelty = 1.0

        # Calculate normalized perplexity
        if perplexity is not None:
            # Use provided perplexity with exponential saturation: σ(H) = 1 - e^{-H/H_0}
            normalized_perplexity = self._exponential_saturation(perplexity)
        else:
            # Approximate: low novelty + repetition suggests high uncertainty
            # This is a heuristic when actual perplexity isn't available
            normalized_perplexity = self._approximate_perplexity(embedding, embedding_novelty)

        # Calculate combined Epiplexity score
        # Ê_t = α·novelty + (1-α)·perplexity
        epiplexity = (
            self.alpha * embedding_novelty +
            (1 - self.alpha) * normalized_perplexity
        )

        # Update state
        self.state.previous_embedding = self.state.current_embedding
        self.state.current_embedding = embedding
        self.state.epiplexity_history.append(epiplexity)
        self.state.total_measurements += 1

        # Calculate Epiplexic Integral (windowed mean)
        window = list(self.state.epiplexity_history)[-self.window_size:]
        epiplexic_integral = sum(window) / len(window) if window else 0.0

        # Determine health status
        status = self._determine_status(
            embedding_novelty,
            normalized_perplexity,
            epiplexic_integral,
        )

        self._track_stagnation(status)

        return EpiplexityResult(
            embedding_novelty=embedding_novelty,
            normalized_perplexity=normalized_perplexity,
            epiplexity=epiplexity,
            epiplexic_integral=epiplexic_integral,
            status=status,
            window_size=len(window),
            threshold=self.threshold,
        )

    def _cosine_similarity(self, a: Sequence[float], b: Sequence[float]) -> float:
        """Calculate cosine similarity between two vectors."""
        if len(a) != len(b):
            raise ValueError("Vectors must have same dimension")

        # Optimization: C-backed map(operator.mul) is faster than zip generator
        dot_product = sum(map(operator.mul, a, b))

        # Optimization: C-backed math.hypot is ~6x faster than math.sqrt(sum(x * x))
        magnitude_a = math.hypot(*a) if a else 0.0
        magnitude_b = math.hypot(*b) if b else 0.0

        if magnitude_a == 0 or magnitude_b == 0:
            return 0.0

        return dot_product / (magnitude_a * magnitude_b)

    def _exponential_saturation(self, h: float) -> float:
        """Exponential saturation σ(H) = 1 - e^{-H/H_0}, mapping [0, ∞) → [0, 1)."""
        return 1.0 - math.exp(-h / self.perplexity_h0)

    def _approximate_perplexity(
        self,
        embedding: Sequence[float],
        novelty: float,
    ) -> float:
        """
        Approximate perplexity when not directly available.

        Heuristic: If output is repetitive (low novelty) but embedding
        is still varying (not converged), model is likely uncertain.
        """
        if self.state.previous_embedding is None:
            return 0.5  # Neutral starting point

        # Check embedding stability (variance of differences)
        history_len = len(self.state.epiplexity_history)
        if history_len < 3:
            return 0.5

        # Use embedding magnitude variance as additional uncertainty signal
        # Optimization: math.hypot is faster for Euclidean norms
        magnitude = math.hypot(*embedding) if embedding else 0.0
        prev_magnitude = math.hypot(*self.state.previous_embedding) if self.state.previous_embedding else 0.0
        magnitude_drift = abs(magnitude - prev_magnitude)

        # Recent trend: increasing epiplexity with low novelty suggests stagnation
        recent = list(self.state.epiplexity_history)[-3:]
        trend = (recent[-1] - recent[0]) / len(recent) if len(recent) > 1 else 0

        # Low novelty with increasing trend = high approximate perplexity
        # Add magnitude drift as additional uncertainty indicator
        if novelty < 0.3 and trend > 0:
            return min(1.0, 0.5 + trend * 2 + magnitude_drift * 0.5)

        return 0.3 + novelty * 0.4 + magnitude_drift * 0.2

    def _approximate_perplexity_from_novelty(self, novelty: float) -> float:
        """Approximate perplexity from novelty alone (distance-provider path).

        Without embeddings, we use the novelty trend as a proxy for
        model uncertainty.  Low novelty with a declining trend suggests
        stagnation; high novelty suggests exploration.
        """
        history = list(self.state.epiplexity_history)
        if len(history) < 2:
            return 0.5

        recent = history[-3:]
        trend = (recent[-1] - recent[0]) / len(recent) if len(recent) > 1 else 0

        if novelty < 0.3 and trend <= 0:
            return min(1.0, 0.6 + abs(trend) * 2)
        return 0.3 + novelty * 0.4

    def _track_stagnation(self, status: HealthStatus) -> None:
        """Update stagnation episode counters."""
        if status in (HealthStatus.STAGNANT, HealthStatus.CRITICAL):
            self.state.current_consecutive_stagnant += 1
            if self.state.current_consecutive_stagnant > self.state.max_consecutive_stagnant:
                self.state.max_consecutive_stagnant = self.state.current_consecutive_stagnant
        else:
            if self.state.current_consecutive_stagnant > 0:
                self.state.stagnant_episodes += 1
            self.state.current_consecutive_stagnant = 0

    def _determine_status(
        self,
        novelty: float,
        perplexity: float,
        integral: float,
    ) -> HealthStatus:
        """
        Determine health status from components.

        Decision logic (per paper Section 5):
        - High novelty (>0.5): EXPLORING (making progress)
        - Low novelty + low perplexity (<0.3): CONVERGING (task complete)
        - Low novelty + high perplexity: STAGNANT (loop — stuck but uncertain)
        - Integral < threshold for critical_duration: CRITICAL (sustained stagnation)
        """
        # Check for critical (sustained low surprise = deep stagnation)
        if (
            integral < self.threshold and
            self.state.current_consecutive_stagnant >= self.critical_duration
        ):
            return HealthStatus.CRITICAL

        # High novelty = exploring
        if novelty > 0.5:
            return HealthStatus.EXPLORING

        # Low novelty + low perplexity = converging to solution
        if novelty < 0.3 and perplexity < 0.3:
            return HealthStatus.CONVERGING

        # Low novelty + high perplexity = stagnant (loop pathology)
        if novelty < 0.3 and perplexity > 0.5:
            return HealthStatus.STAGNANT

        # Low integral = low Bayesian surprise = stagnant
        if integral < self.threshold:
            return HealthStatus.STAGNANT

        return HealthStatus.HEALTHY

    def reset(self) -> None:
        """Reset monitor state."""
        self.state = EpiplexityState()

    def stats(self) -> dict:
        """Return monitoring statistics."""
        history = list(self.state.epiplexity_history)
        return {
            "total_measurements": self.state.total_measurements,
            "stagnant_episodes": self.state.stagnant_episodes,
            "max_consecutive_stagnant": self.state.max_consecutive_stagnant,
            "current_consecutive_stagnant": self.state.current_consecutive_stagnant,
            "mean_epiplexity": sum(history) / len(history) if history else 0.0,
            "max_epiplexity": max(history) if history else 0.0,
            "window_size": self.window_size,
            "threshold": self.threshold,
        }
