# 深度 Q 網路（DQN）學平衡倒單擺：經驗回放＋目標網路（Mnih 等 2015 的做法），環境自己寫、不需安裝 gym
import json
import math
import random
from pathlib import Path

import torch
from torch import nn

def step(s, a):
    """經典 CartPole：台車 1 kg、擺桿 0.1 kg（半長 0.5 m）、推力 ±10 N、每步 0.02 s"""
    x, xd, th, thd = s
    f = 10.0 if a == 1 else -10.0
    c, sn = math.cos(th), math.sin(th)
    tmp = (f + 0.05 * thd * thd * sn) / 1.1
    thacc = (9.8 * sn - c * tmp) / (0.5 * (4 / 3 - 0.1 * c * c / 1.1))
    xacc = tmp - 0.05 * thacc * c / 1.1
    s = (x + 0.02 * xd, xd + 0.02 * xacc, th + 0.02 * thd, thd + 0.02 * thacc)
    return s, abs(s[0]) > 2.4 or abs(s[2]) > 12 * math.pi / 180      # 出界或倒超過 12° 就失敗

random.seed(0); torch.manual_seed(0)
Q = nn.Sequential(nn.Linear(4, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 2))
Q_target = nn.Sequential(nn.Linear(4, 64), nn.ReLU(), nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 2))
Q_target.load_state_dict(Q.state_dict())                               # 目標網路：固定一陣子才更新，讓學習目標穩定
opt = torch.optim.Adam(Q.parameters(), lr=1e-3)
replay, CAP = [], 20000                                                # 經驗回放：存過去的 (s, a, r, s′, done)
def evaluate(net, n=10, seed=100):
    """不探索（永遠選 Q 最大的動作）跑 n 回合，回傳每回合撐的步數"""
    out = []
    for k in range(n):
        rng = random.Random(seed + k)
        s, T = tuple(rng.uniform(-0.05, 0.05) for _ in range(4)), 0
        while T < 500:
            with torch.no_grad():
                s, done = step(s, int(net(torch.tensor(s)).argmax()))
            T += 1
            if done:
                break
        out.append(T)
    return out

steps, lens, tests, best = 0, [], [], (-1, None, 0)
for ep in range(1, 601):
    s, T = tuple(random.uniform(-0.05, 0.05) for _ in range(4)), 0
    while T < 500:
        eps = max(0.05, 1.0 - steps / 10000)                           # ε 從 1 線性降到 0.05
        if random.random() < eps:
            a = random.randrange(2)
        else:
            with torch.no_grad():
                a = int(Q(torch.tensor(s)).argmax())
        s2, done = step(s, a)
        replay.append((s, a, 1.0, s2, float(done)))
        if len(replay) > CAP:
            replay.pop(0)
        s, T, steps = s2, T + 1, steps + 1
        if len(replay) >= 1000:                                        # 從回放中隨機抽一小批，打破資料的時間相關
            S, A, R, S2, D = map(torch.tensor, zip(*random.sample(replay, 64)))
            with torch.no_grad():
                target = R + 0.99 * (1 - D) * Q_target(S2).max(1).values
            q = Q(S).gather(1, A.view(-1, 1)).squeeze(1)
            loss = nn.functional.smooth_l1_loss(q, target)
            opt.zero_grad(); loss.backward(); opt.step()
        if steps % 500 == 0:
            Q_target.load_state_dict(Q.state_dict())
        if done:
            break
    lens.append(T)
    if ep % 50 == 0:
        score = sum(evaluate(Q)) / 10                                  # 每 50 回合測一次，保留表現最好的網路
        tests.append([ep, score])
        if score > best[0]:
            best = (score, {k: v.clone() for k, v in Q.state_dict().items()}, ep)
        print(f"第 {ep:3d} 回合：最近 50 回合平均撐 {sum(lens[-50:]) / 50:5.1f} 步，不探索測試 {score:5.1f} 步")

Q.load_state_dict(best[1])
evals = evaluate(Q, 20, seed=1000)                                      # 用另一組初始狀態測 20 回合
print(f"採用第 {best[2]} 回合的網路，另測 20 回合：平均 {sum(evals) / 20:.1f} 步，撐滿 500 步的有 {evals.count(500)} 回合")
layers = [(m.weight.tolist(), m.bias.tolist()) for m in Q if isinstance(m, nn.Linear)]
Path(__file__).with_name("dqn_cartpole.json").write_text(json.dumps({"layers": layers, "lens": lens, "tests": tests, "best": best[2]}), encoding="utf-8")   # 給網頁動畫