# ANFIS 系統鑑別：用一階 Sugeno 模糊推論系統（寫成網路）學水槽（同模糊控制章）的動態
# 輸入：水位 h(k)、閥開度 u(k)；輸出：一秒後水位的變化量 Δh = h(k+1) − h(k)
# 資料與 tank_narx.py 完全相同（同一個亂數種子），方便和 NARX、ARX 比較
import json
import math
from pathlib import Path

import numpy as np
import torch

torch.set_default_dtype(torch.float64)

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):
            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):
    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)
u_te = steps_input(600, rng); y_te = tank(u_te)
X = torch.tensor(np.c_[y_tr[:-1], u_tr])                                # [h(k), u(k)]
T = torch.tensor(y_tr[1:] - y_tr[:-1])                                  # Δh

class ANFIS:
    """水位 3 個、閥開度 2 個高斯歸屬函數 → 6 條規則；規則 r 的結論 f_r = p_r h + q_r u + s_r"""
    def __init__(self):
        self.c = [torch.tensor([0.0, 1.0, 2.0]), torch.tensor([0.0, 0.62])]      # 中心：等距
        w = lambda a, b, m: (b - a) / (m - 1) / (2 * math.sqrt(2 * math.log(2)))  # 寬度：相鄰兩個交在 0.5
        self.s = [torch.full((3,), w(0, 2, 3)), torch.full((2,), w(0, 0.62, 2))]
        self.rules = [(i, j) for i in range(3) for j in range(2)]
        self.P = torch.zeros(6, 3)

    def norm_firing(self, X):                                           # 第 1～3 層：歸屬度 → 相乘 → 正規化
        mu = [torch.exp(-(X[:, d:d + 1] - self.c[d]) ** 2 / (2 * self.s[d] ** 2)) for d in range(2)]
        W = torch.stack([mu[0][:, i] * mu[1][:, j] for i, j in self.rules], 1)
        return W / W.sum(1, keepdim=True)

    def __call__(self, X):                                              # 第 4～5 層：加權平均各規則的線性結論
        Xa = torch.cat([X, torch.ones(len(X), 1)], 1)
        return (self.norm_firing(X) * (Xa @ self.P.T)).sum(1)

    def lse(self, X, T):
        """前提固定時，輸出對 18 個結論參數是線性的 → 最小平方一次解出"""
        with torch.no_grad():
            Xa = torch.cat([X, torch.ones(len(X), 1)], 1)
            A = (self.norm_firing(X)[:, :, None] * Xa[:, None, :]).reshape(len(X), -1)
            self.P = torch.linalg.solve(A.T @ A + 1e-8 * torch.eye(18), A.T @ T).reshape(6, 3)

def train(net, epochs, lr):
    """混合學習：每輪先用最小平方解結論，再用自動微分＋Adam 調歸屬函數的中心與寬度"""
    prem = net.c + net.s
    for p in prem:
        p.requires_grad_(True)
    opt = torch.optim.Adam(prem, lr=lr)
    hist = []
    for _ in range(epochs):
        net.lse(X, T)
        loss = torch.mean((net(X) - T) ** 2)
        hist.append(math.sqrt(loss.item()))
        opt.zero_grad(); loss.backward(); opt.step()
        with torch.no_grad():
            for s in net.s:
                s.clamp_(min=0.05)
    net.lse(X, T)
    for p in prem:
        p.requires_grad_(False)
    return hist

def free_run(net, y, u):
    """自由模擬：只給初始水位，之後都用模型自己的預測"""
    h = [y[0]]
    with torch.no_grad():
        for k in range(len(u)):
            dh = float(net(torch.tensor([[h[-1], u[k]]]))[0])
            h.append(min(max(h[-1] + dh, 0.0), 2.0))
    return np.array(h)

rm = lambda a, b: float(np.sqrt(np.mean((a - b) ** 2)))
fixed, learned = ANFIS(), ANFIS()
fixed.lse(X, T)                                                         # 只學結論（歸屬函數固定在等距）
hist = train(learned, 300, 0.03)                                        # 混合學習 300 輪
for name, net, n in (("固定歸屬函數（只學結論）", fixed, 18), ("混合學習（也調歸屬函數）", learned, 28)):
    one = net(torch.tensor(np.c_[y_te[:-1], u_te])).detach().numpy() + y_te[:-1]
    print(f"{name}：{n} 個參數，一步預測 RMSE {rm(one, y_te[1:]) * 100:.3f} cm，自由模擬 600 s 的 RMSE {rm(free_run(net, y_te, u_te), y_te) * 100:.3f} cm")
print("混合學習後的水位歸屬函數 中心", np.round(learned.c[0].numpy(), 3), "寬度", np.round(learned.s[0].numpy(), 3))
print("　　　　　閥開度歸屬函數 中心", np.round(learned.c[1].numpy(), 3), "寬度", np.round(learned.s[1].numpy(), 3))
print("固定歸屬函數的規則（Δh = p·h + q·u + s）：")
for (i, j), (p, q, s) in zip(fixed.rules, fixed.P.numpy()):
    print(f"  若水位{'低中高'[i]}且閥開度{'小大'[j]}，Δh = {p:+.3f}h {q:+.3f}u {s:+.3f}")

dump = lambda net: {"c": [c.tolist() for c in net.c], "s": [s.tolist() for s in net.s], "P": net.P.tolist()}
Path(__file__).with_name("tank_anfis.json").write_text(json.dumps({
    "fixed": dump(fixed), "learned": dump(learned), "hist": hist,
    "test": {"u": u_te.tolist(), "y": y_te.tolist(), "fixed": free_run(fixed, y_te, u_te).tolist(),
             "learned": free_run(learned, y_te, u_te).tolist()}}), encoding="utf-8")