"""
03_world_model_planning.py
================================
用最少的程式（只靠 numpy）示範「世界模型 + 在腦中規劃」這整套想法。
對照講義 docs/02_世界模型與規劃.md。

我們會做三件事：
  1. 有一個「真實世界」：2D 平面上的一個點，可以施力推它移動（含摩擦力）。
     任務：從起點到達目標，同時避開中間的障礙物（想像成花瓶）。
  2. 「學」一個世界模型：用隨機亂動蒐集到的資料，擬合一個前向模型
        s[t+1] = f(s[t], a[t])
     （這裡的動力學是線性的，所以用最小平方法 lstsq 幾乎可以完美學到。
       對照訊號與系統：這就是學出一個狀態空間前向/受控體模型 x[n+1]=f(x[n],u[n])。）
  3. 用這個學到的模型「在腦中試走」做規劃（隨機射擊 + 模型預測控制 MPC）：
        - 從現在狀態，隨機產生很多條候選動作序列
        - 全部丟進「學到的模型」往前模擬（完全不碰真實世界）
        - 幫每條打分：離目標多近？有沒有撞到花瓶？動作省不省力？
        - 挑分數最好的，只執行第一步，然後重新觀察、再想一次（這就是 MPC）

執行：python3 03_world_model_planning.py
（不需要 GPU，不需要安裝任何東西，只要 numpy）
"""

import numpy as np

rng = np.random.default_rng(42)

# ---------- 環境設定 ----------
DT = 0.2
FRICTION = 0.85
START = np.array([0.0, 0.0, 0.0, 0.0])   # 狀態 = [x, y, vx, vy]
GOAL = np.array([5.0, 5.0])
OBSTACLE_C = np.array([2.5, 2.5])         # 障礙物（花瓶）圓心
OBSTACLE_R = 1.0                          # 障礙物半徑
STATE_DIM, ACTION_DIM = 4, 2


def true_step(state, action):
    """真實世界的動力學（規劃時不准直接用它，只能用學到的模型）。"""
    action = np.clip(action, -1.0, 1.0)
    x, y, vx, vy = state
    vx = FRICTION * vx + action[0] * DT
    vy = FRICTION * vy + action[1] * DT
    x = x + vx * DT
    y = y + vy * DT
    return np.array([x, y, vx, vy])


# ---------- 第 2 步：蒐集資料並「學」世界模型 ----------
def collect_data(n=6000):
    """從隨機狀態施隨機力，記錄 (狀態, 動作) -> 下一個狀態。"""
    S, A, S2 = [], [], []
    for _ in range(n):
        s = np.concatenate([rng.uniform(-6, 6, 2), rng.uniform(-3, 3, 2)])
        a = rng.uniform(-1, 1, 2)
        S.append(s); A.append(a); S2.append(true_step(s, a))
    return np.array(S), np.array(A), np.array(S2)


def learn_model(S, A, S2):
    """用最小平方法擬合線性前向模型：S2 ≈ [S, A, 1] @ W。回傳 W。"""
    X = np.concatenate([S, A, np.ones((len(S), 1))], axis=1)  # 加一欄常數當 bias
    W, *_ = np.linalg.lstsq(X, S2, rcond=None)
    pred = X @ W
    mse = np.mean((pred - S2) ** 2)
    return W, mse


def model_step_batch(states, actions, W):
    """用學到的模型，平行預測 N 個狀態各走一步（在腦中模擬）。"""
    X = np.concatenate([states, actions, np.ones((len(states), 1))], axis=1)
    return X @ W


# ---------- 第 3 步：用學到的模型做規劃（隨機射擊 + MPC）----------
def rollout_cost(s0, action_seqs, W):
    """
    把 N 條動作序列丟進學到的模型往前模擬，回傳每條的總成本。
    action_seqs: (N, H, 2)
    """
    N, H, _ = action_seqs.shape
    states = np.tile(s0, (N, 1)).astype(float)   # (N, 4)
    cost = np.zeros(N)
    for t in range(H):
        states = model_step_batch(states, action_seqs[:, t, :], W)
        pos = states[:, :2]
        # 1) 離目標越近越好
        cost += np.linalg.norm(pos - GOAL, axis=1)
        # 2) 撞到花瓶 → 大幅懲罰
        d_obs = np.linalg.norm(pos - OBSTACLE_C, axis=1)
        cost += np.where(d_obs < OBSTACLE_R, 200.0, 0.0)
        # 3) 動作省力（小懲罰）
        cost += 0.01 * np.sum(action_seqs[:, t, :] ** 2, axis=1)
    # 終點離目標的額外懲罰（鼓勵真的抵達）
    cost += 10.0 * np.linalg.norm(states[:, :2] - GOAL, axis=1)
    return cost


def plan_action(s0, W, n_samples=2000, horizon=15):
    """隨機射擊：抽很多條動作序列，用模型評分，回傳分數最好那條的『第一個動作』。"""
    action_seqs = rng.uniform(-1, 1, size=(n_samples, horizon, ACTION_DIM))
    costs = rollout_cost(s0, action_seqs, W)
    best = np.argmin(costs)
    return action_seqs[best, 0, :]            # 只取第一步（MPC 精神）


def run(policy, W=None, steps=60):
    """在真實世界跑一段；policy 決定每一步的動作。回傳軌跡。"""
    s = START.copy()
    traj = [s[:2].copy()]
    hit = False
    for _ in range(steps):
        if policy == "plan":
            a = plan_action(s, W)
        else:  # 隨機亂走當對照組
            a = rng.uniform(-1, 1, 2)
        s = true_step(s, a)              # 注意：用真實世界執行（規劃用的是學到的模型）
        traj.append(s[:2].copy())
        if np.linalg.norm(s[:2] - OBSTACLE_C) < OBSTACLE_R:
            hit = True
        if np.linalg.norm(s[:2] - GOAL) < 0.3:
            break
    return np.array(traj), hit


def main():
    print("=" * 64)
    print("任務：從 (0,0) 走到", GOAL.tolist(),
          "，避開圓心", OBSTACLE_C.tolist(), "半徑", OBSTACLE_R, "的花瓶")
    print("=" * 64)

    # 學世界模型
    S, A, S2 = collect_data()
    W, mse = learn_model(S, A, S2)
    print(f"[學世界模型] 用 {len(S)} 筆隨機資料擬合線性前向模型，"
          f"預測誤差 MSE = {mse:.2e}（≈0 代表幾乎完美學到動力學）")

    # 用學到的模型規劃
    traj, hit = run("plan", W, steps=60)
    final_d = np.linalg.norm(traj[-1] - GOAL)
    print("\n[用世界模型規劃 MPC]")
    print(f"  走了 {len(traj)-1} 步，最後位置 ({traj[-1,0]:.2f}, {traj[-1,1]:.2f})")
    print(f"  離目標距離 = {final_d:.2f}（<0.3 視為抵達）")
    print(f"  是否撞到花瓶：{'是 ✗' if hit else '否 ✓'}")

    # 對照組：隨機亂走
    rtraj, rhit = run("random", None, steps=60)
    print("\n[對照組：隨機亂走]")
    print(f"  最後離目標距離 = {np.linalg.norm(rtraj[-1]-GOAL):.2f}")

    print("\n結論：規劃版會『在腦中想像很多條路線、挑一條到得了又不撞花瓶的』，")
    print("      所以能繞開障礙抵達目標；隨機亂走則到不了。")
    print("      把這裡的『線性模型』換成神經網路世界模型，就是真實機器人的做法。")


if __name__ == "__main__":
    main()

"""
06_rl_qlearning.py
================================
強化學習（RL）入門：用 Q-learning 教一個 agent 走迷宮到終點、避開陷阱。
這銜接 docs/02、docs/03 的物理 AI / 世界模型——機器人的「策略」常用 RL 訓練。

和 03（世界模型規劃）的差別：
  - 03 是「有模型」：先學一個世界模型，再在腦中規劃。
  - 這支是「無模型」：agent 不知道環境規則，純粹靠不斷嘗試、從獎勵中學會一張「狀態→動作該多好」的表（Q 表）。
  兩者是 RL 的兩大路線；DreamerV3 那類「世界模型 RL」其實是把兩者結合。

Q-learning 更新公式（每走一步就更新一次）：
    Q(s,a) ← Q(s,a) + α · [ r + γ·max_a' Q(s',a') − Q(s,a) ]

執行：python3 06_rl_qlearning.py   （只需要 numpy）
"""

import numpy as np

rng = np.random.default_rng(0)

# ---------- 迷宮設定 ----------
#  S=起點  G=終點(+1)  X=陷阱(-1)  .=空地
GRID = [
    "S....",
    ".XX.X",
    "...X.",
    ".X...",
    "X...G",
]
ROWS, COLS = len(GRID), len(GRID[0])
ACTIONS = [(-1, 0), (1, 0), (0, -1), (0, 1)]   # 上 下 左 右
ARROW = ["↑", "↓", "←", "→"]

start = goal = None
traps = set()
for r in range(ROWS):
    for c in range(COLS):
        ch = GRID[r][c]
        if ch == "S": start = (r, c)
        elif ch == "G": goal = (r, c)
        elif ch == "X": traps.add((r, c))


def step(state, a):
    """環境：給狀態與動作，回傳 (下一狀態, 獎勵, 是否結束)。agent 事先並不知道這個函數。"""
    dr, dc = ACTIONS[a]
    nr, nc = state[0] + dr, state[1] + dc
    if not (0 <= nr < ROWS and 0 <= nc < COLS):   # 撞牆 → 留在原地、小懲罰
        return state, -0.1, False
    ns = (nr, nc)
    if ns == goal:
        return ns, 1.0, True                      # 到終點
    if ns in traps:
        return ns, -1.0, True                     # 掉陷阱
    return ns, -0.02, False                       # 走一步的小成本（鼓勵走短路）


def train(episodes=4000, alpha=0.1, gamma=0.95):
    Q = np.zeros((ROWS, COLS, len(ACTIONS)))
    for ep in range(episodes):
        s = start
        eps = max(0.05, 1.0 - ep / (episodes * 0.6))   # 探索率隨時間下降
        for _ in range(100):
            if rng.random() < eps:
                a = rng.integers(len(ACTIONS))         # 探索：隨機試
            else:
                a = int(np.argmax(Q[s[0], s[1]]))      # 利用：選目前最好的
            ns, r, done = step(s, a)
            # Q-learning 更新
            best_next = 0.0 if done else np.max(Q[ns[0], ns[1]])
            td = r + gamma * best_next - Q[s[0], s[1], a]
            Q[s[0], s[1], a] += alpha * td
            s = ns
            if done:
                break
    return Q


def greedy_path(Q, max_steps=50):
    s, path = start, [start]
    for _ in range(max_steps):
        a = int(np.argmax(Q[s[0], s[1]]))
        ns, _, done = step(s, a)
        path.append(ns)
        s = ns
        if done:
            break
    return path, (s == goal)


def show_policy(Q):
    print("學到的策略（每格顯示最佳動作；G=終點，X=陷阱）：")
    for r in range(ROWS):
        row = []
        for c in range(COLS):
            if (r, c) == goal: row.append(" G")
            elif (r, c) in traps: row.append(" X")
            else: row.append(" " + ARROW[int(np.argmax(Q[r, c]))])
        print("  " + "".join(row))


def main():
    print(f"迷宮 {ROWS}x{COLS}，起點 {start}，終點 {goal}，陷阱 {sorted(traps)}\n")
    Q = train()
    show_policy(Q)
    path, reached = greedy_path(Q)
    print(f"\n依學到的策略走：{path}")
    print(f"是否到達終點：{'是 ✓' if reached else '否 ✗'}，共 {len(path)-1} 步")
    print("\n重點：agent 一開始什麼都不知道，純靠『嘗試 + 獎勵』學會一張 Q 表，")
    print("      最後能挑出一條避開陷阱、最快到終點的路。這就是強化學習的核心。")


if __name__ == "__main__":
    main()