# 用 Bellman 方程式（動態規劃）解 LQR：MCK 系統（m = 1、c = 1、k = 2，同模糊控制章）離散化後，
# 從 V = 0 開始反覆套用 Bellman 方程式，就是 Riccati 迭代，結果和 scipy 解離散代數 Riccati 方程式相同
import numpy as np
from scipy.linalg import expm, solve_discrete_are

dt = 0.1
Ac, Bc = np.array([[0., 1.], [-2., -1.]]), np.array([[0.], [1.]])
M = expm(np.block([[Ac, Bc], [np.zeros((1, 3))]]) * dt)                 # 零階保持離散化
A, B = M[:2, :2], M[:2, 2:]
Q, R = np.diag([10., 1.]), np.array([[0.1]])                           # 成本 Σ (xᵀQx + uᵀRu)

# V(x) = xᵀPx 代入 V(x) = min_u [xᵀQx + uᵀRu + V(Ax + Bu)]，最小化 u 後得到 P 的更新式
P = np.zeros((2, 2))
for k in range(1, 201):
    K = np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A)                   # 這一步的最佳回授 u = −Kx
    P_new = Q + A.T @ P @ (A - B @ K)
    if np.abs(P_new - P).max() < 1e-9 * np.abs(P_new).max():
        P = P_new
        break
    P = P_new
print(f"Bellman 迭代 {k} 次收斂，P =\n{P.round(4)}")
P_are = solve_discrete_are(A, B, Q, R)
print("與 solve_discrete_are 的最大差：", f"{np.abs(P - P_are).max():.1e}")
K = np.linalg.solve(R + B.T @ P @ B, B.T @ P @ A)
print("最佳回授增益 K =", K.round(4), "，閉迴路特徵值大小", np.abs(np.linalg.eigvals(A - B @ K)).round(4))

x = np.array([1.0, 0.0])                                               # 從位移 1、速度 0 出發
for t in range(30):
    x = (A - B @ K) @ x
x_open = np.linalg.matrix_power(A, 30) @ np.array([1.0, 0.0])
print(f"3 秒後 x = {x.round(4)}（不加控制時 x = {x_open.round(4)}）")