# 物理資訊神經網路（PINN）：水槽關閥排水，只量到前 4 秒的 10 個（有雜訊的）水位，
# 要預測到 10 秒，並同時估計未知的出水係數 c。已知的物理：dh/dt = −c√(2gh)
import json
import math
from pathlib import Path

import numpy as np
import torch
from torch import nn

torch.set_default_dtype(torch.float64)
G, C_TRUE, H0 = 9.8, 0.05, 2.0
truth = lambda t: np.maximum(math.sqrt(H0) - C_TRUE * math.sqrt(G / 2) * t, 0.0) ** 2   # 解析解（只用來產生資料與評分）

rng = np.random.default_rng(1)
t_d = np.linspace(0, 4, 10)
h_d = truth(t_d) + rng.normal(0, 0.02, len(t_d))                         # 量測：標準差 2 cm 的雜訊
td, hd = torch.tensor(t_d)[:, None], torch.tensor(h_d)[:, None]
tc = torch.linspace(0, 10, 101)[:, None].requires_grad_(True)           # 檢查物理的「配置點」：0～10 s
t_all = np.linspace(0, 10, 101)

def make_net(seed):
    torch.manual_seed(seed)
    return nn.Sequential(nn.Linear(1, 32), nn.Tanh(), nn.Linear(32, 32), nn.Tanh(), nn.Linear(32, 1))

def train(physics, iters=6000, lr=2e-3, snap=200):
    net = make_net(0)
    log_c = torch.tensor(math.log(0.1), requires_grad=True)             # 未知參數 c 也當成要學的量（初猜 0.1）
    opt = torch.optim.Adam(list(net.parameters()) + ([log_c] if physics else []), lr=lr)
    h = lambda t: net(t / 10)                                           # 時間先縮放到 0～1
    snaps = []
    for it in range(iters + 1):
        loss = torch.mean((h(td) - hd) ** 2)                            # 資料誤差
        if physics:
            hc = h(tc)
            dh = torch.autograd.grad(hc.sum(), tc, create_graph=True)[0]  # 自動微分求 dh/dt
            res = dh + log_c.exp() * torch.sqrt(2 * G * hc.clamp(min=0) + 1e-12)
            loss = loss + torch.mean(res ** 2)                          # 物理誤差：方程式在配置點上要成立
        if it % snap == 0:
            with torch.no_grad():
                snaps.append({"it": it, "h": h(torch.tensor(t_all)[:, None])[:, 0].tolist(), "c": float(log_c.exp())})
        if it < iters:
            opt.zero_grad(); loss.backward(); opt.step()
    return snaps

nn_s = train(False)
pinn_s = train(True)
far = t_all > 4
rm = lambda h, sel: float(np.sqrt(np.mean((h[sel] - truth(t_all[sel])) ** 2))) * 100
for name, s in (("一般神經網路（只用資料）", nn_s), ("PINN（資料＋物理）", pinn_s)):
    hh = np.array(s[-1]["h"])
    print(f"{name}：0～4 s 的 RMSE {rm(hh, ~far):.2f} cm，4～10 s（沒有資料）的 RMSE {rm(hh, far):.2f} cm")
print(f"PINN 估計的出水係數 c = {pinn_s[-1]['c']:.4f}（真值 {C_TRUE}）")

# 對照：這個例子剛好有解析解，√h 是時間的直線，直接對 √h 做直線擬合也能估 c
a, b = np.polyfit(t_d, np.sqrt(np.maximum(h_d, 0)), 1)
print(f"對照（解析解＋直線擬合）：c = {-a / math.sqrt(G / 2):.4f}")

Path(__file__).with_name("pinn_tank.json").write_text(json.dumps({
    "t": t_all.tolist(), "truth": truth(t_all).tolist(), "td": t_d.tolist(), "hd": h_d.tolist(),
    "nn": nn_s, "pinn": pinn_s}), encoding="utf-8")