# -*- coding: utf-8 -*-
"""Ti-1 氯化车间质量预测 · 模型训练与评估（Issue #69 / PRD §5.3 ①）。

承接 #68（特征工程）产出的 ``FeatureMatrix``：把"特征矩阵 + 质量标签 → 可评估、
可解释的预测模型"这条链路**模板化、可测试**，且与 #73 模型部署解耦。

PRD 设计口径
------------
- 架构表（PRD §5.3）：``质量预测 | 预测 | 入：特征矩阵；
  出：质量指标预测值(纯度/杂质) | ① 质量预测 | 中``。
- 模板化技术路径：模型超参（``alpha`` 正则强度、``target`` 标签列）外置为
  JSON/YAML 超参包，切换模板只改超参包（PRD §5.3「换行业只改 Recipe」）。
- 验收（PRD §10 DoD）：质量预测在客户数据上达约定 R² / MAE 指标。

本模块交付
----------
1. **岭回归 ``RidgeRegression``**：纯标准库最小二乘 + L2 正则（闭式解），
   ``fit(X, y)`` / ``predict(X)`` / ``coef_`` / ``intercept_``。不依赖 numpy。
2. **质量预测模型 ``QualityModel``**：聚合特征名 + 目标列 + 岭回归，提供
   ``fit(matrix, target)`` / ``predict(matrix)`` / ``evaluate(matrix, target)``
   （R² / MAE / RMSE）/ ``explain()``（权重 → 特征贡献，可溯源）。
3. **超参包 ``ModelRecipe``**：``target``/``alpha``/``feature_names`` 外置
   JSON/YAML 加载（零依赖 YAML 子集解析）。
4. **评估指标**：R²、MAE、RMSE 纯标准库实现。

设计要点
--------
- **零运行时依赖**（纯标准库）：矩阵运算手写（不依赖 numpy/sklearn）。
- **可解释**：``explain()`` 输出每个特征的权重 × 方差贡献度，供 #73 接入驾驶舱
  展示"为何预测这个纯度"（对齐 PRD"要求结果可解释、可溯源"）。
- **稳健**：L2 正则避免共线性/过拟合；数据不足时返回明确错误而非崩溃。
"""
from __future__ import annotations

import json
import math
import os
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence, Tuple

NAN = float("nan")


class ModelError(ValueError):
    """质量预测模型错误（数据不足 / 维度不匹配 / 奇异 等）。"""


# ---------------------------------------------------------------------------
# 纯标准库线性代数（最小二乘岭回归闭式解）
# ---------------------------------------------------------------------------


def _matmul_at_a(a: Sequence[Sequence[float]]) -> List[List[float]]:
    """计算 AᵀA（n×n）。"""
    n = len(a[0]) if a else 0
    res = [[0.0] * n for _ in range(n)]
    for row in a:
        for i in range(n):
            ri = row[i]
            if ri == 0.0:
                continue
            for j in range(n):
                res[i][j] += ri * row[j]
    return res


def _matvec_at_b(a: Sequence[Sequence[float]], b: Sequence[float]) -> List[float]:
    """计算 Aᵀb（n）。"""
    n = len(a[0]) if a else 0
    res = [0.0] * n
    for row, y in zip(a, b):
        for i in range(n):
            res[i] += row[i] * y
    return res


def _solve(A: List[List[float]], b: List[float]) -> List[float]:
    """高斯消元解 Ax=b（带部分主元）。A 会被修改。"""
    n = len(A)
    # 增广
    M = [list(A[i]) + [b[i]] for i in range(n)]
    for col in range(n):
        # 主元
        pivot = max(range(col, n), key=lambda r: abs(M[r][col]))
        if abs(M[pivot][col]) < 1e-12:
            raise ModelError("矩阵奇异（特征共线或数据不足），无法求解")
        M[col], M[pivot] = M[pivot], M[col]
        piv = M[col][col]
        for j in range(col, n + 1):
            M[col][j] /= piv
        for r in range(n):
            if r == col:
                continue
            factor = M[r][col]
            if factor == 0.0:
                continue
            for j in range(col, n + 1):
                M[r][j] -= factor * M[col][j]
    return [M[i][n] for i in range(n)]


# ---------------------------------------------------------------------------
# 岭回归
# ---------------------------------------------------------------------------


class RidgeRegression:
    """岭回归（L2 正则最小二乘，闭式解）。纯标准库。

    解：``w = (XᵀX + αI)⁻¹ Xᵀy``。
    """

    def __init__(self, alpha: float = 1.0):
        if alpha < 0:
            raise ModelError(f"alpha 不能为负: {alpha}")
        self.alpha = alpha
        self.coef_: List[float] = []
        self.intercept_: float = 0.0
        self._n_features: int = 0

    def fit(self, X: Sequence[Sequence[float]], y: Sequence[float]) -> "RidgeRegression":
        m = len(X)
        if m == 0:
            raise ModelError("训练集为空")
        n = len(X[0])
        if n == 0:
            raise ModelError("特征数为 0")
        if len(y) != m:
            raise ModelError(f"X/y 行数不匹配: {m} != {len(y)}")
        self._n_features = n
        # 中心化（数值稳定 + 让 intercept 可独立）
        x_mean = [sum(X[i][j] for i in range(m)) / m for j in range(n)]
        y_mean = sum(y) / m
        Xc = [[X[i][j] - x_mean[j] for j in range(n)] for i in range(m)]
        yc = [y[i] - y_mean for i in range(m)]
        # 正规方程 (XᵀX + αI) w = Xᵀy
        A = _matmul_at_a(Xc)
        for i in range(n):
            A[i][i] += self.alpha
        b = _matvec_at_b(Xc, yc)
        self.coef_ = _solve(A, b)
        self.intercept_ = y_mean - sum(self.coef_[j] * x_mean[j] for j in range(n))
        return self

    def predict(self, X: Sequence[Sequence[float]]) -> List[float]:
        if not self.coef_:
            raise ModelError("模型未训练")
        return [self.intercept_ + sum(self.coef_[j] * row[j]
                                      for j in range(self._n_features))
                for row in X]

    def to_dict(self) -> Dict[str, Any]:
        return {"alpha": self.alpha, "coef": list(self.coef_),
                "intercept": self.intercept_,
                "n_features": self._n_features}

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "RidgeRegression":
        m = cls(alpha=float(d.get("alpha", 1.0)))
        m.coef_ = [float(c) for c in d.get("coef", [])]
        m.intercept_ = float(d.get("intercept", 0.0))
        m._n_features = int(d.get("n_features", len(m.coef_)))
        return m


# ---------------------------------------------------------------------------
# 评估指标
# ---------------------------------------------------------------------------


def r2_score(y_true: Sequence[float], y_pred: Sequence[float]) -> float:
    """决定系数 R²。"""
    if len(y_true) != len(y_pred):
        raise ModelError("y_true/y_pred 长度不匹配")
    n = len(y_true)
    if n == 0:
        return NAN
    mean = sum(y_true) / n
    ss_res = sum((y_true[i] - y_pred[i]) ** 2 for i in range(n))
    ss_tot = sum((y_true[i] - mean) ** 2 for i in range(n))
    if ss_tot == 0:
        return 1.0 if ss_res == 0 else 0.0
    return 1.0 - ss_res / ss_tot


def mae_score(y_true: Sequence[float], y_pred: Sequence[float]) -> float:
    if len(y_true) != len(y_pred):
        raise ModelError("y_true/y_pred 长度不匹配")
    n = len(y_true)
    return sum(abs(y_true[i] - y_pred[i]) for i in range(n)) / n if n else NAN


def rmse_score(y_true: Sequence[float], y_pred: Sequence[float]) -> float:
    if len(y_true) != len(y_pred):
        raise ModelError("y_true/y_pred 长度不匹配")
    n = len(y_true)
    if n == 0:
        return NAN
    return math.sqrt(sum((y_true[i] - y_pred[i]) ** 2 for i in range(n)) / n)


# ---------------------------------------------------------------------------
# 质量预测模型（聚合特征 + 目标 + 岭回归）
# ---------------------------------------------------------------------------


@dataclass
class ModelRecipe:
    """模型超参包（外置 JSON/YAML）。"""

    target: str               # 质量标签列名（如 TiCl₄纯度）
    alpha: float = 1.0        # L2 正则强度
    feature_names: List[str] = field(default_factory=list)
    unit: str = ""

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "ModelRecipe":
        return cls(
            target=str(d.get("target", "")).strip(),
            alpha=float(d.get("alpha", 1.0)),
            feature_names=[str(x) for x in d.get("feature_names", [])],
            unit=str(d.get("unit", "")).strip(),
        )

    def validate(self) -> List[str]:
        errs = []
        if not self.target:
            errs.append("target 不能为空")
        if self.alpha < 0:
            errs.append(f"alpha 不能为负: {self.alpha}")
        return errs


@dataclass
class Evaluation:
    """评估结果。"""

    r2: float
    mae: float
    rmse: float
    n_samples: int

    def to_dict(self) -> Dict[str, Any]:
        return {"r2": round(self.r2, 6), "mae": round(self.mae, 6),
                "rmse": round(self.rmse, 6), "n_samples": self.n_samples}

    def passes(self, *, min_r2: float = 0.0, max_mae: float = math.inf) -> bool:
        return (self.r2 >= min_r2 and self.mae <= max_mae
                and not math.isnan(self.r2))


class QualityModel:
    """质量预测模型：特征名 + 目标列 + 岭回归。"""

    def __init__(self, recipe: ModelRecipe):
        errs = recipe.validate()
        if errs:
            raise ModelError("超参包校验失败: " + "; ".join(errs))
        self.recipe = recipe
        self.regression: Optional[RidgeRegression] = None
        self._feature_std: List[float] = []

    @property
    def fitted(self) -> bool:
        return self.regression is not None

    def fit(self, X: Sequence[Sequence[float]], y: Sequence[float]) -> "QualityModel":
        """X 行=样本，列=特征（顺序与 recipe.feature_names 对齐）。"""
        self.regression = RidgeRegression(alpha=self.recipe.alpha).fit(X, y)
        # 记录训练集每列标准差，供 explain() 计算尺度归一化重要性
        m = len(X)
        n = len(X[0]) if X else 0
        if m > 1 and n:
            means = [sum(X[i][j] for i in range(m)) / m for j in range(n)]
            self._feature_std = [
                math.sqrt(sum((X[i][j] - means[j]) ** 2 for i in range(m)) / (m - 1))
                for j in range(n)]
        else:
            self._feature_std = [1.0] * n
        return self

    def predict(self, X: Sequence[Sequence[float]]) -> List[float]:
        if self.regression is None:
            raise ModelError("模型未训练，先 fit()")
        return self.regression.predict(X)

    def evaluate(self, X: Sequence[Sequence[float]],
                 y: Sequence[float]) -> Evaluation:
        pred = self.predict(X)
        return Evaluation(
            r2=r2_score(y, pred), mae=mae_score(y, pred),
            rmse=rmse_score(y, pred), n_samples=len(y))

    def explain(self) -> List[Dict[str, Any]]:
        """特征贡献度（尺度归一化：|权重| × 特征标准差），供可解释性。

        归一化动机：原始权重受特征量纲影响（温度 850℃ vs 配比 30），
        ``|权重|×std`` 才反映特征对预测的实际扰动幅度（与 sklearn
        permutation importance / 标准化系数同思路），可跨特征横向比较。
        """
        if self.regression is None:
            raise ModelError("模型未训练")
        names = self.recipe.feature_names or [
            f"x{i}" for i in range(len(self.regression.coef_))]
        weights = list(self.regression.coef_)
        stds = getattr(self, "_feature_std", None) or [1.0] * len(weights)
        contribs = [abs(weights[i]) * stds[i] for i in range(len(weights))]
        total = sum(contribs) or 1.0
        return [{"feature": names[i], "weight": round(weights[i], 6),
                 "importance": round(contribs[i] / total, 4)}
                for i in range(len(weights))]

    def to_dict(self) -> Dict[str, Any]:
        return {
            "recipe": {"target": self.recipe.target, "alpha": self.recipe.alpha,
                       "feature_names": list(self.recipe.feature_names),
                       "unit": self.recipe.unit},
            "regression": (self.regression.to_dict()
                           if self.regression else None),
            "feature_std": list(self._feature_std),
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "QualityModel":
        model = cls(ModelRecipe.from_dict(d.get("recipe", {})))
        reg = d.get("regression")
        if reg:
            model.regression = RidgeRegression.from_dict(reg)
        model._feature_std = [float(x) for x in d.get("feature_std", [])]
        return model


# ---------------------------------------------------------------------------
# 训练集数据集（特征矩阵 + 标签）
# ---------------------------------------------------------------------------


@dataclass
class TrainingSet:
    """训练集：特征名 + X + y。"""

    feature_names: List[str]
    X: List[List[float]]
    y: List[float]

    @classmethod
    def from_records(cls, records: Sequence[Dict[str, float]],
                     feature_names: Sequence[str],
                     target: str) -> "TrainingSet":
        X, y = [], []
        for r in records:
            if target not in r or math.isnan(r[target]):
                continue
            row = [r.get(fn, NAN) for fn in feature_names]
            if any(math.isnan(v) for v in row):
                continue
            X.append(row)
            y.append(r[target])
        return cls(feature_names=list(feature_names), X=X, y=y)

    def __len__(self) -> int:
        return len(self.X)


# ---------------------------------------------------------------------------
# 超参包加载（零依赖 YAML 子集 / JSON）
# ---------------------------------------------------------------------------


def load_recipe(text: str) -> ModelRecipe:
    text = text.strip()
    if text.startswith("{"):
        data = json.loads(text)
    else:
        data = _parse_yaml_subset(text)
    if not isinstance(data, dict):
        raise ModelError("超参包顶层应为映射")
    return ModelRecipe.from_dict(data)


def _parse_yaml_subset(text: str) -> Any:
    """极简 YAML 子集解析（与 features.py 同款实现，避免跨模块依赖）。"""
    lines: List[str] = []
    for raw in text.splitlines():
        stripped = raw.rstrip()
        if not stripped.strip() or stripped.lstrip().startswith("#"):
            continue
        hi = _find_inline_comment(stripped)
        if hi is not None:
            stripped = stripped[:hi].rstrip()
        if stripped:
            lines.append(stripped)
    parser = _YamlParser(lines)
    return parser.parse_block(0)[0] if lines else {}


def _find_inline_comment(line: str) -> Optional[int]:
    depth = 0
    in_str = False
    for i, ch in enumerate(line):
        if ch == '"':
            in_str = not in_str
        elif not in_str:
            if ch in "[{":
                depth += 1
            elif ch in "]}":
                depth = max(0, depth - 1)
            elif ch == "#" and depth == 0:
                if i == 0 or line[i - 1] in (" ", "\t"):
                    return i
    return None


def _parse_scalar(raw: str) -> Any:
    raw = raw.strip()
    if not raw:
        return None
    if raw.startswith('"') and raw.endswith('"'):
        return raw[1:-1]
    if raw.startswith("[") or raw.startswith("{"):
        try:
            return json.loads(raw)
        except Exception:
            return raw
    low = raw.lower()
    if low == "true":
        return True
    if low == "false":
        return False
    if low in ("null", "~", "none"):
        return None
    try:
        return int(raw)
    except ValueError:
        pass
    try:
        return float(raw)
    except ValueError:
        pass
    return raw


class _YamlParser:
    def __init__(self, lines: List[str]) -> None:
        self.lines = lines
        self.i = 0

    def _indent(self, line: str) -> int:
        return len(line) - len(line.lstrip(" "))

    def parse_block(self, indent: int) -> Tuple[Any, bool]:
        if self.i >= len(self.lines):
            return {}, False
        line = self.lines[self.i]
        cur = self._indent(line)
        if cur < indent:
            return {}, False
        stripped = line.strip()
        if stripped.startswith("- ") or stripped == "-":
            return self._parse_list(cur), True
        return self._parse_mapping(cur), False

    def _parse_mapping(self, indent: int) -> Dict[str, Any]:
        result: Dict[str, Any] = {}
        effective = indent
        if self.i < len(self.lines):
            first = self._indent(self.lines[self.i])
            if first > indent:
                effective = first
        while self.i < len(self.lines):
            line = self.lines[self.i]
            cur = self._indent(line)
            if cur < effective:
                break
            if cur > effective:
                self.i += 1
                continue
            stripped = line.strip()
            if stripped.startswith("- "):
                break
            key, sep, rest = stripped.partition(":")
            if not sep:
                self.i += 1
                continue
            key = key.strip()
            rest = rest.strip()
            self.i += 1
            if rest:
                result[key] = _parse_scalar(rest)
            else:
                if self.i < len(self.lines):
                    nxt = self._indent(self.lines[self.i])
                    if nxt > effective:
                        val, _ = self.parse_block(nxt)
                        result[key] = val
        return result

    def _parse_list(self, indent: int) -> List[Any]:
        items: List[Any] = []
        while self.i < len(self.lines):
            line = self.lines[self.i]
            cur = self._indent(line)
            if cur < indent:
                break
            if cur > indent:
                self.i += 1
                continue
            stripped = line.strip()
            if not stripped.startswith("-"):
                break
            item_text = stripped[1:].strip()
            if not item_text:
                self.i += 1
                items.append(None)
                continue
            if ":" in item_text and not item_text.startswith('"'):
                k, sep, v = item_text.partition(":")
                if sep:
                    item: Dict[str, Any] = {k.strip(): _parse_scalar(v.strip())}
                    self.i += 1
                    if self.i < len(self.lines):
                        child_indent = self._indent(self.lines[self.i])
                        if child_indent > cur:
                            sub, _ = self.parse_block(child_indent)
                            if isinstance(sub, dict):
                                item.update(sub)
                    items.append(item)
                    continue
            items.append(_parse_scalar(item_text))
            self.i += 1
        return items