# 平移不變性：只用畫在「左半邊」的圖形訓練，拿畫在「右半邊」的圖形測試，比較卷積網路與全連接網路
import torch
from torch import nn

def make(n, side, seed):
    """16×16 灰階影像：0＝空心方框、1＝十字、2＝斜線叉，位置限制在左半邊或右半邊，加雜訊"""
    g = torch.Generator().manual_seed(seed)
    X, y = torch.zeros(n, 1, 16, 16), torch.randint(0, 3, (n,), generator=g)
    for i in range(n):
        s = int(torch.randint(4, 7, (1,), generator=g))                # 圖形大小 4～6
        r = int(torch.randint(0, 16 - s + 1, (1,), generator=g))
        c0 = 0 if side == "left" else 8
        c = c0 + int(torch.randint(0, 8 - s + 1, (1,), generator=g))  # 只在左（0～7 行）或右（8～15 行）
        img = X[i, 0]
        if y[i] == 0:
            img[r, c:c + s] = img[r + s - 1, c:c + s] = 1
            img[r:r + s, c] = img[r:r + s, c + s - 1] = 1
        elif y[i] == 1:
            img[r + s // 2, c:c + s] = 1
            img[r:r + s, c + s // 2] = 1
        else:
            for k in range(s):
                img[r + k, c + k] = img[r + k, c + s - 1 - k] = 1
    return X + 0.15 * torch.randn(X.shape, generator=g), y

Xtr, ytr = make(1500, "left", 1)
Xte_l, yte_l = make(500, "left", 2)
Xte_r, yte_r = make(500, "right", 3)

cnn = nn.Sequential(nn.Conv2d(1, 8, 3, padding=1), nn.ReLU(), nn.MaxPool2d(2),     # 16×16 → 8×8
                    nn.Conv2d(8, 16, 3, padding=1), nn.ReLU(),
                    nn.AdaptiveMaxPool2d(1), nn.Flatten(), nn.Linear(16, 3))      # 全域最大池化：不管圖形在哪裡
mlp = nn.Sequential(nn.Flatten(), nn.Linear(256, 32), nn.ReLU(), nn.Linear(32, 3))

for name, net in [("卷積網路 CNN", cnn), ("全連接網路 MLP", mlp)]:
    torch.manual_seed(0)
    for m in net.modules():
        if hasattr(m, "reset_parameters"):
            m.reset_parameters()
    opt = torch.optim.Adam(net.parameters(), lr=0.01)
    for epoch in range(30):
        for idx in torch.randperm(len(Xtr)).split(50):
            opt.zero_grad()
            nn.functional.cross_entropy(net(Xtr[idx]), ytr[idx]).backward()
            opt.step()
    with torch.no_grad():
        acc = lambda X, y: (net(X).argmax(1) == y).float().mean().item()
        n_par = sum(p.numel() for p in net.parameters())
        print(f"{name}：參數 {n_par:5d} 個，左半邊（看過的位置）準確率 {acc(Xte_l, yte_l):.1%}，右半邊（沒看過的位置）{acc(Xte_r, yte_r):.1%}")