# -*- coding: utf-8 -*-
"""Ti-2 跨工序关联寻优 · 模型训练（Issue #80 / PRD §5.3 ④）。

承接 #78/#79 的配方优化建模与求解器，本模块把"上游工序指标 → 下游工序质量"
的**跨工序关联**建成**可训练、可评估、可序列化**的轻量模型，为「数据就绪后」
的跨工序寻优（PRD 架构表：``入：上游(TiCl₄)指标；出：下游(海绵钛)寻优建议``）
提供量化基础。

PRD 设计口径
------------
- 架构表（PRD §5.3 ④）：``跨工序关联寻优 | 关联建模 | 入：上游(TiCl₄)指标；
  出：下游(海绵钛)寻优建议 | 高（需闭环反馈）``。
- 模板化技术路径：默认「固定主干 + 可配置超参」；**新增结构走插件注册而非改
  内核**。故本模块主干为**线性 / 岭回归（纯标准库）**，跨工序的强非线性关联
  （LSTM/GNN）走 recipe 插件，不在本期内核。
- 里程碑表：二期交付（标注数据 ≥ 6 个月，LIMS 对接后补标）。故本期交付**可跑通、
  可测试**的关联模型与训练/评估闭环，数据就绪后即可上线。

本模块交付
----------
1. **``CrossProcessSample``**：跨工序样本（上游特征 Dict + 下游目标 + 批次/时间），
   鸭子类型，便于独立测试。
2. **``RidgeRegression``**：纯标准库岭回归（含截距、L2 正则、闭式解），
   ``fit`` / ``predict`` / 评估（MSE、R²）。
3. **``CrossProcessModel``**：跨工序关联模型——把上游指标映射到下游质量，聚合
   多个下游目标的回归器；``fit`` / ``predict`` / ``evaluate``（R² 报告）/ 序列化
   （零依赖 JSON，便于版本化保存与 #41 模型模板注册机制对接）。
4. **``CrossProcessModelConfig``**：声明式配置（特征清单、目标清单、正则强度、
   训练最小样本数），对齐 PRD「超参包驱动」。

设计要点
--------
- **零第三方依赖**（纯标准库）：与内核既有模块一致，便于离线/隔离网部署。
- **训练/评估分离**：``evaluate`` 在测试段产出每个目标的 R²（拟合优度），对齐
  PRD「关键质量指标预测准确率 ≥ 90%」的评估口径。
- **数据门槛前置校验**：``min_samples`` 不足时拒绝训练（对齐 PRD「监督模型需
  ≥ 6 个月标注」的数据门槛约束，避免低质上线）。
- **与 #78/#79 解耦**：模型只依赖样本的 Dict 特征，输出下游质量预测；上层
  （#81）可把预测喂回 #78 的目标函数做跨工序寻优。
"""
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 CrossProcessError(ValueError):
    """跨工序关联模型错误（特征缺失/样本不足/未训练等）。"""


def _is_num(x: object) -> bool:
    return isinstance(x, (int, float)) and not (isinstance(x, float) and math.isnan(x))


# ---------------------------------------------------------------------------
# 样本
# ---------------------------------------------------------------------------


@dataclass
class CrossProcessSample:
    """跨工序样本：上游特征 + 下游目标（鸭子类型，便于独立测试）。

    ``upstream`` 的 key 为上游测点/特征名（如 ``TiCl4_purity``、``TiCl4_impurity``），
    ``downstream`` 的 key 为下游质量指标（如 ``sponge_titanium_grade``）。
    """

    upstream: Dict[str, float] = field(default_factory=dict)
    downstream: Dict[str, float] = field(default_factory=dict)
    batch: str = ""           # 批次号（可溯源，供 #81 引用）
    timestamp: float = 0.0


# ---------------------------------------------------------------------------
# 岭回归（纯标准库闭式解）
# ---------------------------------------------------------------------------


class RidgeRegression:
    """单目标岭回归（含截距 + L2 正则），闭式解（零第三方依赖）。

    解析解：``w = (XᵀX + λI)⁻¹ Xᵀ y``（``X`` 已含截距列）；用高斯消元解线性
    方程组避免 numpy 依赖。``λ=0`` 即普通最小二乘。
    """

    def __init__(self, alpha: float = 1.0) -> None:
        if alpha < 0:
            raise CrossProcessError("岭回归正则强度 alpha 必须 ≥ 0")
        self.alpha = alpha
        self._feature_names: List[str] = []
        self._weights: List[float] = []  # 末位为截距
        self._fitted = False

    @property
    def fitted(self) -> bool:
        return self._fitted

    @property
    def feature_names(self) -> List[str]:
        return list(self._feature_names)

    def fit(self, X: Sequence[Sequence[float]], y: Sequence[float],
            feature_names: Sequence[str]) -> "RidgeRegression":
        """在训练段拟合（X 不含截距列，内部补）。"""
        n = len(X)
        if n == 0:
            raise CrossProcessError("RidgeRegression.fit 至少需要 1 条样本")
        if n != len(y):
            raise CrossProcessError("X 与 y 样本数不一致")
        p = len(feature_names)
        if any(len(row) != p for row in X):
            raise CrossProcessError("X 列数与 feature_names 不一致")
        # 设计矩阵加截距列（末列恒为 1）
        Xa = [list(row) + [1.0] for row in X]
        # XtX + λI（不对截距正则：最后一行/列不加 λ）
        dim = p + 1
        A = [[0.0] * dim for _ in range(dim)]
        for row in Xa:
            for i in range(dim):
                for j in range(dim):
                    A[i][j] += row[i] * row[j]
        for i in range(p):  # 不对截距正则
            A[i][i] += self.alpha
        # Xty
        b = [0.0] * dim
        for k, row in enumerate(Xa):
            for i in range(dim):
                b[i] += row[i] * y[k]
        # 解 A w = b（高斯消元 + 回代，带部分主元）
        self._weights = _solve_linear(A, b)
        self._feature_names = list(feature_names)
        self._fitted = True
        return self

    def predict_one(self, x: Sequence[float]) -> float:
        if not self._fitted:
            raise CrossProcessError("RidgeRegression 未 fit")
        if len(x) != len(self._feature_names):
            raise CrossProcessError("预测特征数与训练不一致")
        return sum(w * v for w, v in zip(self._weights[:-1], x)) + self._weights[-1]

    def weights_dict(self) -> Dict[str, float]:
        """特征权重 + 截距（可解释，供 #81 引用）。"""
        if not self._fitted:
            return {}
        d = {name: self._weights[i] for i, name in enumerate(self._feature_names)}
        d["__intercept__"] = self._weights[-1]
        return d

    def to_dict(self) -> Dict[str, Any]:
        return {
            "alpha": self.alpha,
            "feature_names": list(self._feature_names),
            "weights": list(self._weights),
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "RidgeRegression":
        m = cls(alpha=float(d.get("alpha", 1.0)))
        m._feature_names = list(d.get("feature_names", []))
        m._weights = [float(w) for w in d.get("weights", [])]
        m._fitted = bool(m._weights)
        return m


# ---------------------------------------------------------------------------
# 线性方程组求解（高斯消元，部分主元）
# ---------------------------------------------------------------------------


def _solve_linear(A: List[List[float]], b: List[float]) -> List[float]:
    """解 A w = b（方阵），高斯消元 + 回代，带部分主元选元。"""
    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 CrossProcessError("正规方程奇异（特征共线性或样本不足）")
        M[col], M[pivot] = M[pivot], M[col]
        # 消元
        piv = M[col][col]
        for r in range(col + 1, n):
            factor = M[r][col] / piv
            if factor != 0.0:
                for c in range(col, n + 1):
                    M[r][c] -= factor * M[col][c]
    # 回代
    w = [0.0] * n
    for i in range(n - 1, -1, -1):
        s = M[i][n] - sum(M[i][j] * w[j] for j in range(i + 1, n))
        w[i] = s / M[i][i]
    return w


# ---------------------------------------------------------------------------
# 跨工序关联模型
# ---------------------------------------------------------------------------


@dataclass
class CrossProcessModelConfig:
    """跨工序关联模型声明式配置（对齐 PRD 超参包驱动）。"""

    upstream_features: List[str] = field(default_factory=list)
    downstream_targets: List[str] = field(default_factory=list)
    alpha: float = 1.0            # 岭回归正则强度
    min_samples: int = 10         # 训练最小样本数（数据门槛前置校验）

    def __post_init__(self) -> None:
        if not self.upstream_features:
            raise CrossProcessError("upstream_features 不能为空")
        if not self.downstream_targets:
            raise CrossProcessError("downstream_targets 不能为空")
        if self.alpha < 0:
            raise CrossProcessError("alpha 必须 ≥ 0")
        if self.min_samples < 2:
            raise CrossProcessError("min_samples 必须 ≥ 2")


class CrossProcessModel:
    """跨工序关联模型：上游指标 → 下游质量（多目标，每目标一个岭回归）。

    训练阶段对每个下游目标拟合一个岭回归；推理阶段给定上游指标预测全部下游目标；
    评估阶段在测试段产出每个目标的 R²（拟合优度）。
    """

    def __init__(self, config: CrossProcessModelConfig) -> None:
        self.config = config
        self._regressors: Dict[str, RidgeRegression] = {}
        self._fitted = False
        # 训练统计（可解释，供 #81 引用）
        self.train_n: int = 0
        self.train_means: Dict[str, float] = {}

    @property
    def fitted(self) -> bool:
        return self._fitted

    # ---- 训练 --------------------------------------------------------

    def fit(self, samples: Sequence[CrossProcessSample]) -> "CrossProcessModel":
        """在训练段对每个下游目标拟合岭回归。"""
        cfg = self.config
        if len(samples) < cfg.min_samples:
            raise CrossProcessError(
                f"训练样本不足：{len(samples)} < min_samples={cfg.min_samples}"
                "（对齐 PRD 监督模型数据门槛）")
        # 构造 X / 每目标 y
        X: List[List[float]] = []
        per_target_y: Dict[str, List[float]] = {t: [] for t in cfg.downstream_targets}
        for s in samples:
            row = []
            ok = True
            for f in cfg.upstream_features:
                v = s.upstream.get(f)
                if not _is_num(v):
                    ok = False
                    break
                row.append(float(v))
            if not ok:
                continue
            # 每个目标都要有值，否则跳过该样本（保持对齐）
            target_vals = {}
            for t in cfg.downstream_targets:
                tv = s.downstream.get(t)
                if not _is_num(tv):
                    ok = False
                    break
                target_vals[t] = float(tv)
            if not ok:
                continue
            X.append(row)
            for t in cfg.downstream_targets:
                per_target_y[t].append(target_vals[t])
        if len(X) < cfg.min_samples:
            raise CrossProcessError(
                f"有效训练样本不足：{len(X)} < {cfg.min_samples}（含缺失值过滤后）")
        self._regressors = {}
        for t in cfg.downstream_targets:
            reg = RidgeRegression(alpha=cfg.alpha)
            reg.fit(X, per_target_y[t], cfg.upstream_features)
            self._regressors[t] = reg
        self.train_n = len(X)
        # 训练段上游特征均值（漂移检测/可解释输入）
        self.train_means = {
            f: sum(row[i] for row in X) / len(X)
            for i, f in enumerate(cfg.upstream_features)
        }
        self._fitted = True
        return self

    # ---- 推理 --------------------------------------------------------

    def predict(self, upstream: Dict[str, float]) -> Dict[str, float]:
        """给定上游指标，预测全部下游目标。"""
        if not self._fitted:
            raise CrossProcessError("CrossProcessModel 未 fit")
        row = []
        for f in self.config.upstream_features:
            v = upstream.get(f)
            if not _is_num(v):
                raise CrossProcessError(f"预测缺少上游特征 {f!r}")
            row.append(float(v))
        return {t: reg.predict_one(row) for t, reg in self._regressors.items()}

    # ---- 评估 --------------------------------------------------------

    def evaluate(self, samples: Sequence[CrossProcessSample]) -> Dict[str, float]:
        """在测试段产出每个目标的 R²（拟合优度）+ 总体 MSE。"""
        if not self._fitted:
            raise CrossProcessError("CrossProcessModel 未 fit，无法 evaluate")
        report: Dict[str, float] = {}
        # 收集每个目标的 真实/预测
        per_target: Dict[str, Tuple[List[float], List[float]]] = {
            t: ([], []) for t in self.config.downstream_targets
        }
        for s in samples:
            try:
                pred = self.predict(s.upstream)
            except CrossProcessError:
                continue
            for t in self.config.downstream_targets:
                truth = s.downstream.get(t)
                if not _is_num(truth):
                    continue
                per_target[t][0].append(float(truth))
                per_target[t][1].append(pred[t])
        all_sq_err = []
        for t, (truths, preds) in per_target.items():
            if not truths:
                report[f"r2_{t}"] = NAN
                continue
            mean_t = sum(truths) / len(truths)
            ss_res = sum((tr - pr) ** 2 for tr, pr in zip(truths, preds))
            ss_tot = sum((tr - mean_t) ** 2 for tr in truths)
            r2 = 1.0 - ss_res / ss_tot if ss_tot > 1e-12 else (1.0 if ss_res < 1e-12 else NAN)
            report[f"r2_{t}"] = r2
            all_sq_err.extend((tr - pr) ** 2 for tr, pr in zip(truths, preds))
        report["mse_overall"] = (sum(all_sq_err) / len(all_sq_err)) if all_sq_err else NAN
        report["n_eval"] = float(sum(len(v[0]) for v in per_target.values()))
        return report

    # ---- 可解释（供 #81） -------------------------------------------

    def feature_weights(self, target: str) -> Dict[str, float]:
        """某下游目标的上游特征权重 + 截距（可解释：上游对下游的影响）。"""
        reg = self._regressors.get(target)
        if reg is None:
            raise CrossProcessError(f"未知下游目标 {target!r}")
        return reg.weights_dict()

    # ---- 序列化 ------------------------------------------------------

    def to_dict(self) -> Dict[str, Any]:
        return {
            "config": {
                "upstream_features": list(self.config.upstream_features),
                "downstream_targets": list(self.config.downstream_targets),
                "alpha": self.config.alpha,
                "min_samples": self.config.min_samples,
            },
            "train_n": self.train_n,
            "train_means": dict(self.train_means),
            "regressors": {t: r.to_dict() for t, r in self._regressors.items()},
        }

    @classmethod
    def from_dict(cls, d: Dict[str, Any]) -> "CrossProcessModel":
        cfg = CrossProcessModelConfig(
            upstream_features=list(d["config"]["upstream_features"]),
            downstream_targets=list(d["config"]["downstream_targets"]),
            alpha=float(d["config"].get("alpha", 1.0)),
            min_samples=int(d["config"].get("min_samples", 10)),
        )
        m = cls(cfg)
        m.train_n = int(d.get("train_n", 0))
        m.train_means = dict(d.get("train_means", {}))
        m._regressors = {t: RidgeRegression.from_dict(rd)
                         for t, rd in d.get("regressors", {}).items()}
        m._fitted = bool(m._regressors)
        return m

    def save(self, path: str) -> None:
        with open(path, "w", encoding="utf-8") as fh:
            json.dump(self.to_dict(), fh, ensure_ascii=False, indent=2)

    @classmethod
    def load(cls, path: str) -> "CrossProcessModel":
        with open(path, "r", encoding="utf-8") as fh:
            return cls.from_dict(json.load(fh))