"""Global stability monitoring for HIEROMEM."""

from __future__ import annotations

from collections import OrderedDict, deque
from dataclasses import dataclass
from typing import Deque, Hashable, Iterable

import numpy as np


@dataclass(slots=True)
class StabilitySnapshot:
    retrieval_entropy: float
    summary_fidelity: float
    summary_distortion: float
    summary_distortion_bound: float
    trust_floor: float
    memory_utilisation: float
    contraction_factor: float
    lyapunov_margin: float
    lyapunov_warning: bool
    stability_score: float


class GlobalStabilityMonitor:
    """Track convergence and bounded growth statistics."""

    def __init__(self, window: int = 32, smoothing: float = 0.1) -> None:
        self.window = int(max(window, 1))
        self.smoothing = float(max(smoothing, 1e-6))
        self._retrieval_entropy: Deque[float] = deque(maxlen=self.window)
        self._summary_fidelity: Deque[float] = deque(maxlen=self.window)
        self._summary_distortion: Deque[float] = deque(maxlen=self.window)
        self._distortion_bounds: Deque[float] = deque(maxlen=self.window)
        self._trust_floor: Deque[float] = deque(maxlen=self.window)
        self._memory_utilisation: Deque[float] = deque(maxlen=self.window)
        self._contraction_factors: Deque[float] = deque(maxlen=self.window)
        self._lyapunov_margins: Deque[float] = deque(maxlen=self.window)
        self._stabilised_weights: OrderedDict[Hashable,
                                              np.ndarray] = OrderedDict()

    # ------------------------------------------------------------------ #
    def observe_retrieval(self, weights: Iterable[float]) -> float:
        weights_arr = np.asarray(list(weights), dtype=float)
        if weights_arr.size == 0:
            entropy = 0.0
        else:
            weights_arr = weights_arr / (np.sum(weights_arr) + self.smoothing)
            entropy = float(-np.sum(weights_arr *
                            np.log(weights_arr + self.smoothing)))
        self._retrieval_entropy.append(entropy)
        return entropy

    def stabilise_weights(
            self,
            weights: Iterable[float],
            *,
            key: Hashable | None = None) -> np.ndarray:
        """Return a smoothed version of ``weights`` to reduce oscillation.

        The optional ``key`` argument scopes the smoothing cache to a specific
        retrieval. Using different keys prevents the state from contaminating
        subsequent queries that just happen to produce arrays of the same shape.
        """

        arr = np.asarray(list(weights), dtype=float)
        if arr.size == 0:
            return arr

        cache_key: Hashable = key if key is not None else "__default__"
        cached = self._stabilised_weights.get(cache_key)
        if cached is None or cached.shape != arr.shape:
            self._remember_stabilised(cache_key, arr.copy())
            return arr

        blended = self.smoothing * arr + (1.0 - self.smoothing) * cached
        self._remember_stabilised(cache_key, blended)
        return blended

    def _remember_stabilised(self, key: Hashable, value: np.ndarray) -> None:
        """Cache stabilised weights and evict old keys to bound memory use."""

        self._stabilised_weights[key] = value
        self._stabilised_weights.move_to_end(key)
        while len(self._stabilised_weights) > self.window:
            self._stabilised_weights.popitem(last=False)

    def observe_summary(
        self,
        fidelity: float,
        merged: int,
        total: int,
        *,
        distortion: float | None = None,
        distortion_bound: float | None = None,
    ) -> float:
        fidelity = float(np.clip(fidelity, 0.0, 1.0))
        utilisation = float(total) / max(total + merged, 1)
        contraction = float(np.clip(utilisation, 0.0, 1.0))
        self._summary_fidelity.append(fidelity)
        self._memory_utilisation.append(utilisation)
        self._contraction_factors.append(contraction)
        if distortion is not None:
            clipped = float(np.clip(distortion, 0.0, 2.0))
            self._summary_distortion.append(clipped)
            margin = float(fidelity - clipped)
            self._lyapunov_margins.append(margin)
        else:
            self._lyapunov_margins.append(float(fidelity))
        if distortion_bound is not None:
            bound = float(max(0.0, distortion_bound))
            self._distortion_bounds.append(bound)
        return fidelity

    def observe_trust(self, trusts: Iterable[float]) -> float:
        trusts_arr = np.asarray(list(trusts), dtype=float)
        if trusts_arr.size == 0:
            floor = 0.0
        else:
            floor = float(np.min(trusts_arr))
        self._trust_floor.append(floor)
        return floor

    # ------------------------------------------------------------------ #
    def stability(self) -> StabilitySnapshot:
        retrieval_entropy = float(
            np.mean(self._retrieval_entropy) if self._retrieval_entropy else 0.0)
        summary_fidelity = float(
            np.mean(self._summary_fidelity) if self._summary_fidelity else 0.0)
        summary_distortion = float(
            np.mean(self._summary_distortion) if self._summary_distortion else 0.0
        )
        summary_distortion_bound = float(
            np.mean(self._distortion_bounds) if self._distortion_bounds else 0.0
        )
        trust_floor = float(np.mean(self._trust_floor)
                            if self._trust_floor else 0.0)
        utilisation = float(np.mean(self._memory_utilisation)
                            if self._memory_utilisation else 0.0)
        contraction_factor = float(
            np.mean(
                self._contraction_factors) if self._contraction_factors else utilisation)
        lyapunov_margin = float(
            np.mean(
                self._lyapunov_margins) if self._lyapunov_margins else summary_fidelity)
        lyapunov_warning = bool(lyapunov_margin < 0.0)

        stability_score = float(
            0.3 * summary_fidelity
            + 0.2 * trust_floor
            + 0.2 * contraction_factor
            + 0.15 * (1.0 / (1.0 + retrieval_entropy))
            + 0.15 * max(0.0, 1.0 + lyapunov_margin)
        )

        return StabilitySnapshot(
            retrieval_entropy=retrieval_entropy,
            summary_fidelity=summary_fidelity,
            summary_distortion=summary_distortion,
            summary_distortion_bound=summary_distortion_bound,
            trust_floor=trust_floor,
            memory_utilisation=utilisation,
            contraction_factor=contraction_factor,
            lyapunov_margin=lyapunov_margin,
            lyapunov_warning=lyapunov_warning,
            stability_score=stability_score,
        )
