import argparse
import os
import torch
import torch.nn as nn
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
from sklearn.metrics import classification_report, confusion_matrix


def build_model(arch):
    if arch == "mobilenet":
        model = models.mobilenet_v2()
        num_ftrs = model.classifier[1].in_features
        model.classifier[1] = nn.Linear(num_ftrs, 2)
    elif arch in ("swin_t", "swin_t_finetune"):
        model = models.swin_t()
        num_ftrs = model.head.in_features
        model.head = nn.Linear(num_ftrs, 2)
    elif arch == "swin_s":
        model = models.swin_s()
        num_ftrs = model.head.in_features
        model.head = nn.Linear(num_ftrs, 2)
    else:
        raise ValueError(f"Unknown arch: {arch}")
    return model


def evaluate(data_dir, model_path, arch):
    data_transforms = transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    test_dataset = datasets.ImageFolder(os.path.join(data_dir, 'test'), data_transforms)
    test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    model = build_model(arch)
    model.load_state_dict(torch.load(model_path, map_location=device))
    model = model.to(device)
    model.eval()

    all_preds = []
    all_labels = []

    with torch.no_grad():
        for inputs, labels in test_loader:
            inputs, labels = inputs.to(device), labels.to(device)
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            all_preds.extend(preds.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    print("\nClassification Report:")
    print(classification_report(all_labels, all_preds, target_names=test_dataset.classes))
    print("\nConfusion Matrix:")
    print(confusion_matrix(all_labels, all_preds))


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--arch", choices=["mobilenet", "swin_t", "swin_t_finetune", "swin_s"], required=True)
    parser.add_argument("--data-dir", default="./data/split")
    parser.add_argument("--weights", required=True, help="Path to model weights file")
    args = parser.parse_args()

    evaluate(args.data_dir, args.weights, args.arch)