import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import transforms
from PIL import Image


class FoodCNN(nn.Module):
    """
    食物分类CNN模型
    基于CIFAR10结构，适配3分类任务
    use_internal_preprocess 是否在模型内部预处理
    训练的时候，
    """
    def __init__(self, use_internal_preprocess=False):
        super(FoodCNN, self).__init__()
        self.use_internal_preprocess = use_internal_preprocess
        
        # 图片预处理变换（仅在推理时使用）
        self.preprocess = transforms.Compose([
            transforms.Resize((32, 32)),
            transforms.ToTensor(),
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])
        
        # 内部预处理变换（用于已经是tensor但未归一化的数据）
        self.internal_preprocess = transforms.Compose([
            # 不包含Resize，因为训练时已经在DataLoader中处理了
            transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
        ])
        
        # 第一个卷积块
        self.conv1 = nn.Conv2d(3, 32, 3, padding=1)
        self.conv2 = nn.Conv2d(32, 32, 3, padding=1)
        self.pool1 = nn.MaxPool2d(2, 2)
        self.dropout1 = nn.Dropout2d(0.25)
        
        # 第二个卷积块
        self.conv3 = nn.Conv2d(32, 64, 3, padding=1)
        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
        self.pool2 = nn.MaxPool2d(2, 2)
        self.dropout2 = nn.Dropout2d(0.25)
        
        # 第三个卷积块
        self.conv5 = nn.Conv2d(64, 128, 3, padding=1)
        self.conv6 = nn.Conv2d(128, 128, 3, padding=1)
        self.pool3 = nn.MaxPool2d(2, 2)
        self.dropout3 = nn.Dropout2d(0.25)
        
        # 全连接层
        self.fc1 = nn.Linear(128 * 4 * 4, 512)
        self.dropout4 = nn.Dropout(0.5)
        self.fc2 = nn.Linear(512, 3)  # 3分类
        
    def preprocess_image(self, image):
        """
        预处理单张图片
        
        Args:
            image: PIL Image 或 numpy array
            
        Returns:
            torch.Tensor: 预处理后的张量，形状为 (1, 3, 32, 32)
        """
        if not isinstance(image, Image.Image):
            # 如果是numpy array，转换为PIL Image
            if hasattr(image, 'shape'):
                image = Image.fromarray(image)
            else:
                raise ValueError("输入必须是PIL Image或numpy array")
        
        # 应用预处理变换
        processed = self.preprocess(image)
        # 添加batch维度
        return processed.unsqueeze(0)
    
    def forward(self, x):
        # 如果启用内部预处理且输入是tensor
        if self.use_internal_preprocess and isinstance(x, torch.Tensor):
            x = self.internal_preprocess(x)
        
        # 第一个卷积块
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        x = self.pool1(x)
        x = self.dropout1(x)
        
        # 第二个卷积块
        x = F.relu(self.conv3(x))
        x = F.relu(self.conv4(x))
        x = self.pool2(x)
        x = self.dropout2(x)
        
        # 第三个卷积块
        x = F.relu(self.conv5(x))
        x = F.relu(self.conv6(x))
        x = self.pool3(x)
        x = self.dropout3(x)
        
        # 展平
        x = x.view(-1, 128 * 4 * 4)
        
        # 全连接层
        x = F.relu(self.fc1(x))
        x = self.dropout4(x)
        x = self.fc2(x)
        
        return x


def create_food_cnn(use_internal_preprocess=False):
    """
    创建食物分类CNN模型
    
    Args:
        use_internal_preprocess: 是否在forward中进行预处理
        
    Returns:
        FoodCNN: 网络模型实例
    """
    return FoodCNN(use_internal_preprocess=use_internal_preprocess)


if __name__ == "__main__":
    # 测试网络
    model = create_food_cnn()
    print(f"模型参数数量: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
    
    # 测试前向传播
    dummy_input = torch.randn(1, 3, 32, 32)
    output = model(dummy_input)
    print(f"输出形状: {output.shape}")
    
    # 测试图片预处理
    try:
        import numpy as np
        # 创建一个测试图片 (RGB格式)
        test_image = Image.fromarray(np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8))
        processed = model.preprocess_image(test_image)
        print(f"预处理后图片形状: {processed.shape}")
        print(f"预处理后数值范围: [{processed.min():.3f}, {processed.max():.3f}]")
    except Exception as e:
        print(f"预处理测试失败: {e}")