"""
baseline/VGG_KNN.py
VGG16 预训练模型特征提取 + KNN 四分类基线
可独立运行，也可被 compare_models.py 导入复用
author: yukun-hh
date: 2026-5-14
"""
import sys, os
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import numpy as np
import matplotlib.pyplot as plt
import matplotlib

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import models, transforms
from tqdm import tqdm

from sklearn.neighbors import KNeighborsClassifier
from sklearn.metrics import (
    accuracy_score, f1_score,
    confusion_matrix, ConfusionMatrixDisplay,
    classification_report,
)

from Dataloader import RobustImageFolder

matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans']
matplotlib.rcParams['axes.unicode_minus'] = False


CLASS_NAMES = ['厨余垃圾', '可回收物', '其他垃圾', '有害垃圾']


def load_vgg16_extractor(device):
    try:
        model = models.vgg16(weights='IMAGENET1K_V1')
    except TypeError:
        model = models.vgg16(pretrained=True)
    model.classifier = nn.Identity()
    model = model.to(device).eval()
    for param in model.parameters():
        param.requires_grad = False
    return model


def extract_features(model, loader, device):
    model.eval()
    all_features = []
    all_labels = []
    with torch.no_grad():
        for images, labels in tqdm(loader, desc='Extracting features'):
            images = images.to(device)
            feats = model(images)
            all_features.append(feats.cpu().numpy())
            all_labels.append(labels.numpy())
    return np.concatenate(all_features), np.concatenate(all_labels)


class VGGKNNBaseline:
    def __init__(self, k=5, device='cpu',
                 data_root='../trash_division_data/ultimate_4_class/',
                 image_size=256, batch_size=32, num_workers=4):
        self.k = k
        self.device = device
        self.data_root = data_root
        self.image_size = image_size
        self.batch_size = batch_size
        self.num_workers = num_workers
        self.extractor = load_vgg16_extractor(device)
        self.knn = KNeighborsClassifier(n_neighbors=k, n_jobs=-1)

    def _get_loader(self, split):
        transform = transforms.Compose([
            transforms.Resize((self.image_size, self.image_size)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
        ])
        dataset = RobustImageFolder(
            root=os.path.join(self.data_root, split),
            transform=transform,
        )
        print(f"  {split}: {len(dataset)} 张")
        return DataLoader(dataset, batch_size=self.batch_size,
                          shuffle=False, num_workers=self.num_workers,
                          pin_memory=True, drop_last=False)

    def fit(self, train_loader=None):
        if train_loader is None:
            train_loader = self._get_loader('train')
        print("  提取训练集特征 ...")
        train_feats, train_labels = extract_features(self.extractor, train_loader, self.device)
        self.knn.fit(train_feats, train_labels)

    def predict(self, val_loader=None):
        if val_loader is None:
            val_loader = self._get_loader('val')
        print("  提取验证集特征 ...")
        val_feats, val_labels = extract_features(self.extractor, val_loader, self.device)
        preds = self.knn.predict(val_feats)
        probs = self.knn.predict_proba(val_feats)
        return val_labels, preds, probs


if __name__ == '__main__':
    DATA_ROOT   = '../trash_division_data/ultimate_4_class/'
    BATCH_SIZE  = 32
    IMAGE_SIZE  = 256
    NUM_WORKERS = 4
    K           = 5

    device = torch.device('cuda' if torch.cuda.is_available()
                          else 'xpu' if hasattr(torch, 'xpu') and torch.xpu.is_available()
                          else 'cpu')
    print(f"Device: {device}")

    baseline = VGGKNNBaseline(k=K, device=device,
                               data_root=DATA_ROOT, image_size=IMAGE_SIZE,
                               batch_size=BATCH_SIZE, num_workers=NUM_WORKERS)

    train_loader = baseline._get_loader('train')
    val_loader = baseline._get_loader('val')

    baseline.fit(train_loader)
    y_true, y_preds, y_probs = baseline.predict(val_loader)

    acc = accuracy_score(y_true, y_preds)
    macro_f1 = f1_score(y_true, y_preds, average='macro')
    print(f"\n验证集 Accuracy: {acc:.4f}")
    print(f"验证集 Macro-F1: {macro_f1:.4f}")
    print(f"\n分类报告:\n{classification_report(y_true, y_preds, target_names=CLASS_NAMES)}")

    cm = confusion_matrix(y_true, y_preds)
    fig, ax = plt.subplots(figsize=(8, 7))
    ConfusionMatrixDisplay(cm, display_labels=CLASS_NAMES).plot(
        ax=ax, cmap='Blues', values_format='d', xticks_rotation=30)
    ax.set_title(f'Baseline Confusion Matrix (VGG16 + KNN, K={K})', fontsize=14)
    plt.tight_layout()
    out_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'vgg_knn_confusion_matrix.png')
    plt.savefig(out_path, dpi=150, bbox_inches='tight')
    plt.show()
    print(f"混淆矩阵已保存: {out_path}")