"""温度控制系统模型预测控制器。

本模块实现论文中描述的 MPC 算法，包含以下组件：
  - 参考轨迹生成
  - 状态估计器
  - 基于 LSTM 的预测模型封装
  - 滚动优化器（最小化二次代价函数 J）
  - 反馈校正

使用方法：
    controller = MPCController.from_checkpoint('model.pt', mpc_params)
    u = controller.step(history_records, yr=target_temp)
"""
from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Callable

import numpy as np

# ── 可选依赖：PyTorch ────────────────────────────────────────────────────────
_TORCH_AVAILABLE = False
try:
    import torch
    import torch.nn as nn

    _TORCH_AVAILABLE = True

    class _LSTMModel(nn.Module):
        def __init__(self, input_size: int, hidden_size: int, num_layers: int) -> None:
            super().__init__()
            dropout = 0.2 if num_layers > 1 else 0.0
            self.lstm = nn.LSTM(
                input_size, hidden_size, num_layers,
                batch_first=True, dropout=dropout,
            )
            self.fc = nn.Linear(hidden_size, 1)

        def forward(self, x: 'torch.Tensor') -> 'torch.Tensor':
            out, _ = self.lstm(x)
            return self.fc(out[:, -1, :]).squeeze(-1)

except ImportError:
    pass


# ── 特征列定义（须与训练器保持一致）────────────────────────────────────────────
FEATURE_COLS = ['current', 'voltage', 'set_temperature', 'actual_temperature']
TARGET_COL = 'actual_temperature'
TARGET_IDX = FEATURE_COLS.index(TARGET_COL)
CURRENT_IDX = FEATURE_COLS.index('current')


# ─────────────────────────────────────────────────────────────────────────────
# MPC 参数数据类
# ─────────────────────────────────────────────────────────────────────────────

@dataclass
class MPCParams:
    """MPC 控制器的全部可调参数。

    属性说明：
        P          : 预测时域长度（步数）。应覆盖温度系统的主要动态响应时间。
        M          : 控制时域长度（步数），M <= P。M 越小越稳定、计算量越小。
        Q          : 跟踪误差权重（标量或长度为 P 的列表）。
                     Q 越大跟踪越快，但可能超调。
        R          : 控制增量权重（标量或长度为 M 的列表）。
                     R 越大电流变化越平滑，响应越慢。
        alpha      : 参考轨迹柔化系数（0 < α < 1）。
                     越接近1，趋近设定值越缓慢平滑。
        u_min      : 控制量下限（电流，A），须 ≥ 0。
        u_max      : 控制量上限（电流，A）。
        du_min     : 每步控制增量下限（A/步）。
        du_max     : 每步控制增量上限（A/步）。
        y_min      : 预测温度下限（°C），安全下界。
        y_max      : 预测温度上限（°C），安全上界。
        correction_gain : 反馈校正增益（0~1）。控制预测误差对后续预测的修正强度。
        optimizer  : 'scipy'（默认，无梯度 SLSQP）或 'gradient'
                     （PyTorch 自动微分，torch 可用时更快）。
    """
    P: int = 20
    M: int = 5
    Q: float = 10.0
    R: float = 0.1
    alpha: float = 0.8
    u_min: float = 0.0
    u_max: float = 10.0
    du_min: float = -2.0
    du_max: float = 2.0
    y_min: float = -50.0
    y_max: float = 200.0
    correction_gain: float = 0.5
    optimizer: str = 'scipy'


# ─────────────────────────────────────────────────────────────────────────────
# LSTMPredictor  –  封装已保存的 .pt 检查点
# ─────────────────────────────────────────────────────────────────────────────

class LSTMPredictor:
    """加载已训练的 LSTM 检查点，并提供 predict() 预测接口。"""

    def __init__(self, checkpoint: dict) -> None:
        if not _TORCH_AVAILABLE:
            raise RuntimeError('PyTorch 未安装，无法加载 LSTM 模型。')

        self.seq_len: int = checkpoint['seq_len']
        self.feature_cols: list[str] = checkpoint['feature_cols']
        self.data_min = np.array(checkpoint['data_min'], dtype=np.float32)
        self.data_max = np.array(checkpoint['data_max'], dtype=np.float32)
        self.data_range = self.data_max - self.data_min
        self.data_range[self.data_range == 0] = 1.0

        self._model = _LSTMModel(
            checkpoint['input_size'],
            checkpoint['hidden_size'],
            checkpoint['num_layers'],
        )
        self._model.load_state_dict(checkpoint['model_state'])
        self._model.eval()

    @classmethod
    def from_file(cls, path: 'Path | str') -> 'LSTMPredictor':
        ckpt = torch.load(str(path), map_location='cpu', weights_only=False)
        return cls(ckpt)

    # ── 归一化辅助方法 ──────────────────────────────────────────────────

    def _norm(self, arr: np.ndarray) -> np.ndarray:
        return (arr - self.data_min) / self.data_range

    def _denorm_target(self, val: float) -> float:
        idx = TARGET_IDX
        return val * self.data_range[idx] + self.data_min[idx]

    # ── 单步预测 ──────────────────────────────────────────────────────────

    def predict_one(self, seq_norm: np.ndarray) -> float:
        """基于归一化后的 (seq_len, features) 序列预测下一时刻温度。

        返回原始单位（°C）的预测温度。
        """
        x = torch.tensor(seq_norm[np.newaxis], dtype=torch.float32)  # 形状 (1, T, F)
        with torch.no_grad():
            y_norm = self._model(x).item()
        return self._denorm_target(y_norm)

    # ── 多步滚动预测 ──────────────────────────────────────────────────────────

    def rollout(
        self,
        history: np.ndarray,
        u_sequence: np.ndarray,
        correction_offset: float = 0.0,
    ) -> np.ndarray:
        """给定控制序列，滚动仳真未来 P 步温度。

        参数：
            history           : (seq_len, n_features) 近期测量值，原始单位。
            u_sequence        : (P,) 未来电流序列，原始单位。
            correction_offset : 每步预测叠加的加法校正量（反馈校正，原始单位）。

        返回：
            y_pred : (P,) 预测温度，原始单位。
        """
        P = len(u_sequence)
        seq = history.copy().astype(np.float32)  # (seq_len, F)，原始单位
        y_pred = np.empty(P, dtype=np.float32)

        for i in range(P):
            seq_norm = self._norm(seq)
            y_next = self.predict_one(seq_norm) + correction_offset
            y_pred[i] = y_next

            # 滑动窗口向前推进一步
            new_row = seq[-1].copy()
            new_row[CURRENT_IDX] = u_sequence[i]
            new_row[TARGET_IDX] = y_next
            seq = np.vstack([seq[1:], new_row])

        return y_pred


# ─────────────────────────────────────────────────────────────────────────────
# 状态估计器
# ─────────────────────────────────────────────────────────────────────────────

class StateEstimator:
    """以近期测量值的滑动窗口作为“扩展状态”。

    由于加热器热状态、航天器接收热量等内部状态无法直接测量，
    使用最近 seq_len 步的温度与电流采样值作为充分代理状态。
    """

    def __init__(self, seq_len: int, n_features: int) -> None:
        self.seq_len = seq_len
        self.n_features = n_features
        self._buffer: list[np.ndarray] = []

    def update(self, measurement: np.ndarray) -> None:
        """追加一条新测量记录（长度为 n_features 的一维数组）。"""
        self._buffer.append(measurement.astype(np.float32))
        if len(self._buffer) > self.seq_len:
            self._buffer.pop(0)

    def ready(self) -> bool:
        return len(self._buffer) >= self.seq_len

    def get_state(self) -> np.ndarray:
        """返回 (seq_len, n_features) 历史状态数组。"""
        if not self.ready():
            raise RuntimeError(
                f'状态估计器尚未就绪：已有 {len(self._buffer)} 条记录，'
                f'需要至少 {self.seq_len} 条。'
            )
        return np.array(self._buffer[-self.seq_len:], dtype=np.float32)


# ─────────────────────────────────────────────────────────────────────────────
# MPC 控制器
# ─────────────────────────────────────────────────────────────────────────────

class MPCController:
    """温度控制完整 MPC 控制回路。

    实时循环典型用法::

        ctrl = MPCController.from_checkpoint('model.pt', MPCParams())

        for measurement in sensor_stream:
            ctrl.update_measurement(measurement)
            if ctrl.ready():
                u = ctrl.step(yr=target_temperature)
                send_to_power_supply(u)

    批量仳真模式::

        u = ctrl.step_from_records(records, yr=35.0)
    """

    def __init__(self, predictor: LSTMPredictor, params: MPCParams) -> None:
        self.predictor = predictor
        self.p = params

        self._estimator = StateEstimator(predictor.seq_len, len(FEATURE_COLS))
        self._u_prev: float = 0.0             # 上一时刻控制量
        self._correction_offset: float = 0.0  # 反馈校正偏置

    # ── factory ───────────────────────────────────────────────────────────────

    @classmethod
    def from_checkpoint(
        cls,
        path: 'Path | str',
        params: 'MPCParams | None' = None,
    ) -> 'MPCController':
        predictor = LSTMPredictor.from_file(path)
        return cls(predictor, params or MPCParams())

    # ── 测量值输入 ────────────────────────────────────────────────────────────────

    def update_measurement(self, record: 'dict | np.ndarray') -> None:
        """将一条新传感器测量值送入状态估计器。

        参数：
            record : 键与 FEATURE_COLS 匹配的字典，或长度为 len(FEATURE_COLS) 的一维 numpy 数组。
        """
        if isinstance(record, dict):
            row = np.array(
                [float(record.get(c, 0) or 0) for c in FEATURE_COLS],
                dtype=np.float32,
            )
        else:
            row = np.asarray(record, dtype=np.float32)
        self._estimator.update(row)

    def ready(self) -> bool:
        """历史数据积累足够后返回 True，此时可执行 MPC 控制步。"""
        return self._estimator.ready()

    # ── 参考轨迹 ───────────────────────────────────────────────────────────────────

    def _reference_trajectory(self, y_k: float, yr: float) -> np.ndarray:
        """生成柔化参考轨迹：w(k+j) = α^j · y(k) + (1 - α^j) · yr，j = 1..P"""
        j = np.arange(1, self.p.P + 1, dtype=np.float64)
        alpha_j = self.p.alpha ** j
        return alpha_j * y_k + (1.0 - alpha_j) * yr

    # ── 代价函数 ───────────────────────────────────────────────────────────────────

    def _cost(
        self,
        du_sequence: np.ndarray,  # (M,) 待优化的控制增量序列
        history: np.ndarray,      # (seq_len, F) 当前状态
        w: np.ndarray,            # (P,) 参考轨迹
    ) -> float:
        """二次型 MPC 代价函数  J = Σ Q·(y−w)² + Σ R·(Δu)²"""
        u_seq = self._build_u_sequence(du_sequence)
        y_pred = self.predictor.rollout(history, u_seq, self._correction_offset)

        tracking = float(np.sum(self.p.Q * (y_pred - w) ** 2))
        effort   = float(np.sum(self.p.R * du_sequence ** 2))
        return tracking + effort

    def _build_u_sequence(self, du_sequence: np.ndarray) -> np.ndarray:
        """将 M 个控制增量展开为 P 步控制量序列，并执行限幅。"""
        u_seq = np.empty(self.p.P, dtype=np.float64)
        u = self._u_prev
        for i in range(self.p.P):
            du = du_sequence[i] if i < self.p.M else 0.0
            u = float(np.clip(u + du, self.p.u_min, self.p.u_max))
            u_seq[i] = u
        return u_seq

    # ── 滚动优化器 ─────────────────────────────────────────────────────────────────────

    def _optimize(self, history: np.ndarray, w: np.ndarray) -> np.ndarray:
        """求解最优控制增量序列，返回形状 (M,) 的数组。"""
        from scipy.optimize import minimize

        M = self.p.M
        x0 = np.zeros(M)

        bounds = [(self.p.du_min, self.p.du_max)] * M

        # 不等式约束：u_min ≤ u(k+i) ≤ u_max（通过累加 Δu 链构造）
        constraints = []
        for i in range(M):
            def u_lower(x, idx=i):
                return self._u_prev + float(np.sum(x[: idx + 1])) - self.p.u_min

            def u_upper(x, idx=i):
                return self.p.u_max - (self._u_prev + float(np.sum(x[: idx + 1])))

            constraints += [
                {'type': 'ineq', 'fun': u_lower},
                {'type': 'ineq', 'fun': u_upper},
            ]

        result = minimize(
            self._cost,
            x0,
            args=(history, w),
            method='SLSQP',
            bounds=bounds,
            constraints=constraints,
            options={'maxiter': 200, 'ftol': 1e-6},
        )
        return result.x

    # ── 反馈校正 ─────────────────────────────────────────────────────────────────────

    def _update_correction(self, y_actual: float, y_predicted: float) -> None:
        """利用 k 时刻预测误差更新加法校正偏置。"""
        error = y_actual - y_predicted
        self._correction_offset += self.p.correction_gain * error

    # ── 主控制步 ───────────────────────────────────────────────────────────────────────

    def step(self, yr: float, y_actual: 'float | None' = None) -> float:
        """执行一次 MPC 控制步，返回控制量 u(k)。

        参数：
            yr       : 温度设定值（°C）。
            y_actual : 本步实际温度测量值，用于反馈校正。
                       为 None 时使用状态缓冲区最后一个值。

        返回：
            u : 最优电流指令（A），输出至程控电源。
        """
        if not self.ready():
            raise RuntimeError(
                '控制器尚未就绪，请先调用 update_measurement() 积累足够历史数据。'
            )

        history = self._estimator.get_state()   # (seq_len, F)
        y_k = float(history[-1, TARGET_IDX])    # 当前温度

        # ── 反馈校正 ──────────────────────────────────────────────────────────────
        if y_actual is not None:
            y_pred_k = self.predictor.predict_one(self.predictor._norm(history))
            self._update_correction(y_actual, y_pred_k)

        # ── 参考轨迹 ──────────────────────────────────────────────────────────────
        w = self._reference_trajectory(y_k, yr)  # (P,)

        # ── 滚动优化 ──────────────────────────────────────────────────────────────
        du_opt = self._optimize(history, w)

        # 仅执行第一步增量（滚动时域原则）
        u_opt = float(np.clip(self._u_prev + du_opt[0], self.p.u_min, self.p.u_max))
        self._u_prev = u_opt

        return u_opt

    # ── 批量仿真 ─────────────────────────────────────────────────────────────────────

    def simulate(
        self,
        records: list[dict],
        yr: float,
        on_step: 'Callable[[int, float, float, float], None] | None' = None,
    ) -> dict:
        """对一组历史记录执行闭环 MPC 仿真。

        参数：
            records  : 按时序排列的测量字典列表（格式与训练数据相同）。
            yr       : 温度设定值。
            on_step  : 可选回调函数 callback(step_idx, y_actual, u_cmd, w_ref)。

        返回：
            包含键 'u_history'、'y_history'、'w_history'、'error_history' 的字典。
        """
        u_hist, y_hist, w_hist, err_hist = [], [], [], []

        for i, rec in enumerate(records):
            self.update_measurement(rec)

            y_actual = float(rec.get('actual_temperature') or 0)

            if not self.ready():
                u_hist.append(0.0)
                y_hist.append(y_actual)
                w_hist.append(yr)
                err_hist.append(yr - y_actual)
                continue

            u = self.step(yr=yr, y_actual=y_actual)
            w_k = self._reference_trajectory(y_actual, yr)[0]

            u_hist.append(u)
            y_hist.append(y_actual)
            w_hist.append(w_k)
            err_hist.append(yr - y_actual)

            if on_step:
                on_step(i, y_actual, u, w_k)

        return {
            'u_history':    u_hist,
            'y_history':    y_hist,
            'w_history':    w_hist,
            'error_history': err_hist,
        }

    # ── 便捷接口：从原始记录执行单步控制 ────────────────────────────────────────

    def step_from_records(self, records: list[dict], yr: float) -> float:
        """将最近 seq_len 条记录送入状态估计器后执行一次控制步，返回控制量。

        适用于无需在外部维护状态的一次性调用。
        内部状态估计器缓冲区在多次调用之间不会被清空，
        因此也可用于滚动方式调用。
        """
        for rec in records:
            self.update_measurement(rec)
        return self.step(yr=yr)
