"""
EzVibe 情绪状态机
=================
设计文档对应章节：核心模块结构 - 情绪系统（状态转移模型）

状态集合
    S = {idle, happy, focused, annoyed, sleepy}

状态转移矩阵 P (5×5)
    P[i][j] = P(s_{t+1}=j | s_t=i)
    按索引顺序: idle, happy, focused, annoyed, sleepy

上下文调制
    P'_{ij} = P_{ij} * f(C)
    f(C) 为事件驱动增益因子

实现策略
    1. 概率归一化 (Softmax)
    2. 蒙特卡洛随机采样
    3. 最小驻留时间限制（防止情绪频繁抖动）
"""

from __future__ import annotations

import logging
import math
import random
import time
from enum import Enum
from typing import Optional

import numpy as np

logger = logging.getLogger(__name__)


# ================================================================
# 1. 状态定义
# ================================================================

class EmotionState(str, Enum):
    """
    五种情绪状态。

    idle    — 待机/发呆
    happy   — 开心
    focused — 专注（正在工作/思考）
    annoyed — 烦躁（被打断/久坐未活动）
    sleepy  — 困倦/疲劳
    """
    IDLE    = "idle"
    HAPPY   = "happy"
    FOCUSED = "focused"
    ANNOYED = "annoyed"
    SLEEPY  = "sleepy"

    @classmethod
    def from_string(cls, s: str) -> "EmotionState":
        """字符串 → EmotionState，兼容大小写。"""
        s = s.strip().lower()
        for state in cls:
            if state.value == s:
                return state
        raise ValueError(f"Unknown emotion state: {s!r}")

    @classmethod
    def all_values(cls) -> list[str]:
        return [s.value for s in cls]


# ================================================================
# 2. 默认转移矩阵（文档原生数据）
# ================================================================

# 顺序: idle, happy, focused, annoyed, sleepy
DEFAULT_TRANSITION_MATRIX: list[list[float]] = [
    # idle   →  idle   happy  focused annoyed sleepy
    [0.4,    0.2,   0.2,   0.1,   0.1  ],  # from idle
    [0.2,    0.5,   0.1,   0.1,   0.1  ],  # from happy
    [0.1,    0.2,   0.5,   0.1,   0.1  ],  # from focused
    [0.1,    0.1,   0.1,   0.6,   0.1  ],  # from annoyed
    [0.2,    0.1,   0.1,   0.1,   0.5  ],  # from sleepy
]

# 状态索引映射（避免魔法数字）
_STATE_INDEX: dict[EmotionState, int] = {
    EmotionState.IDLE:    0,
    EmotionState.HAPPY:   1,
    EmotionState.FOCUSED: 2,
    EmotionState.ANNOYED: 3,
    EmotionState.SLEEPY:  4,
}
_INDEX_TO_STATE: list[EmotionState] = [
    EmotionState.IDLE,
    EmotionState.HAPPY,
    EmotionState.FOCUSED,
    EmotionState.ANNOYED,
    EmotionState.SLEEPY,
]


# ================================================================
# 3. 上下文调制因子
# ================================================================

class ContextBoost:
    """
    事件驱动权重调整因子。

    每次调用 update(event) 时，根据事件类型对转移矩阵
    相应行的目标状态概率进行增强。

    使用加法增益：
        boost_factor[to_state] += gain

    然后在采样前用 Softmax 重新归一化该行。
    """

    # 事件 → {目标状态: 增益值}
    # 增益越大，切换到该状态的概率越高
    EVENT_BOOSTS: dict[str, dict[EmotionState, float]] = {
        # 长时间高频工作 → sleepy
        "long_work_session": {
            EmotionState.SLEEPY:   2.0,
            EmotionState.FOCUSED:  0.8,
        },
        # 用户处于专注软件（IDE/文档）→ focused
        "user_focused": {
            EmotionState.FOCUSED:  2.5,
            EmotionState.IDLE:     0.3,
            EmotionState.ANNOYED:  0.2,
        },
        # 频繁打扰/无视健康提醒 → annoyed
        "reminder_ignored": {
            EmotionState.ANNOYED:  3.0,
            EmotionState.HAPPY:    0.1,
        },
        # 用户表扬/夸奖 → happy
        "user_praise": {
            EmotionState.HAPPY:    3.0,
            EmotionState.ANNOYED:  0.1,
        },
        # 用户主动交互 → 唤醒 idle → happy
        "user_interact": {
            EmotionState.HAPPY:    1.5,
            EmotionState.IDLE:     1.0,
            EmotionState.ANNOYED:  0.3,
        },
        # 久坐检测 → sleepy + annoyed
        "sedentary_too_long": {
            EmotionState.SLEEPY:   1.5,
            EmotionState.ANNOYED:  1.0,
            EmotionState.IDLE:     0.5,
        },
        # 正常待机时间流逝 → idle
        "time_passes": {
            EmotionState.IDLE:     1.2,
            EmotionState.HAPPY:    0.5,
        },
        # 用户喝了水/运动 → happy
        "user_healthy_action": {
            EmotionState.HAPPY:    2.0,
            EmotionState.ANNOYED:  0.1,
            EmotionState.SLEEPY:    0.3,
        },
    }

    @classmethod
    def get_boosts(cls, event: str) -> dict[EmotionState, float]:
        """根据事件类型返回增益字典。"""
        return cls.EVENT_BOOSTS.get(event, {})


# ================================================================
# 4. 核心引擎
# ================================================================

class EmotionEngine:
    """
    情绪状态机引擎。

    参数
    ----
    transition_matrix : list[list[float]] | None
        5×5 状态转移矩阵，行为当前状态，列为目标状态。
        默认使用文档中的矩阵。
    min_residence_seconds : float
        最小驻留时间（秒）。在此时间内不允许切换状态，防止抖动。
        默认 5 秒。
    seed : int | None
        随机种子，用于单元测试复现。

    示例
    ----
    >>> engine = EmotionEngine()
    >>> engine.get_state()
    'idle'
    >>> engine.update("user_praise")
    'happy'
    >>> engine.get_state()
    'happy'
    """

    def __init__(
        self,
        transition_matrix: list[list[float]] | None = None,
        min_residence_seconds: float = 5.0,
        seed: int | None = None,
    ) -> None:
        # ── 转移矩阵验证 ──────────────────────────────────────
        raw = transition_matrix if transition_matrix is not None else DEFAULT_TRANSITION_MATRIX
        try:
            self._P = np.array(raw, dtype=np.float64)
        except Exception as exc:
            raise ValueError(
                f"transition_matrix must be a 5×5 matrix, got {type(raw).__name__}"
            ) from exc

        if self._P.shape != (5, 5):
            raise ValueError(
                f"transition_matrix must be shape (5,5), got {self._P.shape}"
            )

        # 确保每行归一化为概率分布（浮点误差容差 1e-6）
        row_sums = self._P.sum(axis=1, keepdims=True)
        if np.any(row_sums == 0):
            raise ValueError("transition_matrix rows must not sum to zero")
        self._P = self._P / row_sums

        # ── 状态初始化 ──────────────────────────────────────
        self._current_state: EmotionState = EmotionState.IDLE
        self._state_since: float = time.monotonic()   # 当前状态开始的绝对时间
        self._min_residence: float = min_residence_seconds

        # ── 随机数 ─────────────────────────────────────────
        self._rng = random.Random(seed)
        self._np_rng = np.random.default_rng(seed)

        # ── 事件历史（用于调试/分析）────────────────────────
        self._history: list[dict] = []

        logger.info(
            "EmotionEngine initialized | state=%s | min_residence=%.1fs",
            self._current_state.value,
            self._min_residence,
        )

    # ----------------------------------------------------------------
    # 公开 API
    # ----------------------------------------------------------------

    def get_state(self) -> str:
        """返回当前情绪状态字符串。"""
        return self._current_state.value

    def get_state_object(self) -> EmotionState:
        """返回当前情绪状态（EmotionState 枚举）。"""
        return self._current_state

    def get_display_name(self, state: str | None = None) -> str:
        """
        返回情绪状态的中文/英文友好名称。

        参数
        ----
        state : str | None
            目标状态字符串，None 则返回当前状态对应的名称。
        """
        display_map: dict[str, tuple[str, str]] = {
            "idle":    ("待机", "Idle"),
            "happy":   ("开心", "Happy"),
            "focused": ("专注", "Focused"),
            "annoyed": ("烦躁", "Annoyed"),
            "sleepy":  ("困倦", "Sleepy"),
        }
        key = (state if state is not None else self._current_state.value).lower()
        zh, en = display_map.get(key, (key, key))
        return f"{zh} ({en})"

    def get_all_states(self) -> list[str]:
        """返回所有状态名称列表。"""
        return EmotionState.all_values()

    def update(self, event: str, context: dict | None = None) -> str:
        """
        根据事件触发一次状态转移。

        流程
        1. 检查驻留时间，若不足则拒绝转移。
        2. 查找事件对应的增益因子。
        3. 将增益叠加到当前行的转移概率。
        4. Softmax 归一化。
        5. 蒙特卡洛采样选出下一状态。
        6. 更新状态并记录历史。

        参数
        ----
        event : str
            事件类型，需匹配 ContextBoost.EVENT_BOOSTS 中的 key。
        context : dict | None
            可选上下文信息（目前预留，供未来扩展）。

        返回
        ----
        str
            新的情绪状态字符串。
        """
        now = time.monotonic()
        elapsed = now - self._state_since

        # ── 驻留时间检查 ────────────────────────────────────
        if elapsed < self._min_residence:
            logger.debug(
                "State transition blocked by residence time "
                "(elapsed=%.2fs < %.2fs, event=%s)",
                elapsed, self._min_residence, event,
            )
            return self._current_state.value

        prev_state = self._current_state
        row_idx = _STATE_INDEX[prev_state]

        # ── 获取事件增益 ────────────────────────────────────
        boosts = ContextBoost.get_boosts(event)

        # ── 构建带增益的概率向量 ────────────────────────────
        probs = self._P[row_idx].copy()

        for target_state, gain in boosts.items():
            col_idx = _STATE_INDEX[target_state]
            probs[col_idx] += gain

        # ── Softmax 归一化 ─────────────────────────────────
        probs = self._softmax(probs)

        # ── 蒙特卡洛采样 ───────────────────────────────────
        next_state = self._sample(probs)

        # ── 更新状态 ───────────────────────────────────────
        if next_state != prev_state:
            self._current_state = next_state
            self._state_since = now
            logger.info(
                "State transition: %s → %s (event=%s, boosts=%s)",
                prev_state.value, next_state.value, event, boosts,
            )
        else:
            logger.debug(
                "State unchanged: %s (event=%s, probs=%s)",
                prev_state.value, event, probs.round(3).tolist(),
            )

        # ── 记录历史 ───────────────────────────────────────
        self._history.append({
            "timestamp":    now,
            "event":        event,
            "prev_state":   prev_state.value,
            "curr_state":   self._current_state.value,
            "transition":   prev_state != self._current_state,
            "elapsed":      elapsed,
            "boosts":       boosts,
        })

        # 历史保留最近 1000 条
        if len(self._history) > 1000:
            self._history = self._history[-1000:]

        return self._current_state.value

    def get_residence_time(self) -> float:
        """返回当前状态已持续的时间（秒）。"""
        return time.monotonic() - self._state_since

    def is_residence_ready(self) -> bool:
        """返回当前状态是否已满足最小驻留时间。"""
        return self.get_residence_time() >= self._min_residence

    def get_transition_matrix(self) -> list[list[float]]:
        """返回原始转移矩阵（列表形式）。"""
        return self._P.tolist()

    def get_transition_probabilities(self) -> dict[str, float]:
        """
        返回当前状态到所有状态的转移概率。
        诊断用。
        """
        row_idx = _STATE_INDEX[self._current_state]
        probs = self._P[row_idx]
        return {
            _INDEX_TO_STATE[i].value: float(probs[i])
            for i in range(5)
        }

    def get_history(self, last_n: int = 50) -> list[dict]:
        """返回最近 N 条状态历史。"""
        return self._history[-last_n:]

    def force_state(self, state: str | EmotionState) -> None:
        """
        强制设置状态（用于重置/测试，不触发状态转移逻辑）。
        """
        if isinstance(state, str):
            state = EmotionState.from_string(state)
        self._current_state = state
        self._state_since = time.monotonic()
        logger.info("State forcibly set to: %s", state.value)

    def tick(self) -> str:
        """
        无事件驱动的时钟推进。

        当无事发生时，每隔一定时间（建议 60s）调用一次，
        触发自然的状态回归/衰减。
        """
        return self.update("time_passes")

    # ----------------------------------------------------------------
    # 内部实现
    # ----------------------------------------------------------------

    @staticmethod
    def _softmax(x: np.ndarray) -> np.ndarray:
        """
        Numerically stable Softmax。

        softmax_i = exp(x_i - max(x)) / sum(exp(x_j - max(x)))
        """
        x = np.asarray(x, dtype=np.float64)
        x_shifted = x - x.max()
        exp_x = np.exp(x_shifted)
        return exp_x / exp_x.sum()

    def _sample(self, probs: np.ndarray) -> EmotionState:
        """
        蒙特卡洛采样：根据概率向量随机选取下一状态。

        算法：累计分布函数（CDF）+ 二分搜索。
        等价于 numpy.random.choice，效率更高。
        """
        probs = np.asarray(probs, dtype=np.float64)
        if not np.isclose(probs.sum(), 1.0):
            probs = probs / probs.sum()  # 容错

        cumsum = np.cumsum(probs)
        r = self._rng.random()  # [0, 1)
        idx = int(np.searchsorted(cumsum, r))
        # searchsorted 可能返回 n（越界），取模安全兜底
        idx = min(idx, len(_INDEX_TO_STATE) - 1)
        return _INDEX_TO_STATE[idx]

    # ----------------------------------------------------------------
    # 序列化（用于持久化）
    # ----------------------------------------------------------------

    def to_dict(self) -> dict:
        """将当前状态导出为可序列化字典。"""
        return {
            "current_state":    self._current_state.value,
            "state_since":      self._state_since,
            "min_residence":    self._min_residence,
            "transition_matrix": self._P.tolist(),
        }

    @classmethod
    def from_dict(cls, d: dict) -> "EmotionEngine":
        """从字典恢复引擎实例（保留状态，矩阵可能不同）。"""
        engine = cls(
            transition_matrix=d.get("transition_matrix", DEFAULT_TRANSITION_MATRIX),
            min_residence_seconds=d.get("min_residence", 5.0),
        )
        engine.force_state(d["current_state"])
        engine._state_since = d.get("state_since", time.monotonic())
        return engine