# 用學到的模型做模型預測控制（MPC）：控制水槽水位（同模糊控制章），和 PI 控制比較
# 需要先執行 tank_narx.py 產生 tank_narx.json（裡面有 NARX 網路與 ARX 係數）
import json
import math
from pathlib import Path

import torch

torch.set_default_dtype(torch.float64)
HERE = Path(__file__).parent
d = json.loads((HERE / "tank_narx.json").read_text(encoding="utf-8"))
W = [(torch.tensor(w), torch.tensor(b)) for w, b in d["layers"]]
mu, sd, theta = torch.tensor(d["mu"]), torch.tensor(d["sd"]), torch.tensor(d["arx"])

def true_model(y, ym1, u, um1):
    """真實的水槽（取樣 1 s，每秒 RK4 十步）；可一次算一批"""
    h = y.clone()
    f = lambda x: 0.5 * u - 0.05 * torch.sqrt(2 * 9.8 * x.clamp(min=0))
    for _ in range(10):
        k1 = f(h); k2 = f(h + 0.05 * k1); k3 = f(h + 0.05 * k2); k4 = f(h + 0.1 * k3)
        h = (h + 0.1 / 6 * (k1 + 2 * k2 + 2 * k3 + k4)).clamp(0, 2)
    return h

def narx_model(y, ym1, u, um1):
    h = (torch.stack([y, ym1, u, um1], -1) - mu) / sd
    for k, (w, b) in enumerate(W):
        h = h @ w.T + b
        if k < len(W) - 1:
            h = torch.tanh(h)
    return y + h[..., 0]

def arx_model(y, ym1, u, um1):
    return theta[0] * y + theta[1] * ym1 + theta[2] * u + theta[3] * um1 + theta[4]

def mpc(model, y, ym1, u_prev, r, N=15, m=3, lam=0.5, d=0.0):
    """閥開度在預測區間內分兩段：前 m 秒 u1、之後 u2。先在 0～1 間隔 0.05 試 21×21 組，再在最好的附近間隔 0.005 細搜。
    成本 J = Σ (r − ŷ − d)² + λ(u1 − u_前)²，只執行 u1（滾動時域）；d 為偏差修正量（量測值 − 模型輸出）"""
    def best_of(cand):
        u1, u2 = cand[:, 0], cand[:, 1]
        yk, ykm1, up = torch.full_like(u1, y), torch.full_like(u1, ym1), torch.full_like(u1, u_prev)
        J = lam * (u1 - u_prev) ** 2
        for j in range(N):
            u = u1 if j < m else u2
            yn = model(yk, ykm1, u, up).clamp(0, 2)
            J = J + (r - yn - d) ** 2
            ykm1, yk, up = yk, yn, u
        return cand[int(J.argmin())]
    g = torch.arange(21) * 0.05
    c = best_of(torch.cartesian_prod(g, g))
    off = (torch.arange(11) - 5) * 0.005
    return float(best_of(torch.cartesian_prod((c[0] + off).clamp(0, 1), (c[1] + off).clamp(0, 1)))[0])

ref = [0.5] * 20 + [1.5] * 100 + [0.8] * 100 + [1.8] * 100           # 設定值：0.5 → 1.5 → 0.8 → 1.8 m
u_ss = lambda h: 0.1 * math.sqrt(2 * 9.8 * h)                           # 維持水位 h 所需的閥開度

MODELS = {"MPC（真實模型）": true_model, "MPC（NARX）": narx_model, "MPC（ARX）": arx_model,
          "MPC（NARX＋偏差修正）": narx_model, "MPC（ARX＋偏差修正）": arx_model}

def run(controller):
    y, u_prev, us = [0.5, 0.5], u_ss(0.5), []
    I = u_ss(0.5) / 0.08
    ym = [0.5, 0.5]                                                     # 偏差修正：模型和真實系統並行，用同樣的輸入往前跑
    for k, r in enumerate(ref):
        if controller == "PI":                                          # PI：在 h = 1 m 線性化設計，兩極點都在 −0.2
            e = r - y[-1]
            u_raw = 0.578 * e + 0.08 * (I + e)
            u = min(max(u_raw, 0.0), 1.0)
            if u == u_raw:
                I += e                                                  # 閥飽和時停止積分（防積分飽和）
        elif "偏差修正" in controller:
            d = y[-1] - ym[-1]                                          # 把模型的誤差當成一個固定的干擾
            u = mpc(MODELS[controller], ym[-1], ym[-2], u_prev, r, d=d)
        else:
            u = mpc(MODELS[controller], y[-1], y[-2], u_prev, r)
        ym.append(float(MODELS.get(controller, true_model)(torch.tensor(ym[-1]), torch.tensor(ym[-2]), torch.tensor(u), torch.tensor(u_prev))))
        y.append(float(true_model(torch.tensor(y[-1]), None, torch.tensor(u), None)))
        us.append(u); u_prev = u
    return y[1:], us

out = {"ref": ref}
for name in ("PI", "MPC（真實模型）", "MPC（NARX）", "MPC（ARX）", "MPC（NARX＋偏差修正）", "MPC（ARX＋偏差修正）"):
    y, us = run(name)
    iae = sum(abs(a - b) for a, b in zip(y[1:], ref))
    over = max(y[k + 1] - ref[k] for k in list(range(20, 120)) + list(range(220, 320)))
    out[name] = {"y": y, "u": us}
    ss = [y[k] - ref[k - 1] for k in (120, 220, 320)]                   # 各段結束時的誤差（穩態誤差）
    print(f"{name:16s}：IAE {iae:5.2f} m·s，往上調時最大超越 {max(over, 0) * 100:4.1f} cm，"
          f"各段結束時的誤差 {", ".join(f"{e * 100:+.1f}" for e in ss)} cm")
(HERE / "tank_mpc.json").write_text(json.dumps(out, ensure_ascii=False), encoding="utf-8")