# 用遞迴網路做系統鑑別：只看輸入力 F(t) 與位移 x(t) 的資料，學出 MCK 系統（m = 1、c = 1、k = 2，同模糊控制章）
import json
import math
from pathlib import Path

import torch
from torch import nn

DT, N = 0.05, 400                                  # 取樣 0.05 s，每段 20 s

def simulate(F):
    """m ẍ + c ẋ + k x = F，RK4、每個取樣區間內 F 固定"""
    x = v = 0.0
    out = []
    for f in F:
        out.append(x)
        acc = lambda x_, v_: f - 1.0 * v_ - 2.0 * x_
        k1 = (v, acc(x, v)); k2 = (v + DT / 2 * k1[1], acc(x + DT / 2 * k1[0], v + DT / 2 * k1[1]))
        k3 = (v + DT / 2 * k2[1], acc(x + DT / 2 * k2[0], v + DT / 2 * k2[1])); k4 = (v + DT * k3[1], acc(x + DT * k3[0], v + DT * k3[1]))
        x += DT / 6 * (k1[0] + 2 * k2[0] + 2 * k3[0] + k4[0]); v += DT / 6 * (k1[1] + 2 * k2[1] + 2 * k3[1] + k4[1])
    return out

def random_input(g):
    """訓練用輸入：每 1～3 秒隨機換一個 −2～2 的定值"""
    F, level, left = [], 0.0, 0
    for _ in range(N):
        if left == 0:
            level, left = float(4 * torch.rand(1, generator=g) - 2), int(torch.randint(20, 61, (1,), generator=g))
        F.append(level); left -= 1
    return F

g = torch.Generator().manual_seed(0)
Ftr = [random_input(g) for _ in range(40)]
U = torch.tensor(Ftr).unsqueeze(-1)                                 # 形狀（40 段, 400 步, 1）
Y = torch.tensor([simulate(F) for F in Ftr]).unsqueeze(-1)
tests = {"step": [0.0] * 20 + [1.0] * (N - 20),                    # t = 1 s 起單位步階（同模糊章）
         "sine": [math.sin(0.8 * k * DT) * 1.5 for k in range(N)]}  # 訓練時沒出現過的正弦輸入

class SeqNet(nn.Module):
    def __init__(self, cell):
        super().__init__()
        self.rnn = {"RNN": nn.RNN, "LSTM": nn.LSTM}[cell](1, 16, batch_first=True)
        self.out = nn.Linear(16, 1)
    def forward(self, u):
        h, _ = self.rnn(u)                                          # 每一步的隱藏狀態
        return self.out(h)                                          # 每一步預測一個 x

result = {"t": [k * DT for k in range(N)], "true": {k: simulate(F) for k, F in tests.items()}}
for cell in ("RNN", "LSTM"):
    torch.manual_seed(0)
    net = SeqNet(cell)
    opt = torch.optim.Adam(net.parameters(), lr=0.01)
    for epoch in range(800):
        opt.zero_grad()
        loss = nn.functional.mse_loss(net(U), Y)
        loss.backward()                                             # 沿時間倒傳遞（BPTT）
        nn.utils.clip_grad_norm_(net.parameters(), 1.0)             # 梯度裁剪：避免梯度爆炸
        opt.step()
    with torch.no_grad():
        result[cell] = {}
        msg = []
        for k, F in tests.items():
            pred = net(torch.tensor(F).view(1, N, 1)).view(-1)
            rmse = (pred - torch.tensor(result["true"][k])).pow(2).mean().sqrt().item()
            result[cell][k] = pred.tolist()
            msg.append(f"{k} 測試 RMSE {rmse:.4f}")
    print(f"{cell:4s}：訓練 mse {loss.item():.2e}，" + "，".join(msg))
Path(__file__).with_name("mck_lstm.json").write_text(json.dumps(result), encoding="utf-8")   # 給網頁的圖使用