# 系統鑑別：只用量測的「閥開度 u → 水位 h」資料，學出水槽（同模糊控制章）的動態模型
# 線性 ARX 與神經網路 NARX 比較；兩者都用 y(k+1) = f(y(k), y(k−1), u(k), u(k−1))
import json
import math
from pathlib import Path

import numpy as np
import torch
from torch import nn

def tank(u, h0=0.5):
    """真實系統（只用來產生資料）：dh/dt = 0.5u − 0.05√(2gh)，取樣 1 s，槽高 2 m"""
    h, out = h0, [h0]
    for uk in u:
        for _ in range(10):                                            # 每秒內 RK4 十步
            f = lambda x: 0.5 * uk - 0.05 * math.sqrt(2 * 9.8 * max(x, 0.0))
            k1 = f(h); k2 = f(h + 0.05 * k1); k3 = f(h + 0.05 * k2); k4 = f(h + 0.1 * k3)
            h = min(max(h + 0.1 / 6 * (k1 + 2 * k2 + 2 * k3 + k4), 0.0), 2.0)
        out.append(h)
    return np.array(out)

def steps_input(n, rng, hi=0.62):
    """隨機階梯輸入：每 10～60 s 換一個 0～0.62 的閥開度（涵蓋 0～2 m 的操作範圍）"""
    u = []
    while len(u) < n:
        u += [rng.uniform(0, hi)] * int(rng.integers(10, 61))
    return np.array(u[:n])

rng = np.random.default_rng(0)
u_tr = steps_input(3000, rng); y_tr = tank(u_tr)                        # 訓練資料：50 分鐘
u_te = steps_input(600, rng); y_te = tank(u_te)                         # 測試資料：另外 10 分鐘
X = lambda y, u: np.stack([y[1:-1], y[:-2], u[1:], u[:-1]], 1)          # [y(k), y(k−1), u(k), u(k−1)]
Xtr, Ytr = X(y_tr, u_tr), y_tr[2:]

# 線性 ARX：最小平方
A = np.c_[Xtr, np.ones(len(Xtr))]
theta = np.linalg.lstsq(A, Ytr, rcond=None)[0]
arx = lambda x: np.r_[x, 1.0] @ theta

# 神經網路 NARX：輸入先標準化，輸出水位的變化量 Δy
mu, sd = Xtr.mean(0), Xtr.std(0)
torch.manual_seed(0)
net = nn.Sequential(nn.Linear(4, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))
opt = torch.optim.Adam(net.parameters(), lr=3e-3)
xt = torch.tensor((Xtr - mu) / sd, dtype=torch.float32)
dt_ = torch.tensor(Ytr - Xtr[:, 0], dtype=torch.float32).unsqueeze(1)
for epoch in range(3000):                                               # 一步預測的誤差（序列–平行模式）
    loss = nn.functional.mse_loss(net(xt), dt_)
    opt.zero_grad(); loss.backward(); opt.step()
narx = lambda x: x[0] + net(torch.tensor((np.asarray(x) - mu) / sd, dtype=torch.float32)).item()

def one_step(model, y, u):
    return np.array([model(x) for x in X(y, u)])

def free_run(model, y, u):
    """自由模擬（平行模式）：只給前兩個真實水位，之後都拿模型自己的預測回授"""
    p = list(y[:2])
    for k in range(1, len(u)):
        p.append(min(max(model([p[k], p[k - 1], u[k], u[k - 1]]), 0.0), 2.0))
    return np.array(p)

rm = lambda a, b: float(np.sqrt(np.mean((a - b) ** 2)))
print("ARX 係數 [a1 a2 b1 b2 c] =", np.round(theta, 4))
for name, m in (("線性 ARX", arx), ("神經網路 NARX", narx)):
    e1 = rm(one_step(m, y_te, u_te), y_te[2:])
    fr = free_run(m, y_te, u_te)
    print(f"{name}：測試資料的一步預測 RMSE {e1:.4f} m，自由模擬 600 s 的 RMSE {rm(fr, y_te[:len(fr)]):.4f} m")

layers = [[m.weight.tolist(), m.bias.tolist()] for m in net if isinstance(m, nn.Linear)]
fr_n, fr_a = free_run(narx, y_te, u_te), free_run(arx, y_te, u_te)
Path(__file__).with_name("tank_narx.json").write_text(json.dumps({
    "layers": layers, "mu": mu.tolist(), "sd": sd.tolist(), "arx": theta.tolist(),
    "test": {"u": u_te.tolist(), "y": y_te.tolist(), "narx": fr_n.tolist(), "arx": fr_a.tolist()}}), encoding="utf-8")