"""
# flake8: noqa
metrics.py
-----------
Quantitative diagnostics and theoretical validation metrics for the HIEROMEM core engine.

Implements analytic and diagnostic utilities to evaluate whether
the memory system satisfies the theoretical constraints of the
framework, including:

  • Bounded growth (Theorem 3)
  • Trust convergence (Lemma 1)
  • Summarisation fidelity (Theorem 2)
  • Global stability (Theorem 4, via GSR)
  • Cognitive activity ratios (Insert / Retrieve / Forget frequencies)
  • Retrieval and narrative continuity fidelity (empirical diagnostics)

This module is diagnostic-only — it has no side effects.
It interfaces with `Memory`, `EventLog`, and `GovernanceSafetyRegion`
to provide real-time indicators of systemic health.

Author: Sai Rishi Kiran Mannava
"""

from __future__ import annotations
from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import Any, Optional, Sequence

import numpy as np

from hieromem_core.events import EventLog, EventType
from hieromem_core.governance import GovernanceSafetyRegion
from hieromem_core.memory import Memory


# --------------------------------------------------------------------------- #
# Data structure for metric snapshot
# --------------------------------------------------------------------------- #
@dataclass
class MetricSnapshot:
    timestamp: datetime
    bounded_growth_ratio: float
    trust_mean: float
    trust_std: float
    trust_min: float
    trust_max: float
    summarisation_frequency: float
    forgetting_rate: float
    stability_score: float
    risk_level: str
    total_episodes: int

    def to_dict(self) -> dict[str, Any]:
        return {
            "timestamp": self.timestamp.isoformat(),
            "bounded_growth_ratio": self.bounded_growth_ratio,
            "trust_mean": self.trust_mean,
            "trust_std": self.trust_std,
            "trust_min": self.trust_min,
            "trust_max": self.trust_max,
            "summarisation_frequency": self.summarisation_frequency,
            "forgetting_rate": self.forgetting_rate,
            "stability_score": self.stability_score,
            "risk_level": self.risk_level,
            "total_episodes": self.total_episodes,
        }


# --------------------------------------------------------------------------- #
# Core Metric Engine
# --------------------------------------------------------------------------- #
class MemoryMetrics:
    """Computes theoretical and empirical metrics for a running memory instance."""

    def __init__(
        self,
        memory: Memory,
        event_log: Optional[EventLog] = None,
        gsr: Optional[GovernanceSafetyRegion] = None,
        window_minutes: int = 10,
    ):
        self.memory = memory
        self.event_log = event_log or getattr(memory, "event_log", None)
        self.gsr = gsr or getattr(memory, "gsr", None)
        self.window_minutes = window_minutes

    # ------------------------------------------------------------------ #
    def compute(self) -> MetricSnapshot:
        """Compute a live snapshot of theoretical + empirical metrics."""
        episodes = getattr(self.memory, "episodes", [])
        Nmax = max(1, getattr(self.memory, "Nmax", len(episodes) or 1))
        trusts = np.array([ep.trust for ep in episodes]
                          ) if episodes else np.zeros(1)

        # --- (Theorem 3) Bounded Growth ---
        bounded_growth_ratio = len(episodes) / Nmax

        # --- (Lemma 1) Trust Convergence Stats ---
        trust_mean = float(np.mean(trusts))
        trust_std = float(np.std(trusts))
        trust_min = float(np.min(trusts))
        trust_max = float(np.max(trusts))

        # --- Cognitive Activity Ratios ---
        summarisation_frequency, forgetting_rate = self._activity_ratios()

        # --- (Theorem 4) Global Stability (GSR) ---
        stability_score = self.gsr.stability_score() if self.gsr else float("nan")
        risk_level = self.gsr.risk_level() if self.gsr else "unknown"

        return MetricSnapshot(
            timestamp=datetime.now(),
            bounded_growth_ratio=bounded_growth_ratio,
            trust_mean=trust_mean,
            trust_std=trust_std,
            trust_min=trust_min,
            trust_max=trust_max,
            summarisation_frequency=summarisation_frequency,
            forgetting_rate=forgetting_rate,
            stability_score=stability_score,
            risk_level=risk_level,
            total_episodes=len(episodes),
        )

    # ------------------------------------------------------------------ #
    def _activity_ratios(self) -> tuple[float, float]:
        """Compute summarisation and forgetting frequencies from recent EventLog."""
        if not self.event_log or not hasattr(self.event_log, "all"):
            return 0.0, 0.0

        window = timedelta(minutes=self.window_minutes)
        cutoff = datetime.now() - window
        recent = [e for e in self.event_log.all() if e.timestamp > cutoff]

        insert_count = sum(e.event_type == EventType.INSERT for e in recent)
        summarise_count = sum(
            e.event_type == EventType.SUMMARISE for e in recent)
        forget_count = sum(e.event_type == EventType.FORGET for e in recent)

        summarisation_frequency = summarise_count / max(1, insert_count)
        forgetting_rate = forget_count / \
            max(1, len(getattr(self.memory, "episodes", [])))
        return summarisation_frequency, forgetting_rate

    # ------------------------------------------------------------------ #
    def report(self, precision: int = 3) -> str:
        """Return formatted human-readable metric report."""
        snap = self.compute()
        return (
            f"\n[HIEROMEM METRICS @ {snap.timestamp.strftime('%H:%M:%S')}]\n"
            f"  Episodes:            {snap.total_episodes}\n"
            f"  Bounded Growth:      {snap.bounded_growth_ratio:.{precision}f}  (<= 1 expected)\n"
            f"  Trust Mean:          {snap.trust_mean:.{precision}f}\n"
            f"  Trust Std:           {snap.trust_std:.{precision}f}\n"
            f"  Trust Min / Max:     {snap.trust_min:.{precision}f} / {snap.trust_max:.{precision}f}\n"
            f"  Summarisation Freq:  {snap.summarisation_frequency:.{precision}f}\n"
            f"  Forgetting Rate:     {snap.forgetting_rate:.{precision}f}\n"
            f"  Stability Score:     {snap.stability_score:.{precision}f} ({snap.risk_level})\n")

    def to_dict(self) -> dict[str, Any]:
        """Return latest metrics as a dictionary."""
        return self.compute().to_dict()

    def bounded_growth_violation(self, tolerance: float = 1.05) -> bool:
        """Check if memory exceeds the theoretical capacity (Nmax)."""
        episodes = getattr(self.memory, "episodes", [])
        nmax_attr = getattr(self.memory, "Nmax", 1)
        nmax_value = nmax_attr if isinstance(
            nmax_attr, int) else len(episodes) or 1
        ratio = len(episodes) / max(1, nmax_value)
        return ratio > tolerance


# --------------------------------------------------------------------------- #
# Stand-alone theoretical validation functions
# --------------------------------------------------------------------------- #
def bounded_growth_check(memory: Memory) -> bool:
    """
    Theorem 3 — Bounded Growth:
    Memory size must asymptotically remain <= Nmax once summarisation is active.
    """
    n = len(getattr(memory, "episodes", []))
    nmax_attr = getattr(memory, "Nmax", None)
    if isinstance(nmax_attr, int):
        nmax = nmax_attr
    else:
        nmax = n if n else 1
    return n <= nmax


def trust_convergence_check(memory: Memory, epsilon: float = 0.05) -> bool:
    """
    Lemma 1 — Trust Convergence:
    After sufficient iterations, variance of trust values should converge below ε.
    """
    trusts = np.array([ep.trust for ep in getattr(memory, "episodes", [])])
    if len(trusts) < 2:
        return True
    variance = float(np.var(trusts))
    return variance < epsilon


def summarisation_fidelity_score(
    original_vecs: Sequence[np.ndarray], summary_vec: np.ndarray
) -> float:
    """
    Theorem 2 — Summarisation Fidelity:
    The summary embedding should stay within the convex hull of its constituents.
    We approximate this by measuring average cosine similarity.
    """
    if not original_vecs:
        return 0.0
    sims = [np.dot(v, summary_vec) / (np.linalg.norm(v) * \
                   np.linalg.norm(summary_vec)) for v in original_vecs]
    return float(np.mean(sims))


def retrieval_fidelity_score(
    query_vec: np.ndarray, recalled_vecs: Sequence[np.ndarray]
) -> float:
    """
    Empirical Recall Fidelity:
    Measures average cosine similarity between query and retrieved episodes.
    """
    if not recalled_vecs:
        return 0.0
    sims = [np.dot(query_vec, r) / (np.linalg.norm(query_vec)
                                    * np.linalg.norm(r)) for r in recalled_vecs]
    return float(np.mean(sims))


def narrative_continuity_score(memory: Memory) -> float:
    """
    Empirical Narrative Continuity:
    Measures temporal coherence by checking similarity of consecutive episode embeddings.
    """
    episodes = getattr(memory, "episodes", [])
    if len(episodes) < 2:
        return 1.0
    vecs = [ep.e for ep in episodes]
    sims = [
        np.dot(vecs[i], vecs[i + 1]) /
        (np.linalg.norm(vecs[i]) * np.linalg.norm(vecs[i + 1]) + 1e-8)
        for i in range(len(vecs) - 1)
    ]
    return float(np.mean(sims))
