# 硬性”类别：属于哪个类别； “软性”类别：属于每个类别的概率
# 表示分类数据的简单方法：独热编码（one-hot encoding）。
# 独热编码是一个向量，它的分量和类别一样多。
# 类别对应的分量设置为1，其他所有分量设置为0
# 与线性回归一样，softmax回归也是一个单层神经网络，softmax回归的输出层也是全连接层


# 问题：每次输入是一个2*2的灰度图像
# 我们可以用一个标量表示每个像素值，每个图像对应四个特征。 此外，假设每个图像属于类别“猫”“鸡”和“狗”中的一个。
# 我们有4个特征和3个可能的输出类别， 我们将需要12个标量来表示权重（带下标的）， 3个标量来表示偏置（带下标的）
# 要将输出视为概率，我们必须保证在任何数据上的输出都是非负的且总和为1
# softmax函数能够将未规范化的预测变换为非负数并且总和为1，同时让模型保持可导的性质。
# 为了完成这一目标，我们首先对每个未规范化的预测求幂，这样可以确保输出非负。
# 为了确保最终输出的概率值总和为1，我们再让每个求幂后的结果除以它们的总和
# 我们需要一个训练的目标函数，来激励模型精准地估计概率。 例如， 在分类器输出0.5的所有样本中，我们希望这些样本是刚好有一半实际上属于预测的类别。 这个属性叫做校准（calibration）。
# 损失函数：最大似然估计
# softmax运算获取一个向量并将其映射为概率。
# softmax回归适用于分类问题，它使用了softmax运算中输出类别的概率分布。
# 交叉熵是一个衡量两个概率分布之间差异的很好的度量，它测量给定模型编码数据所需的比特数。
# 交叉熵：主观概率为Q的观察者在看到根据概率P生成的数据时的预期惊异

# 实现softmax由三个步骤组成：
# 对每个项求幂（使用exp）；
# 对每一行求和（小批量中每个样本是一行），得到每个样本的规范化常数；
# 将每一行除以其规范化常数，确保结果的和为1。

# softmax_pytorch_fmnist.py  --  纯 PyTorch 的 softmax 回归示例（可在 PyCharm 运行）
import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
import torch.nn.functional as F
import matplotlib.pyplot as plt
import numpy as np

# ----- 超参数 -----
batch_size = 256
lr = 0.1
num_epochs = 10
num_inputs = 784  # 28*28
num_outputs = 10  # 10 classes

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print("Device:", device)

# ----- 数据集 -----
transform = transforms.Compose([
    transforms.ToTensor(),       # [0,1], shape (1,28,28)
])

train_dataset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=transform)
train_iter = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=2, pin_memory=True)
test_iter = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=2, pin_memory=True)

# ----- 参数（手写模型） -----
W = torch.randn(num_inputs, num_outputs, device=device) * 0.01
W.requires_grad_(True)
b = torch.zeros(num_outputs, device=device, requires_grad=True)

# ----- 辅助函数 -----
def softmax(X):
    # X: (batch_size, num_outputs)
    X_exp = torch.exp(X - X.max(dim=1, keepdim=True)[0])  # 稳定数值
    partition = X_exp.sum(dim=1, keepdim=True)
    return X_exp / partition

def net(X):
    X = X.reshape(-1, num_inputs)
    return softmax(torch.matmul(X, W) + b)

def cross_entropy(y_hat, y):
    # y_hat: (batch, num_outputs)  ; y: (batch,)
    # 避免 log(0) -> 数值稳定
    return -torch.log(y_hat[torch.arange(len(y_hat)), y] + 1e-12)

def accuracy(y_hat, y):
    if y_hat.dim() > 1 and y_hat.shape[1] > 1:
        y_hat_labels = y_hat.argmax(dim=1)
    else:
        y_hat_labels = y_hat
    cmp = (y_hat_labels.type(y.dtype) == y)
    return float(cmp.type(torch.float32).sum())

def evaluate_accuracy(net_fn, data_iter):
    net_fn_eval = net_fn  # 在这里 net 是函数，用 torch.no_grad() 即可
    acc_sum = 0.0
    n = 0
    with torch.no_grad():
        for X, y in data_iter:
            X = X.to(device)
            y = y.to(device)
            y_hat = net_fn_eval(X)
            acc_sum += accuracy(y_hat, y)
            n += y.numel()
    return acc_sum / n

# ----- 优化器 -----
optimizer = torch.optim.SGD([W, b], lr=lr)

# ----- 训练函数 -----
def train_epoch(net_fn, train_iter, loss_fn, updater):
    # 返回 (训练损失, 训练精度)
    train_loss_sum = 0.0
    train_acc_sum = 0.0
    n = 0
    for X, y in train_iter:
        X = X.to(device)
        y = y.to(device)
        y_hat = net_fn(X)
        l = loss_fn(y_hat, y)   # shape (batch,)
        # 梯度
        updater.zero_grad()
        l.mean().backward()
        updater.step()
        train_loss_sum += float(l.sum().item())
        train_acc_sum += accuracy(y_hat, y)
        n += y.numel()
    return train_loss_sum / n, train_acc_sum / n

def train(net_fn, train_iter, test_iter, loss_fn, num_epochs, updater):
    for epoch in range(num_epochs):
        train_loss, train_acc = train_epoch(net_fn, train_iter, loss_fn, updater)
        test_acc = evaluate_accuracy(net_fn, test_iter)
        print(f"Epoch {epoch+1}: train loss {train_loss:.4f}, train acc {train_acc:.4f}, test acc {test_acc:.4f}")

# ----- 训练 -----
# train(net, train_iter, test_iter, cross_entropy, num_epochs, optimizer)

# ----- 预测与显示 -----
labels_map = ['t-shirt', 'trouser', 'pullover', 'dress', 'coat',
              'sandal', 'shirt', 'sneaker', 'bag', 'ankle boot']

def get_fashion_mnist_labels(indices):
    return [labels_map[i] for i in indices]

def show_images(imgs, titles=None, nrows=1, ncols=6, figsize=(10, 3)):
    plt.figure(figsize=figsize)
    for i in range(len(imgs)):
        plt.subplot(nrows, ncols, i+1)
        plt.imshow(imgs[i].squeeze(), cmap='gray')
        if titles:
            plt.title(titles[i])
        plt.axis('off')
    plt.tight_layout()
    plt.show()

def predict_and_show(net_fn, test_iter, n=6):
    for X, y in test_iter:
        X = X.to(device)
        y = y.to(device)
        break
    preds = net_fn(X).argmax(dim=1).cpu().numpy()
    trues = y.cpu().numpy()
    imgs = X.cpu().numpy()
    titles = [f"{get_fashion_mnist_labels([t])[0]}\n{get_fashion_mnist_labels([p])[0]}" for t,p in zip(trues[:n], preds[:n])]
    show_images(imgs[:n], titles=titles, nrows=1, ncols=n, figsize=(12,2))

# predict_and_show(net, test_iter, n=6)

if __name__ == '__main__':
    # 把下面三行放进这个 if 块中
    # 在 macOS 或 Windows 上运行 .py 文件时，Python 使用的是 spawn 模式（不是 Linux 的 fork 模式）
    # 否则在启动 DataLoader 子进程时，Python 会“重新执行”整个文件，导致递归创建进程、直接崩溃。

    # 训练
    train(net, train_iter, test_iter, cross_entropy, num_epochs, optimizer)
    # 预测
    predict_and_show(net, test_iter, n=6)


# FashionMNIST 的每张图像是 1 × 28 × 28，即 单通道灰度图。
# 显示时，matplotlib 的 imshow() 默认用灰度色表（灰阶）。

softmax回归的简洁实现：

import matplotlib
matplotlib.use('TkAgg')  # 或 'Qt5Agg' / 'MacOSX'，取决于系统
import matplotlib.pyplot as plt
import torch
from torch import nn, optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from d2l import torch as d2l  # d2l 的 PyTorch 版

# ===================== 超参数 =====================
batch_size = 256
num_epochs = 10
lr = 0.1
num_inputs = 28*28
num_outputs = 10
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# ===================== 数据加载 =====================
train_dataset = datasets.FashionMNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())
test_dataset = datasets.FashionMNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())

train_iter = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_iter = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

# ===================== 模型 =====================
# 保持和 MXNet 版本一样风格：Sequential + Dense
net = nn.Sequential()
net.add_module('flatten', nn.Flatten())
net.add_module('linear', nn.Linear(num_inputs, num_outputs))
net.to(device)

# ===================== 损失函数和优化器 =====================
loss = nn.CrossEntropyLoss(reduction='none')  # PyTorch 交叉熵
trainer = optim.SGD(net.parameters(), lr=lr)

# ===================== 训练 =====================
# 使用 d2l 的 PyTorch 版本 train_ch3
d2l.train_ch3(net, train_iter, test_iter, loss, num_epochs, trainer)

# 以上只有训练内容，无画图