import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
from PIL import Image
import matplotlib.pyplot as plt

class CatDogClassifier:
    def __init__(self, img_size=128, batch_size=32):
        """
        初始化分类器
        img_size: 图片尺寸
        batch_size: 每次处理多少张图片
        """
        self.img_size = img_size
        self.batch_size = batch_size
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # 定义图片预处理方式
        self.transform = transforms.Compose([
            transforms.Resize((self.img_size, self.img_size)),  # 统一图片大小
            transforms.ToTensor(),                              # 转成数字
            transforms.Normalize(mean=[0.485, 0.456, 0.406],    # 标准化
                               std=[0.229, 0.224, 0.225])
        ])
        
        # 创建AI模型
        self.model = self._create_model()
        self.model.to(self.device)
        
        print(f"✅ 分类器初始化完成，使用设备: {self.device}")

    def _create_model(self):
        """创建AI模型（这是我们的"大脑"）"""
        class SimpleCNN(nn.Module):
            def __init__(self, img_size=128):
                super().__init__()
                # 这些是AI的"神经元层"
                # 第一层：看图片的细节
                self.conv1 = nn.Conv2d(3, 32, 3, padding=1)  # 3个颜色通道→32个特征
                self.relu1 = nn.ReLU()                       # 让AI学会"思考"
                self.pool1 = nn.MaxPool2d(2, 2)              # 抓住重要特征
                
                # 第二层：看图片的形状
                self.conv2 = nn.Conv2d(32, 64, 3, padding=1) # 32个特征→64个特征
                self.relu2 = nn.ReLU()
                self.pool2 = nn.MaxPool2d(2, 2)
                
                # 第三层：做最终判断
                pool_size = img_size // 4  # 经过两次池化，图片变小了
                self.fc1 = nn.Linear(64 * pool_size * pool_size, 512)  # 全连接层
                self.relu3 = nn.ReLU()
                self.fc2 = nn.Linear(512, 2)  # 输出2个结果（猫或狗）

            def forward(self, x):
                # 这是AI"看"图片的过程
                x = self.pool1(self.relu1(self.conv1(x)))  # 第一层处理
                x = self.pool2(self.relu2(self.conv2(x)))  # 第二层处理
                x = x.view(x.size(0), -1)                 # 展平成一行
                x = self.relu3(self.fc1(x))               # 第三层处理
                x = self.fc2(x)                           # 输出结果
                return x

        return SimpleCNN(self.img_size)

    def train_model(self, data_dir, num_epochs=10, learning_rate=0.001):
        """
        训练模型
        data_dir: 图片数据存放的文件夹
        num_epochs: 训练多少轮
        learning_rate: 学习速度
        """
        
        # 加载图片数据
        dataset = datasets.ImageFolder(
            root=data_dir,        # 图片文件夹路径
            transform=self.transform  # 预处理方式
        )
        
        data_loader = DataLoader(
            dataset,
            batch_size=self.batch_size,
            shuffle=True,          # 随机打乱图片顺序
            num_workers=2
        )
        
        print(f"✅ 数据加载完成！")
        print(f"数据集大小: {len(dataset)}")
        print(f"类别: {dataset.classes}")
        
        # 设置AI的学习规则
        criterion = nn.CrossEntropyLoss()    # 判断AI预测得对不对
        optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)  # 告诉AI怎么改自己
        
        print(f"🚀 开始训练，共 {num_epochs} 轮...")
        
        for epoch in range(num_epochs):
            running_loss = 0.0
            correct = 0
            total = 0
            
            for i, (images, labels) in enumerate(data_loader):
                # 把图片和标签放到GPU上
                images, labels = images.to(self.device), labels.to(self.device)
                
                # 清空之前的记忆
                optimizer.zero_grad()
                
                # AI预测
                outputs = self.model(images)
                
                # 计算误差
                loss = criterion(outputs, labels)
                
                # AI根据误差调整自己
                loss.backward()
                optimizer.step()
                
                # 统计正确率
                running_loss += loss.item()
                _, predicted = torch.max(outputs.data, 1)
                total += labels.size(0)
                correct += (predicted == labels).sum().item()
                
                # 每100次显示一次进度
                if i % 100 == 99:
                    avg_loss = running_loss / 100
                    accuracy = 100 * correct / total
                    print(f'Epoch [{epoch+1}/{num_epochs}], Step [{i+1}], Loss: {avg_loss:.4f}, Accuracy: {accuracy:.2f}%')
                    running_loss = 0.0
            
            # 每轮结束后显示准确率
            epoch_acc = 100 * correct / total
            print(f'Epoch [{epoch+1}/{num_epochs}] 完成, 准确率: {epoch_acc:.2f}%')
        
        print("✅ 训练完成！")
        return dataset  # 返回数据集，用于预测时知道类别名

    def predict(self, image_path, dataset=None):
        """预测单张图片"""
        if not os.path.exists(image_path):
            print(f"❌ 图片不存在: {image_path}")
            return None
        
        try:
            # 加载图片
            img = Image.open(image_path)
            img_tensor = self.transform(img).unsqueeze(0).to(self.device)  # 预处理并添加批次维度
            
            # AI预测
            with torch.no_grad():  # 不需要计算梯度（测试时）
                self.model.eval()  # 设置为评估模式
                output = self.model(img_tensor)
                probabilities = torch.softmax(output, dim=1)[0]  # 转成概率
                predicted_class_idx = torch.argmax(output, 1).item()  # 找最大概率的类别
            
            # 获取类别名称
            class_names = dataset.classes if dataset else ['cat', 'dog']
            predicted_class = class_names[predicted_class_idx]
            confidence = probabilities[predicted_class_idx].item()
            
            return {
                'class': predicted_class,      # 预测的类别
                'confidence': confidence,      # 置信度
                'all_probabilities': {        # 所有类别的概率
                    class_names[i]: probabilities[i].item() 
                    for i in range(len(class_names))
                }
            }
            
        except Exception as e:
            print(f"❌ 预测出错: {str(e)}")
            return None

    def show_prediction(self, image_path, dataset=None):
        """显示预测结果和图片"""
        result = self.predict(image_path, dataset)
        if result is None:
            return
        
        # 显示图片
        img = Image.open(image_path)
        plt.figure(figsize=(8, 6))
        plt.imshow(img)
        plt.title(f"预测: {result['class']} (置信度: {result['confidence']:.3f})")
        plt.axis('off')
        plt.show()
        
        # 显示概率分布
        print(f"预测结果:")
        for class_name, prob in result['all_probabilities'].items():
            print(f"  {class_name}: {prob:.3f}")

    def save_model(self, path="cat_dog_classifier.pth"):
        """保存训练好的模型"""
        torch.save(self.model.state_dict(), path)
        print(f"✅ 模型已保存到: {path}")

    def load_model(self, path="cat_dog_classifier.pth"):
        """加载已训练的模型"""
        if not os.path.exists(path):
            print(f"❌ 模型文件不存在: {path}")
            return False
        
        self.model.load_state_dict(torch.load(path, map_location=self.device))
        print(f"✅ 模型已加载: {path}")
        return True

# 1. 创建分类器实例
classifier = CatDogClassifier(
    img_size=128,      # 图片大小
    batch_size=32      # 每次处理32张图片
)

# 2. 训练模型（需要你的数据集）
# 假设你的图片放在 /content/PetImages 文件夹下
# 文件结构应该是：
# PetImages/
# ├── Cat/
# │   ├── cat1.jpg
# │   └── cat2.jpg
# └── Dog/
#     ├── dog1.jpg
#     └── dog2.jpg

 dataset = classifier.train_model(
     data_dir="/content/PetImages",
     num_epochs=5,
     learning_rate=0.001
 )

# 3. 保存模型
 classifier.save_model("my_cat_dog_model.pth")

# 4. 预测图片
 result = classifier.predict("/content/PetImages/Cat/cat.1.jpg", dataset)
 print(f"预测结果: {result}")

# 5. 显示预测结果
 classifier.show_prediction("/content/PetImages/Cat/cat.1.jpg", dataset)