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),interpolation=transforms.InterpolationMode.BILINEAR),
            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")
        
        # 转换为tensor（不做resize）
        tensor = transforms.ToTensor()(image)
        # 添加batch维度
        tensor = tensor.unsqueeze(0)
        
        # 使用与移动端相同的插值方法缩放到32x32
        tensor = F.interpolate(tensor, size=(32, 32), mode='bilinear', align_corners=False)
        
        # ImageNet标准化
        mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
        tensor = (tensor - mean) / std
        
        return tensor
    
    def mobile_preprocess(self, x):
        """
        移动端预处理（TorchScript兼容）
        
        Args:
            x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-255
            
        Returns:
            torch.Tensor: 预处理后的tensor，形状为 [batch, 3, 32, 32]
        """
        # 归一化到 [0, 1]
        # 好像安卓已经做了归一化了。
        # x = x.float() / 255.0
        
        # 缩放到 32x32
        x = F.interpolate(x, size=(32, 32), mode='bilinear', align_corners=False)
        # 缩放到 32x32 - 使用 align_corners=True 来匹配 PIL 的默认行为
        # x = F.interpolate(x, size=(32, 32), mode='bilinear', align_corners=True)

        # ImageNet标准化
        mean = torch.tensor([0.485, 0.456, 0.406], device=x.device).view(1, 3, 1, 1)
        std = torch.tensor([0.229, 0.224, 0.225], device=x.device).view(1, 3, 1, 1)
        x = (x - mean) / std
        
        return x
    
    def _forward_network(self, x):
        """
        网络的核心前向传播（不包含预处理）
        
        Args:
            x: 预处理后的tensor，形状为 [batch, 3, 32, 32]
            
        Returns:
            torch.Tensor: 模型输出
        """
        # 第一个卷积块
        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 forward(self, x):
        # 如果启用内部预处理且输入是tensor
        if self.use_internal_preprocess and isinstance(x, torch.Tensor):
            x = self.internal_preprocess(x)
        
        return self._forward_network(x)
    
    def forward_mobile(self, x):
        """
        移动端前向传播（包含预处理）
        
        Args:
            x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-255
            
        Returns:
            torch.Tensor: 模型输出
        """
        # 移动端预处理
        x = self.mobile_preprocess(x)
        
        # 执行网络前向传播
        return self._forward_network(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)


def create_mobile_food_cnn():
    """
    创建移动端食物分类CNN模型
    
    Returns:
        FoodCNN: 配置为移动端使用的网络模型实例
    """
    
    class MobileFoodCNN(FoodCNN):
        """
        移动端食物分类模型
        重写forward方法以包含预处理
        """
        def __init__(self):
            super().__init__(use_internal_preprocess=False)
        
        def forward(self, x):
            """
            移动端前向传播（自动包含预处理）
            
            Args:
                x: 输入tensor，形状为 [batch, 3, height, width]，值范围 0-255
                
            Returns:
                torch.Tensor: 模型输出
            """
            return self.forward_mobile(x)
    
    return MobileFoodCNN()


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}")