"""Transformer 记忆预测服务：回忆概率、遗忘曲线、复习间隔建议。"""

from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Optional

import numpy as np
from sqlalchemy.orm import Session

from models import QuizRecord, User, Word
from services.memory_transformer.encoding import (
    build_sequence_matrix,
    parse_iso,
    word_en,
)
from services.memory_transformer.network import MiniTransformer, init_random_weights
from services.memory_visual_service import retention_percent, stability_hours

WEIGHTS_PATH = Path(__file__).resolve().parents[2] / "data" / "memory_transformer" / "weights.json"
DEFAULT_WEIGHTS_PATH = Path(__file__).resolve().parents[2] / "data" / "memory_transformer" / "default_weights.json"


def _hours_since_review(word: Word, now: datetime) -> float:
    last_at = parse_iso(word.last_reviewed_at or word.created_at)
    return max(0.0, (now - last_at).total_seconds() / 3600.0)


def _formula_recall(word: Word, hours_since: float) -> float:
    return retention_percent(hours_since, stability_hours(word)) / 100.0


def _stability_days_from_target_r(target_r: float, hours_since: float) -> int:
    """达到目标回忆率所需的额外稳定化间隔（天），用于推荐复习。"""
    target_r = max(0.55, min(0.95, target_r))
    if hours_since <= 0:
        return 1
    s_hours = -hours_since / np.log(target_r)
    days = int(max(1, min(30, round(s_hours / 24.0))))
    return days


class MemoryTransformerService:
    def __init__(self) -> None:
        self._model: Optional[MiniTransformer] = None
        self._loaded_path: Optional[Path] = None

    def _load_model(self) -> MiniTransformer:
        path = WEIGHTS_PATH if WEIGHTS_PATH.is_file() else DEFAULT_WEIGHTS_PATH
        if self._model is not None and self._loaded_path == path:
            return self._model
        if path.is_file():
            self._model = MiniTransformer.load_json(path)
        else:
            m = MiniTransformer()
            m.weights = init_random_weights(
                d_model=m.d_model,
                n_heads=m.n_heads,
                d_ff=m.d_ff,
                n_layers=m.n_layers,
            )
            self._model = m
        self._loaded_path = path
        return self._model

    def model_ready(self) -> bool:
        return WEIGHTS_PATH.is_file()

    def predict_for_word(
        self,
        db: Session,
        user: User,
        word: Word,
        horizon_hours: Optional[list[float]] = None,
    ) -> dict:
        now = datetime.now(timezone.utc)
        records = (
            db.query(QuizRecord)
            .filter(QuizRecord.user_id == user.id, QuizRecord.word_id == word.id)
            .order_by(QuizRecord.created_at.asc())
            .all()
        )
        seq, valid_len = build_sequence_matrix(word, records, now)
        model = self._load_model()

        p_now = model.predict_proba(seq, valid_len)
        hours_since = _hours_since_review(word, now)
        p_formula = _formula_recall(word, hours_since)

        # 只有真实训练权重才参与融合；默认随机权重仅保留接口兼容性。
        n_events = max(0, valid_len - 1)
        blend = min(1.0, n_events / 5.0) if self.model_ready() else 0.0
        recall_now = round((blend * p_now + (1.0 - blend) * p_formula) * 100, 1)

        if horizon_hours is None:
            horizon_hours = [0, 6, 12, 24, 48, 72, 120, 168, 240, 336]
        curve = []
        for dh in horizon_hours:
            future = now + timedelta(hours=dh)
            seq_f, valid_f = build_sequence_matrix(word, records, future)
            p_f = model.predict_proba(seq_f, valid_f)
            h_total = hours_since + dh
            p_form_f = _formula_recall(word, h_total)
            recall = round((blend * p_f + (1.0 - blend) * p_form_f) * 100, 1)
            curve.append(
                {
                    "hours_ahead": dh,
                    "recall_percent": recall,
                    "forgetting_percent": round(100.0 - recall, 1),
                }
            )

        target_r = 0.75
        recommended_days = _stability_days_from_target_r(
            target_r, hours_since * max(0.3, recall_now / 100.0)
        )
        half_life_hours = None
        for pt in curve:
            if pt["recall_percent"] <= 50.0:
                half_life_hours = pt["hours_ahead"]
                break

        return {
            "word_id": word.id,
            "en": word_en(word),
            "recall_now_percent": recall_now,
            "formula_recall_percent": round(p_formula * 100, 1),
            "model_recall_percent": round(p_now * 100, 1),
            "blend_weight": round(blend, 2),
            "event_count": n_events,
            "model_trained": self.model_ready(),
            "recommended_review_days": recommended_days,
            "half_life_hours": half_life_hours,
            "curve": curve,
            "attention_hint": self._attention_hint(model, seq, valid_len),
        }

    def _attention_hint(
        self,
        model: MiniTransformer,
        seq: list[list[float]],
        valid_len: int,
    ) -> list[dict]:
        """返回 CLS 对序列位置的关注度（便于解释「哪些 token 影响预测」）。"""
        if valid_len <= 1:
            return [{"index": 0, "label": "词项", "weight": 1.0}]
        x = np.array(seq[:valid_len], dtype=np.float64)
        w = model.weights
        x = x @ w["in_proj"] + w["in_bias"]
        Q = x[0:1] @ w["L0.Wq"] + w["L0.Bq"]
        K = x @ w["L0.Wk"] + w["L0.Bk"]
        sl = slice(0, model.d_k)
        scores = (Q[:, sl] @ K[:, sl].T)[0] / np.sqrt(model.d_k)
        attn = np.exp(scores - scores.max())
        attn = attn / attn.sum()
        labels = ["词项(CLS)"] + [f"事件{i}" for i in range(1, valid_len)]
        return [
            {"index": i, "label": labels[i], "weight": round(float(attn[i]), 3)}
            for i in range(valid_len)
        ]


memory_transformer_service = MemoryTransformerService()