import paddle
import paddle.nn as nn
import paddle.nn.functional as F
import numpy as np

# 检查 GPU 并设置设备
if paddle.device.is_compiled_with_cuda():
    paddle.device.set_device('gpu')
    print("使用 GPU 进行训练")
else:
    paddle.device.set_device('cpu')
    print("使用 CPU 进行训练")

# 定义一个简单的 CNN 模型
class SimpleCNN(nn.Layer):
    def __init__(self, num_classes=10):
        super().__init__()
        self.conv1 = nn.Conv2D(in_channels=1, out_channels=16, kernel_size=3, padding=1)
        self.pool1 = nn.MaxPool2D(kernel_size=2, stride=2)
        
        self.conv2 = nn.Conv2D(in_channels=16, out_channels=32, kernel_size=3, padding=1)
        self.pool2 = nn.MaxPool2D(kernel_size=2, stride=2)
        
        self.flatten = nn.Flatten()
        self.fc1 = nn.Linear(in_features=32 * 7 * 7, out_features=128)
        self.dropout = nn.Dropout(p=0.5)
        self.fc2 = nn.Linear(in_features=128, out_features=num_classes)
    
    def forward(self, x):
        x = F.relu(self.conv1(x))
        x = self.pool1(x)
        x = F.relu(self.conv2(x))
        x = self.pool2(x)
        x = self.flatten(x)
        x = F.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.fc2(x)
        return x

# 创建模拟数据（MNIST 风格的数据）
def create_dummy_data():
    # 模拟 32 个 28x28 的灰度图像
    batch_size = 32
    train_data = np.random.randn(batch_size, 1, 28, 28).astype('float32')
    train_labels = np.random.randint(0, 10, (batch_size, 1)).astype('int64')
    
    return paddle.to_tensor(train_data), paddle.to_tensor(train_labels)

# 训练函数
def train_cnn():
    # 初始化模型
    model = SimpleCNN(num_classes=10)
    
    # 创建模拟数据
    train_data, train_labels = create_dummy_data()
    
    # 定义损失函数和优化器
    criterion = nn.CrossEntropyLoss()
    optimizer = paddle.optimizer.Adam(
        parameters=model.parameters(),
        learning_rate=0.001
    )
    
    # 训练循环
    epochs = 10
    for epoch in range(epochs):
        # 前向传播
        outputs = model(train_data)
        
        # 计算损失
        loss = criterion(outputs, train_labels.squeeze())
        
        # 计算准确率
        predictions = paddle.argmax(outputs, axis=1)
        accuracy = paddle.mean(paddle.cast(predictions == train_labels.squeeze(), 'float32'))
        
        # 反向传播和优化
        loss.backward()
        optimizer.step()
        optimizer.clear_grad()
        
        # 使用 .item() 方法获取标量值
        loss_value = loss.item()
        accuracy_value = accuracy.item()
        
        print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss_value:.4f}, Accuracy: {accuracy_value:.4f}')
    
    print("\nCNN 训练完成！")
    return model

# 运行训练
if __name__ == "__main__":
    model = train_cnn()
    
    # 测试模型
    test_data = paddle.to_tensor(np.random.randn(5, 1, 28, 28).astype('float32'))
    with paddle.no_grad():
        predictions = model(test_data)
        predicted_classes = paddle.argmax(predictions, axis=1)
        print(f"\n测试数据预测结果: {predicted_classes.numpy()}")