import os
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
from tqdm import tqdm  # 引入了实时进度条神器


def train_model():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"当前使用的计算设备: {device}")

    if device.type == 'cpu':
        print("⚠️ 警告：当前正在使用 CPU 训练，4000张图片预计每轮需要 15-30 分钟，请保持耐心！")

    data_transforms = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    data_dir = './data'
    image_dataset = datasets.ImageFolder(data_dir, data_transforms)
    # CPU 训练比较慢，我们把 batch_size 稍微调大一点点到 16
    dataloader = DataLoader(image_dataset, batch_size=16, shuffle=True)

    print(f"总计训练图片数量: {len(image_dataset)} 张\n")

    print("正在加载 MobileNetV2 模型...")
    model = models.mobilenet_v2(weights=models.MobileNet_V2_Weights.IMAGENET1K_V1)

    num_ftrs = model.classifier[1].in_features
    model.classifier[1] = nn.Linear(num_ftrs, 2)
    model = model.to(device)

    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=1e-4)

    # 为了用 CPU 能快点看到结果，我们先设为 3 轮
    num_epochs = 3
    print("\n--- 开始模型微调 ---")

    for epoch in range(num_epochs):
        model.train()
        running_loss = 0.0
        corrects = 0

        # 【核心修改】：用 tqdm 包装 dataloader，生成实时进度条
        progress_bar = tqdm(dataloader, desc=f"第 {epoch + 1}/{num_epochs} 轮", leave=False, colour='green')

        for inputs, labels in progress_bar:
            inputs = inputs.to(device)
            labels = labels.to(device)

            optimizer.zero_grad()
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            loss = criterion(outputs, labels)

            loss.backward()
            optimizer.step()

            running_loss += loss.item() * inputs.size(0)
            corrects += torch.sum(preds == labels.data)

            # 让进度条实时显示当前的误差值
            progress_bar.set_postfix({'loss': f"{loss.item():.4f}"})

        epoch_loss = running_loss / len(image_dataset)
        epoch_acc = corrects.double() / len(image_dataset)
        print(f"✅ 第 {epoch + 1}/{num_epochs} 轮完成 | 平均损失: {epoch_loss:.4f} | 准确率: {epoch_acc:.4f}")

    save_path = 'mobilenet_finetuned.pth'
    torch.save(model.state_dict(), save_path)
    print(f"\n🎉 训练完成！模型权重已保存至: {save_path}")


if __name__ == '__main__':
    train_model()