#迁移学习 (Transfer Learning)：利用预训练模型快速构建高性能模型

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torchvision.models as models
from torchvision.models import (
    ResNet50_Weights,
    VGG16_Weights,
    MobileNet_V3_Small_Weights
)

def build_transfer_model(
    model_name: str = "resnet50",
    num_classes: int = 10,
    strategy: str = "feature_extract",  # 可选: feature_extract, fine_tune, full_tune
    freeze_layers: int = 4  # 仅fine_tune有效，ResNet50有children()[:8]是骨干网络
):
    """
    构建迁移学习模型
    """
    # 1. 加载带预训练权重的模型
    if model_name == "resnet50":
        weights = ResNet50_Weights.IMAGENET1K_V1
        model = models.resnet50(weights=weights)
        in_features = model.fc.in_features
        head = "fc"
    elif model_name == "vgg16":
        weights = VGG16_Weights.IMAGENET1K_V1
        model = models.vgg16(weights=weights)
        in_features = model.classifier[6].in_features
        head = "classifier"
    elif model_name == "mobilenet_v3_small":
        weights = MobileNet_V3_Small_Weights.IMAGENET1K_V1
        model = models.mobilenet_v3_small(weights=weights)
        in_features = model.classifier[3].in_features
        head = "classifier"
    else:
        raise ValueError(f"不支持的模型: {model_name}")

    # 2. 根据策略冻结参数
    if strategy == "feature_extract":
        for param in model.parameters():
            param.requires_grad = False
    elif strategy == "fine_tune":
        children = list(model.children()) if hasattr(model, "children") else []
        for child in children[:freeze_layers]:
            for param in child.parameters():
                param.requires_grad = False
    # full_tune不冻结任何参数

    # 3. 替换任务特定的分类头
    if model_name == "resnet50":
        model.fc = nn.Linear(in_features, num_classes)
    elif model_name == "vgg16":
        model.classifier[6] = nn.Linear(in_features, num_classes)
    elif model_name == "mobilenet_v3_small":
        model.classifier[3] = nn.Linear(in_features, num_classes)

    return model, weights.transforms()  # 直接获取预训练模型要求的预处理！

def get_data_loaders(
    dataset_path: str = "./data",
    preprocess: transforms.Compose = None,
    batch_size: int = 32,
    train_augment: bool = True
):
    """
    获取训练/验证数据加载器
    """
    # 训练集：用预训练预处理 + 数据增强（可选）
    train_transforms = [preprocess]
    if train_augment:
        train_transforms.insert(0, transforms.RandomHorizontalFlip(p=0.5))
        train_transforms.insert(0, transforms.Resize(256))
        train_transforms.insert(1, transforms.RandomCrop(224))
    train_transform = transforms.Compose(train_transforms)

    # 验证/测试集：只用预训练预处理
    val_transform = preprocess

    # 加载ImageFolder格式的数据集（最常用：根目录下是class文件夹）
    train_dataset = datasets.ImageFolder(f"{dataset_path}/train", transform=train_transform)
    val_dataset = datasets.ImageFolder(f"{dataset_path}/val", transform=val_transform)

    # 构建DataLoader
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    return train_loader, val_loader, train_dataset.classes

def train_model(
    model: nn.Module,
    train_loader: DataLoader,
    val_loader: DataLoader,
    num_epochs: int = 10,
    lr: float = 0.001,
    device: str = "cuda" if torch.cuda.is_available() else "cpu"
):
    """
    训练迁移学习模型
    """
    model.to(device)
    criterion = nn.CrossEntropyLoss()
    # 仅优化可训练的参数！
    optimizer = optim.Adam(filter(lambda p: p.requires_grad, model.parameters()), lr=lr)
    # 简单的学习率调度（可选）
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='max', patience=3, factor=0.5)

    best_val_acc = 0.0
    best_model_state = model.state_dict()

    for epoch in range(num_epochs):
        # 训练阶段
        model.train()
        train_loss, train_correct, train_total = 0.0, 0, 0
        for inputs, labels in train_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            optimizer.zero_grad()
            outputs = model(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            train_loss += loss.item() * inputs.size(0)
            _, preds = outputs.max(1)
            train_total += labels.size(0)
            train_correct += preds.eq(labels).sum().item()

        train_loss /= train_total
        train_acc = train_correct / train_total

        # 验证阶段
        model.eval()
        val_loss, val_correct, val_total = 0.0, 0, 0
        with torch.no_grad():
            for inputs, labels in val_loader:
                inputs, labels = inputs.to(device), labels.to(device)
                outputs = model(inputs)
                loss = criterion(outputs, labels)

                val_loss += loss.item() * inputs.size(0)
                _, preds = outputs.max(1)
                val_total += labels.size(0)
                val_correct += preds.eq(labels).sum().item()

        val_loss /= val_total
        val_acc = val_correct / val_total
        scheduler.step(val_acc)

        # 更新最佳模型
        if val_acc > best_val_acc:
            best_val_acc = val_acc
            best_model_state = model.state_dict().copy()

        # 打印日志
        print(f"Epoch {epoch+1}/{num_epochs} | "
              f"Train Loss: {train_loss:.4f} Acc: {train_acc:.4f} | "
              f"Val Loss: {val_loss:.4f} Acc: {val_acc:.4f}")

    print(f"\n训练完成！最佳验证准确率: {best_val_acc:.4f}")
    model.load_state_dict(best_model_state)
    return model

if __name__ == "__main__":
    # 1. 构建模型（以特征提取为例）
    model, preprocess = build_transfer_model(
        model_name="resnet50",
        num_classes=5,  # 假设你的私有数据集有5类
        strategy="feature_extract"
    )

    # 2. 加载数据（确保你的数据集是ImageFolder格式：./data/train/类1, ./data/val/类1...）
    train_loader, val_loader, classes = get_data_loaders(
        dataset_path="./data",
        preprocess=preprocess,
        batch_size=32
    )
    print(f"类别列表: {classes}")

    # 3. 训练模型
    trained_model = train_model(
        model=model,
        train_loader=train_loader,
        val_loader=val_loader,
        num_epochs=10,
        lr=0.001
    )

    # 4. 保存最佳模型
    torch.save(trained_model.state_dict(), "best_transfer_model.pth")