# 用 PyTorch 重做 Matlab 的倒傳遞範例：newff([-1 2; 0 5], [3,1], {'tansig','purelin'}, 'traingd')
import torch
from torch import nn

torch.manual_seed(0)
p = torch.tensor([[-1., 0.], [-1., 5.], [2., 0.], [2., 5.]])   # 每一列是一筆輸入（Matlab 是每一行）
t = torch.tensor([[-1.], [-1.], [1.], [1.]])

def make_net():
    return nn.Sequential(nn.Linear(2, 3), nn.Tanh(), nn.Linear(3, 1))   # tansig = tanh，purelin = 不加函數

for name, Opt, lr in [("SGD（同 traingd）", torch.optim.SGD, 0.05), ("Adam", torch.optim.Adam, 0.05)]:
    torch.manual_seed(1)
    net = make_net()
    opt = Opt(net.parameters(), lr=lr)
    loss_fn = nn.MSELoss()                       # 同 Matlab 的 mse
    for epoch in range(1, 301):
        opt.zero_grad()                          # 清掉上一步的梯度
        loss = loss_fn(net(p), t)                # 前向＋誤差
        loss.backward()                          # 倒傳遞
        opt.step()                               # 更新權值
        if epoch in (1, 50, 300):
            print(f"{name:14s} 第 {epoch:3d} 回合 mse = {loss.item():.2e}")
    print("  訓練後輸出：", [round(v, 3) for v in net(p).squeeze().tolist()])