"""
Entropy Monitor - Measures system entropy using multiple metrics
"""

from __future__ import annotations

import math
import statistics
from collections import Counter
from dataclasses import dataclass, field
from datetime import datetime
from typing import Any, Dict, List, Optional, Tuple


@dataclass(frozen=True)
class EntropyMonitorConfig:
    """Configuration for EntropyMonitor.

    Goals:
    - Keep classic metrics stable (decision/dispersion/communication/combined in [0,1]).
    - Add calibratable signals: baseline per workload and z-scores.
    - Avoid "no-signal" saturation (e.g., always returning 0.0 due to missing agent fields).
    """

    # Combined entropy weights (must sum ~1.0). Defaults kept for backward compat.
    weights: Dict[str, float] = field(
        default_factory=lambda: {
            "decision": 0.4,
            "dispersion": 0.3,
            "communication": 0.3,
        }
    )

    # History handling.
    max_history_size: int = 100
    baseline_window: int = 50
    min_baseline_samples: int = 10

    # Numerical stability for z-scores.
    epsilon_std: float = 1e-6
    zscore_clip: float = 6.0

    # Signal quality thresholds.
    min_overall_coverage: float = 0.5

    # Health scoring.
    combined_healthy_range: Tuple[float, float] = (0.2, 0.7)
    health_weights: Dict[str, float] = field(
        default_factory=lambda: {
            "range": 0.4,
            "decision_z": 0.2,
            "dispersion_z": 0.2,
            "communication_z": 0.2,
        }
    )


class EntropyMonitor:
    """Monitors and measures entropy across multiple dimensions."""

    def __init__(self, config: Optional[EntropyMonitorConfig] = None):
        self.config = config or EntropyMonitorConfig()

        # Classic history API (backward compatible): stores the default workload.
        self.metrics_history: List[Dict[str, Any]] = []
        self.last_measurement: Optional[Dict[str, Any]] = None

        # Internal histories keyed by workload id for calibration.
        self._workload_histories: Dict[str, List[Dict[str, Any]]] = {}

    def measure_system_entropy(
        self,
        agents_state: List[Any],
        *,
        workload_id: Optional[str] = None,
        timestamp: Optional[datetime] = None,
        record: bool = True,
    ) -> Dict[str, Any]:
        """Measure system entropy.

        Classic metrics:
        - decision: Shannon entropy of decisions
        - dispersion: state dispersion
        - communication: communication complexity
        - combined: weighted average

        Calibratable additions:
        - baseline per workload (rolling window) and z-scores
        - signal quality flags to avoid "no signal" appearing as "perfectly stable"
        - health score composed from range + anomaly signals
        """

        wid = workload_id or "default"
        ts = timestamp or datetime.now()

        # Extract data from agents.
        decisions: List[Any] = []
        states: List[Any] = []
        messages: List[Any] = []
        got_decision = 0
        got_state = 0
        got_messages = 0

        for agent in agents_state:
            d, ok_d = self._extract_decision(agent)
            s, ok_s = self._extract_state(agent)
            msg, ok_m = self._extract_messages(agent)

            if ok_d:
                got_decision += 1
                decisions.append(d)
            if ok_s:
                got_state += 1
                states.append(s)
            if ok_m:
                got_messages += 1
                messages.append(msg)

        agent_count = len(agents_state)
        decision_cov = (got_decision / agent_count) if agent_count else 0.0
        state_cov = (got_state / agent_count) if agent_count else 0.0
        msg_cov = (got_messages / agent_count) if agent_count else 0.0
        overall_cov = (decision_cov + state_cov + msg_cov) / 3.0 if agent_count else 0.0

        issues: List[str] = []
        if agent_count > 0 and overall_cov < self.config.min_overall_coverage:
            issues.append("low_signal_coverage")
        if agent_count > 0 and got_decision == 0 and got_state == 0 and got_messages == 0:
            issues.append("no_signal")

        # Calculate individual metrics (0..1). If we have no usable signal, return neutral
        # values (0.5) and avoid recording the sample as baseline training data.
        if "no_signal" in issues:
            decision_entropy = 0.5
            dispersion = 0.5
            communication = 0.5
        else:
            decision_entropy = self._shannon_entropy(decisions)
            dispersion = self._calculate_dispersion(states)
            communication = self._communication_complexity(messages)

        # Combined entropy (weighted average, configurable). Defaults preserved.
        w = self.config.weights
        w_sum = float(
            w.get("decision", 0.0) + w.get("dispersion", 0.0) + w.get("communication", 0.0)
        )
        if w_sum <= 0:
            w = {"decision": 0.4, "dispersion": 0.3, "communication": 0.3}
            w_sum = 1.0

        combined = (
            decision_entropy * float(w.get("decision", 0.0))
            + dispersion * float(w.get("dispersion", 0.0))
            + communication * float(w.get("communication", 0.0))
        ) / w_sum

        baseline = self._compute_baseline(wid)
        z = self._z_scores(
            {
                "decision": decision_entropy,
                "dispersion": dispersion,
                "communication": communication,
                "combined": combined,
            },
            baseline,
        )
        health = self._health_score(
            combined=combined,
            z=z,
            signal_ok=("no_signal" not in issues),
            coverage=overall_cov,
            issues=issues,
        )

        result: Dict[str, Any] = {
            "decision": float(min(max(decision_entropy, 0.0), 1.0)),
            "dispersion": float(min(max(dispersion, 0.0), 1.0)),
            "communication": float(min(max(communication, 0.0), 1.0)),
            "combined": float(min(max(combined, 0.0), 1.0)),
            "timestamp": ts,
            "workload_id": wid,
            "signal": {
                "agent_count": agent_count,
                "decision_coverage": float(decision_cov),
                "state_coverage": float(state_cov),
                "message_coverage": float(msg_cov),
                "overall_coverage": float(overall_cov),
                "issues": list(issues),
                "signal_ok": ("no_signal" not in issues),
            },
            "baseline": baseline,
            "z": z,
            "health": health,
        }

        if record:
            self._record_measurement(wid, result)

        self.last_measurement = result
        return result

    def _shannon_entropy(self, decisions: List[Any]) -> float:
        """Calculates Shannon entropy of decisions."""
        if not decisions:
            return 0.0

        # Count occurrences.
        counts = Counter(str(d) for d in decisions)
        total = len(decisions)

        # Calculate Shannon entropy.
        entropy = 0.0
        for count in counts.values():
            p = count / total
            if p > 0:
                entropy -= p * math.log2(p)

        # Normalize to 0-1 range (assuming max entropy for uniform distribution).
        max_entropy = math.log2(len(counts)) if len(counts) > 1 else 1.0
        normalized = entropy / max_entropy if max_entropy > 0 else 0.0

        return min(normalized, 1.0)

    def _calculate_dispersion(self, states: List[Any]) -> float:
        """Calculates how dispersed agent states are."""
        if not states:
            return 0.0

        try:
            # Convert states to numeric values.
            numeric_states: List[float] = []
            for state in states:
                if isinstance(state, (int, float)):
                    numeric_states.append(float(state))
                elif hasattr(state, "__dict__"):
                    # Use hash of state dict as numeric representation.
                    numeric_states.append(hash(str(state.__dict__)) % 1000 / 1000.0)
                else:
                    numeric_states.append(hash(str(state)) % 1000 / 1000.0)

            # Calculate standard deviation (normalized).
            if len(numeric_states) > 1:
                std_dev = statistics.pstdev(numeric_states)
                # Normalize assuming states are in 0-1 range.
                return min(std_dev * 2, 1.0)
            return 0.0

        except Exception:
            # Fallback to simple diversity measure.
            unique_states = len(set(str(s) for s in states))
            return unique_states / len(states) if states else 0.0

    def _communication_complexity(self, messages: List[Any]) -> float:
        """Measures communication complexity/volume."""
        total_messages = sum(len(msg) if isinstance(msg, list) else 1 for msg in messages)

        # Normalize: 0 messages = 0 entropy, 100+ messages = 1.0 entropy.
        normalized = min(total_messages / 100.0, 1.0)
        return normalized

    def _extract_decision(self, agent: Any) -> Tuple[Any, bool]:
        """Extracts last decision from agent."""
        if isinstance(agent, dict):
            if "last_decision" in agent:
                return agent.get("last_decision"), True
            if "last_action" in agent:
                return agent.get("last_action"), True

        if hasattr(agent, "last_decision"):
            return getattr(agent, "last_decision"), True
        if hasattr(agent, "last_action"):
            return getattr(agent, "last_action"), True
        return None, False

    def _extract_state(self, agent: Any) -> Tuple[Any, bool]:
        """Extracts current state from agent."""
        if isinstance(agent, dict):
            if "current_state" in agent:
                return agent.get("current_state"), True
            if "state" in agent:
                return agent.get("state"), True

        if hasattr(agent, "current_state"):
            return getattr(agent, "current_state"), True
        if hasattr(agent, "state"):
            return getattr(agent, "state"), True
        return None, False

    def _extract_messages(self, agent: Any) -> Tuple[List[Any], bool]:
        """Extracts messages from agent."""
        if isinstance(agent, dict):
            if "messages" in agent:
                msg = agent.get("messages")
                if msg is None:
                    return [], False
                return (msg if isinstance(msg, list) else [msg]), True
            if "message_history" in agent:
                msg = agent.get("message_history")
                if msg is None:
                    return [], False
                return (msg if isinstance(msg, list) else [msg]), True

        if hasattr(agent, "messages"):
            msg = getattr(agent, "messages")
            if msg is None:
                return [], False
            return (msg if isinstance(msg, list) else [msg]), True
        if hasattr(agent, "message_history"):
            msg = getattr(agent, "message_history")
            if msg is None:
                return [], False
            return (msg if isinstance(msg, list) else [msg]), True
        return [], False

    def get_history(self, last_n: int = None) -> List[Dict[str, float]]:
        """Returns entropy history."""
        if last_n:
            return self.metrics_history[-last_n:]
        return self.metrics_history

    def get_trend(self) -> str:
        """Analyzes recent trend in entropy."""
        if len(self.metrics_history) < 3:
            return "insufficient_data"

        recent = [m["combined"] for m in self.metrics_history[-5:]]

        if recent[-1] > recent[0] + 0.1:
            return "increasing"
        if recent[-1] < recent[0] - 0.1:
            return "decreasing"
        return "stable"

    def measure_entropy(self, agents_state: List[Any]) -> float:
        """Measure entropy and return combined value as float (backward compatible)."""
        metrics = self.measure_system_entropy(agents_state)
        return float(metrics["combined"])

    def _record_measurement(self, workload_id: str, result: Dict[str, Any]) -> None:
        # Always keep classic history for inspection.
        self.metrics_history.append(result)

        # Do not train baseline on "no_signal" samples.
        if result.get("signal", {}).get("signal_ok", True):
            self._workload_histories.setdefault(workload_id, []).append(result)

        # Cap classic history.
        if len(self.metrics_history) > int(self.config.max_history_size):
            self.metrics_history = self.metrics_history[-int(self.config.max_history_size) :]

        # Cap per-workload history used for calibration.
        hist = self._workload_histories.get(workload_id)
        if hist is not None and len(hist) > int(self.config.max_history_size):
            self._workload_histories[workload_id] = hist[-int(self.config.max_history_size) :]

    def _compute_baseline(self, workload_id: str) -> Dict[str, Any]:
        hist = self._workload_histories.get(workload_id, [])
        window = hist[-int(self.config.baseline_window) :]
        n = len(window)
        if n <= 0:
            return {
                "ready": False,
                "n": 0,
                "mean": {},
                "std": {},
                "window": int(self.config.baseline_window),
            }

        def _col(key: str) -> List[float]:
            vals: List[float] = []
            for r in window:
                v = r.get(key)
                if isinstance(v, (int, float)):
                    vals.append(float(v))
            return vals

        means: Dict[str, float] = {}
        stds: Dict[str, float] = {}
        for k in ("decision", "dispersion", "communication", "combined"):
            xs = _col(k)
            if not xs:
                continue
            means[k] = float(statistics.fmean(xs))
            if len(xs) >= 2:
                std = float(statistics.pstdev(xs))
            else:
                std = 0.0
            stds[k] = max(std, float(self.config.epsilon_std))

        ready = n >= int(self.config.min_baseline_samples)
        return {
            "ready": bool(ready),
            "n": int(n),
            "mean": means,
            "std": stds,
            "window": int(self.config.baseline_window),
        }

    def _z_scores(self, values: Dict[str, float], baseline: Dict[str, Any]) -> Dict[str, float]:
        if not baseline.get("ready"):
            return {"decision": 0.0, "dispersion": 0.0, "communication": 0.0, "combined": 0.0}

        means = baseline.get("mean", {}) or {}
        stds = baseline.get("std", {}) or {}
        z: Dict[str, float] = {}

        for k in ("decision", "dispersion", "communication", "combined"):
            v = float(values.get(k, 0.0))
            mu = float(means.get(k, v))
            sd = float(stds.get(k, float(self.config.epsilon_std)))
            zv = (v - mu) / sd if sd > 0 else 0.0

            clip = float(self.config.zscore_clip)
            if zv > clip:
                zv = clip
            elif zv < -clip:
                zv = -clip

            z[k] = float(zv)

        return z

    def _health_score(
        self,
        *,
        combined: float,
        z: Dict[str, float],
        signal_ok: bool,
        coverage: float,
        issues: List[str],
    ) -> Dict[str, Any]:
        # Range health: best when combined is within a healthy band; fades to 0 outside.
        lo, hi = self.config.combined_healthy_range
        c = float(combined)

        if c < lo and lo > 0:
            range_health = max(0.0, 1.0 - (lo - c) / lo)
        elif c > hi and hi < 1.0:
            range_health = max(0.0, 1.0 - (c - hi) / (1.0 - hi))
        else:
            range_health = 1.0

        def z_health(key: str) -> float:
            return 1.0 / (1.0 + abs(float(z.get(key, 0.0))))

        components = {
            "range_health": float(range_health),
            "decision_z_health": float(z_health("decision")),
            "dispersion_z_health": float(z_health("dispersion")),
            "communication_z_health": float(z_health("communication")),
        }

        w = self.config.health_weights
        w_range = float(w.get("range", 0.0))
        w_dec = float(w.get("decision_z", 0.0))
        w_disp = float(w.get("dispersion_z", 0.0))
        w_comm = float(w.get("communication_z", 0.0))
        w_sum = w_range + w_dec + w_disp + w_comm
        if w_sum <= 0:
            w_range, w_dec, w_disp, w_comm = 0.4, 0.2, 0.2, 0.2
            w_sum = 1.0

        score = (
            components["range_health"] * w_range
            + components["decision_z_health"] * w_dec
            + components["dispersion_z_health"] * w_disp
            + components["communication_z_health"] * w_comm
        ) / w_sum

        # Penalize extremely low coverage even if metrics look fine.
        if not signal_ok:
            score *= 0.5
        if "low_signal_coverage" in issues:
            score *= max(0.2, float(coverage))

        score = float(min(max(score, 0.0), 1.0))

        return {
            "score": score,
            "components": components,
            "coverage": float(coverage),
            "issues": list(issues),
            "explain": (
                "Health score combines an absolute healthy-range check for combined entropy "
                "with anomaly detection (z-scores) per metric. Low or absent signal reduces the score."
            ),
        }
