from pathlib import Path
from urllib.request import urlretrieve
from zipfile import ZipFile

project_root = Path(r"C:\Users\81052\Documents\AI & GIS")

day5_raw_dir = project_root / "data" / "raw" / "day5"
zip_path = day5_raw_dir / "hymenoptera_data.zip"
data_dir = day5_raw_dir / "hymenoptera_data"

source_url = (
    "https://download.pytorch.org/tutorial/"
    "hymenoptera_data.zip"
)

day5_raw_dir.mkdir(parents=True, exist_ok=True)



if not data_dir.exists():
    print("正在解压……")
    with ZipFile(zip_path, mode="r") as archive:
        archive.extractall(day5_raw_dir)
else:
    print("数据目录已存在，不重复解压。")

image_extensions = {".jpg", ".jpeg", ".png"}

for split in ["train", "val"]:
    print(f"\n{split}:")
    
    for class_name in ["ants", "bees"]:
        class_dir = data_dir / split / class_name
        
        image_count = sum(
            1
            for path in class_dir.iterdir()
            if path.suffix.lower() in image_extensions
        )
        
        print(f"  {class_name}: {image_count} images")

assert data_dir.exists()
assert (data_dir / "train" / "ants").exists()
assert (data_dir / "train" / "bees").exists()
assert (data_dir / "val" / "ants").exists()
assert (data_dir / "val" / "bees").exists()

print("\nDay 5 数据目录检查通过：", data_dir)

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.models import ResNet18_Weights

# 明确指定固定版本，避免 DEFAULT 将来指向其他权重。
weights = ResNet18_Weights.IMAGENET1K_V1

# validation 使用权重自带的确定性预处理。
val_transform = weights.transforms()

# train 保持相同尺寸、插值和归一化参数，
# 但加入随机裁剪与水平翻转作为数据增强。
train_transform = transforms.Compose([
    transforms.RandomResizedCrop(
        size=val_transform.crop_size,
        interpolation=val_transform.interpolation,
        antialias=True
    ),
    transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=val_transform.mean,
        std=val_transform.std
    )
])

print("与权重绑定的 validation transform:")
print(val_transform)

train_dataset = datasets.ImageFolder(
    data_dir / "train",
    transform=train_transform
)

val_dataset = datasets.ImageFolder(
    data_dir / "val",
    transform=val_transform
)

assert train_dataset.class_to_idx == val_dataset.class_to_idx

print("\n类别名称:", train_dataset.classes)
print("类别映射:", train_dataset.class_to_idx)
print("train samples:", len(train_dataset))
print("validation samples:", len(val_dataset))

# ---------- 检查单条 sample ----------
sample_image, sample_label = train_dataset[0]

print("\n单条 sample:")
print("image shape:", sample_image.shape)
print("image dtype:", sample_image.dtype)
print("label:", sample_label, type(sample_label))
print(
    "归一化后数值范围:",
    sample_image.min().item(),
    "到",
    sample_image.max().item()
)

# ---------- 建立 DataLoader ----------
batch_size = 16

train_generator = torch.Generator().manual_seed(42)

train_loader = DataLoader(
    train_dataset,
    batch_size=batch_size,
    shuffle=True,
    num_workers=0,
    generator=train_generator
)

val_loader = DataLoader(
    val_dataset,
    batch_size=batch_size,
    shuffle=False,
    num_workers=0
)

# ---------- 检查一个 batch ----------
batch_images, batch_labels = next(iter(train_loader))

print("\n一个 batch:")
print("images shape:", batch_images.shape)
print("images dtype:", batch_images.dtype)
print("labels shape:", batch_labels.shape)
print("labels dtype:", batch_labels.dtype)
print("前 8 个 labels:", batch_labels[:8])

assert sample_image.shape == (3, 224, 224)
assert batch_images.shape[1:] == (3, 224, 224)
assert batch_labels.ndim == 1

print("\nDay 5 sample/batch shape 检查通过。")

from torch import nn
from torchvision.models import resnet18

# GPU 可用就使用 GPU，否则退回 CPU。
device = torch.device(
    "cuda" if torch.cuda.is_available() else "cpu"
)

# 1. 创建 ResNet18，并加载指定的 ImageNet 预训练权重。
# 第一次运行时可能下载约 45 MB 的权重到 PyTorch 缓存。
model = resnet18(weights=weights)

print("原始分类层:")
print(model.fc)

# 原分类层接收的特征数量。
backbone_feature_count = model.fc.in_features
class_count = len(train_dataset.classes)

print("\nbackbone 输出特征数:", backbone_feature_count)
print("当前任务类别数:", class_count)

# 2. 冻结目前模型中的所有参数。
for parameter in model.parameters():
    parameter.requires_grad = False

# 3. 创建新的分类头。
# 在固定 seed 后创建，使新分类头的随机初始值可复现。
torch.manual_seed(42)

model.fc = nn.Linear(
    in_features=backbone_feature_count,
    out_features=class_count
)

# 新建的 Linear 参数默认 requires_grad=True。
model = model.to(device)

print("\n替换后的分类层:")
print(model.fc)
print("计算设备:", device)

# ---------- 检查哪些参数可以训练 ----------
total_parameter_count = sum(
    parameter.numel()
    for parameter in model.parameters()
)

trainable_parameter_count = sum(
    parameter.numel()
    for parameter in model.parameters()
    if parameter.requires_grad
)

frozen_parameter_count = (
    total_parameter_count - trainable_parameter_count
)

print("\n全部参数:", total_parameter_count)
print("冻结参数:", frozen_parameter_count)
print("可训练参数:", trainable_parameter_count)

print(
    "第一层卷积是否可训练:",
    model.conv1.weight.requires_grad
)

print(
    "新分类层 weight 是否可训练:",
    model.fc.weight.requires_grad
)

print(
    "新分类层 bias 是否可训练:",
    model.fc.bias.requires_grad
)

# ---------- 用 4 张图片验证 forward shape ----------
model.eval()

example_images = batch_images[:4].to(device)

with torch.no_grad():
    example_logits = model(example_images)

print("\nforward shape 检查:")
print("模型输入:", example_images.shape)
print("模型输出 logits:", example_logits.shape)
print("前两张影像的 logits:\n", example_logits[:2])

assert example_images.shape == (4, 3, 224, 224)
assert example_logits.shape == (4, 2)
assert trainable_parameter_count == 512 * 2 + 2
assert model.conv1.weight.requires_grad is False
assert model.fc.weight.requires_grad is True

print("\nResNet18 冻结与分类层替换检查通过。")

from pathlib import Path

# ---------- 训练配置 ----------
learning_rate = 0.001
momentum = 0.9
num_epochs = 12

loss_fn = nn.CrossEntropyLoss()

# 关键：optimizer 只管理新的分类头。
optimizer = torch.optim.SGD(
    model.fc.parameters(),
    lr=learning_rate,
    momentum=momentum
)

history = {
    "train_loss": [],
    "train_accuracy": [],
    "val_loss": [],
    "val_accuracy": []
}

best_val_loss = float("inf")

checkpoint_path = (
    project_root
    / "outputs"
    / "day5"
    / "best_resnet18_head.pt"
)

checkpoint_path.parent.mkdir(
    parents=True,
    exist_ok=True
)

print("训练设备:", device)
print("optimizer 管理的参数量:", sum(
    parameter.numel()
    for group in optimizer.param_groups
    for parameter in group["params"]
))
print("checkpoint:", checkpoint_path)

# ---------- 正式训练 ----------
for epoch in range(num_epochs):
    # ===== Training =====
    # Backbone 保持 eval，避免 BatchNorm 统计值变化。
    # 分类头保持 train，并且梯度仍然开启。
    model.eval()
    model.fc.train()

    train_loss_sum = 0.0
    train_correct = 0
    train_sample_count = 0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        optimizer.zero_grad()

        logits = model(images)
        loss = loss_fn(logits, labels)

        loss.backward()
        optimizer.step()

        batch_sample_count = labels.size(0)

        train_loss_sum += (
            loss.item() * batch_sample_count
        )

        train_correct += (
            logits.argmax(dim=1) == labels
        ).sum().item()

        train_sample_count += batch_sample_count

    train_loss = (
        train_loss_sum / train_sample_count
    )

    train_accuracy = (
        train_correct / train_sample_count
    )

    # ===== Validation =====
    model.eval()

    val_loss_sum = 0.0
    val_correct = 0
    val_sample_count = 0

    with torch.no_grad():
        for images, labels in val_loader:
            images = images.to(device)
            labels = labels.to(device)

            logits = model(images)
            loss = loss_fn(logits, labels)

            batch_sample_count = labels.size(0)

            val_loss_sum += (
                loss.item() * batch_sample_count
            )

            val_correct += (
                logits.argmax(dim=1) == labels
            ).sum().item()

            val_sample_count += batch_sample_count

    val_loss = val_loss_sum / val_sample_count
    val_accuracy = val_correct / val_sample_count

    history["train_loss"].append(train_loss)
    history["train_accuracy"].append(train_accuracy)
    history["val_loss"].append(val_loss)
    history["val_accuracy"].append(val_accuracy)

    # validation loss 刷新最低纪录时保存。
    if val_loss < best_val_loss:
        best_val_loss = val_loss

        torch.save(
            {
                "epoch": epoch + 1,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "validation_loss": val_loss,
                "validation_accuracy": val_accuracy,
                "class_to_idx": train_dataset.class_to_idx,
                "weights_name": weights.name,
                "backbone_frozen": True,
                "learning_rate": learning_rate,
                "momentum": momentum,
                "batch_size": batch_size
            },
            checkpoint_path
        )

        checkpoint_status = "已保存新最佳模型"
    else:
        checkpoint_status = ""

    print(
        f"Epoch {epoch + 1:02d}/{num_epochs} | "
        f"train loss {train_loss:.4f}, "
        f"acc {train_accuracy:.4f} | "
        f"val loss {val_loss:.4f}, "
        f"acc {val_accuracy:.4f} "
        f"{checkpoint_status}"
    )

print("\n最低 validation loss:", best_val_loss)
print("最佳 checkpoint:", checkpoint_path)

# ---------- 读取 checkpoint ----------
checkpoint = torch.load(
    checkpoint_path,
    map_location=device,
    weights_only=True
)

print("checkpoint 字段:")
print(checkpoint.keys())

print("\n保存的 epoch:", checkpoint["epoch"])
print("保存的 validation loss:", checkpoint["validation_loss"])
print("保存的 validation accuracy:", checkpoint["validation_accuracy"])
print("类别映射:", checkpoint["class_to_idx"])
print("预训练权重版本:", checkpoint["weights_name"])
print("backbone frozen:", checkpoint["backbone_frozen"])

# ---------- 重建相同的模型结构 ----------
# weights=None：只创建结构，不重新下载或加载预训练权重。
# checkpoint 的 state_dict 随后会覆盖全部模型参数。
restored_model = resnet18(weights=None)

restored_feature_count = (
    restored_model.fc.in_features
)

for parameter in restored_model.parameters():
    parameter.requires_grad = False

restored_model.fc = nn.Linear(
    in_features=restored_feature_count,
    out_features=len(checkpoint["class_to_idx"])
)

load_result = restored_model.load_state_dict(
    checkpoint["model_state_dict"]
)

restored_model = restored_model.to(device)
restored_model.eval()

print("\n参数加载结果:", load_result)

# ---------- 重新计算全部 validation ----------
restored_val_loss_sum = 0.0
restored_val_correct = 0
restored_val_sample_count = 0

first_batch_predictions = None
first_batch_labels = None

with torch.no_grad():
    for batch_index, (images, labels) in enumerate(
        val_loader
    ):
        images = images.to(device)
        labels = labels.to(device)

        logits = restored_model(images)
        loss = loss_fn(logits, labels)

        batch_sample_count = labels.size(0)

        restored_val_loss_sum += (
            loss.item() * batch_sample_count
        )

        predictions = logits.argmax(dim=1)

        restored_val_correct += (
            predictions == labels
        ).sum().item()

        restored_val_sample_count += (
            batch_sample_count
        )

        if batch_index == 0:
            first_batch_predictions = (
                predictions.cpu()
            )
            first_batch_labels = labels.cpu()

restored_val_loss = (
    restored_val_loss_sum
    / restored_val_sample_count
)

restored_val_accuracy = (
    restored_val_correct
    / restored_val_sample_count
)

print("\n恢复后 validation:")
print("sample count:", restored_val_sample_count)
print("correct:", restored_val_correct)
print("loss:", restored_val_loss)
print("accuracy:", restored_val_accuracy)

# ---------- 把数字标签翻译成类别名 ----------
idx_to_class = {
    class_index: class_name
    for class_name, class_index
    in checkpoint["class_to_idx"].items()
}

print("\n第一个 validation batch 的前 8 项:")

for index in range(8):
    predicted_index = (
        first_batch_predictions[index].item()
    )

    true_index = first_batch_labels[index].item()

    print(
        f"{index}: "
        f"predicted={idx_to_class[predicted_index]}, "
        f"true={idx_to_class[true_index]}"
    )

# ---------- 验证恢复结果 ----------
assert load_result.missing_keys == []
assert load_result.unexpected_keys == []

assert abs(
    restored_val_loss
    - checkpoint["validation_loss"]
) < 1e-6

assert abs(
    restored_val_accuracy
    - checkpoint["validation_accuracy"]
) < 1e-12

print("\n最佳 checkpoint 恢复验证通过。")

import matplotlib.pyplot as plt

assert len(history["train_loss"]) == num_epochs
assert len(history["val_loss"]) == num_epochs

epochs = list(range(1, num_epochs + 1))

best_epoch = checkpoint["epoch"]
best_index = best_epoch - 1

curve_path = (
    checkpoint_path.parent
    / "training_curves.png"
)

fig, axes = plt.subplots(
    1,
    2,
    figsize=(13, 4.5)
)

# ---------- Loss ----------
axes[0].plot(
    epochs,
    history["train_loss"],
    marker="o",
    label="train"
)

axes[0].plot(
    epochs,
    history["val_loss"],
    marker="o",
    label="validation"
)

axes[0].axvline(
    best_epoch,
    color="red",
    linestyle="--",
    alpha=0.7,
    label=f"best checkpoint: epoch {best_epoch}"
)

axes[0].scatter(
    best_epoch,
    history["val_loss"][best_index],
    color="red",
    s=70,
    zorder=5
)

axes[0].set_title("ResNet18 Transfer Learning: Loss")
axes[0].set_xlabel("Epoch")
axes[0].set_ylabel("Sample-weighted loss")
axes[0].set_xticks(epochs)
axes[0].grid(alpha=0.3)
axes[0].legend()

# ---------- Accuracy ----------
axes[1].plot(
    epochs,
    history["train_accuracy"],
    marker="o",
    label="train"
)

axes[1].plot(
    epochs,
    history["val_accuracy"],
    marker="o",
    label="validation"
)

axes[1].axvline(
    best_epoch,
    color="red",
    linestyle="--",
    alpha=0.7,
    label=f"selected epoch {best_epoch}"
)

axes[1].scatter(
    best_epoch,
    history["val_accuracy"][best_index],
    color="red",
    s=70,
    zorder=5
)

all_accuracies = (
    history["train_accuracy"]
    + history["val_accuracy"]
)

accuracy_lower_limit = max(
    0.0,
    min(all_accuracies) - 0.05
)

axes[1].set_title("ResNet18 Transfer Learning: Accuracy")
axes[1].set_xlabel("Epoch")
axes[1].set_ylabel("Accuracy")
axes[1].set_xticks(epochs)
axes[1].set_ylim(accuracy_lower_limit, 1.01)
axes[1].grid(alpha=0.3)
axes[1].legend()

fig.suptitle(
    "Frozen ResNet18 Backbone + Trainable 2-Class Head",
    fontsize=14
)

plt.tight_layout()

fig.savefig(
    curve_path,
    dpi=160,
    bbox_inches="tight"
)

plt.show()

print("训练曲线已保存：", curve_path)