"""
Active Inference Engine — Free Energy Principle implementation.
Governs agent behavior through Variational Free Energy (VFE) minimization.
"""

import math
import random
import time
from typing import Any


class FreeEnergyEngine:
    def __init__(self, alpha: float = 0.85, learning_rate: float = 0.1):
        self.alpha = alpha
        self.learning_rate = learning_rate
        self.vfe = 0.0  # Variational Free Energy
        self.efe = 0.0  # Expected Free Energy
        self.surprisal_history: list[float] = []
        self.precision = 1.0

    def compute_surprisal(self, predicted: float, observed: float) -> float:
        error = observed - predicted
        return 0.5 * (error ** 2) / self.precision

    def update_vfe(self, surprisal: float):
        self.vfe = self.alpha * self.vfe + (1.0 - self.alpha) * surprisal
        self.surprisal_history.append(surprisal)
        if len(self.surprisal_history) > 100:
            self.surprisal_history.pop(0)
        self.precision = 1.0 / (self._estimate_variance() + 1e-8)

    def _estimate_variance(self) -> float:
        if len(self.surprisal_history) < 2:
            return 1.0
        mean = sum(self.surprisal_history) / len(self.surprisal_history)
        var = sum((s - mean) ** 2 for s in self.surprisal_history) / len(self.surprisal_history)
        return var

    def compute_efe(self, epistemic_value: float, pragmatic_value: float,
                    exploration_bonus: float = 0.0) -> float:
        self.efe = -epistemic_value - pragmatic_value - exploration_bonus
        return self.efe

    def temperature(self) -> float:
        """Adaptive temperature based on free energy."""
        factor = 1.0 + 0.5 * math.tanh(self.vfe - 1.0)
        return max(0.4, min(1.4, 0.8 * factor))

    def confidence_from_vfe(self) -> float:
        """Convert VFE to a confidence score (0-1)."""
        base = 1.0 / (1.0 + abs(self.vfe))
        return max(0.0, min(1.0, base))

    def exploration_urgency(self) -> float:
        """How much the agent should explore vs exploit (0-1)."""
        if len(self.surprisal_history) < 5:
            return 0.5
        recent = self.surprisal_history[-5:]
        avg = sum(recent) / len(recent)
        return min(1.0, avg * 2)

    def should_explore(self) -> bool:
        return self.exploration_urgency() > 0.6

    def state(self) -> dict[str, Any]:
        return {
            "vfe": round(self.vfe, 4),
            "efe": round(self.efe, 4),
            "precision": round(self.precision, 4),
            "temperature": round(self.temperature(), 3),
            "confidence": round(self.confidence_from_vfe(), 3),
            "exploration_urgency": round(self.exploration_urgency(), 3),
            "should_explore": self.should_explore(),
        }


class EmpowermentMetric:
    """Channel capacity between actions and observations.
    High empowerment = agent has control over outcomes.
    """

    def __init__(self, window: int = 10):
        self.window = window
        self.action_outcomes: list[tuple[str, str]] = []

    def record(self, action: str, outcome: str):
        self.action_outcomes.append((action, outcome))
        if len(self.action_outcomes) > self.window * 10:
            self.action_outcomes = self.action_outcomes[-self.window * 10:]

    def compute(self) -> float:
        if len(self.action_outcomes) < 2:
            return 0.5
        recent = self.action_outcomes[-self.window:]
        unique_actions = set(a for a, _ in recent)
        unique_outcomes = set(o for _, o in recent)
        if len(unique_actions) == 0:
            return 0.0
        return len(unique_outcomes) / (len(unique_actions) + 1e-6) / 2