import torch
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, classification_report

from model import build_model
from dataloader import get_dataloaders
from utils import get_device

CSV_PATH = "data_processed/metadata_final.csv"
IMG_DIR = "data_processed/images"
CHECKPOINT_PATH = "checkpoints/best_model.pth"

device = get_device()

df = pd.read_csv(CSV_PATH)
num_classes = df["label_id"].nunique()

model = build_model(num_classes, device)
model.load_state_dict(torch.load(CHECKPOINT_PATH))
model.eval()

_, val_loader = get_dataloaders(
    csv_path=CSV_PATH,
    images_dir=IMG_DIR,
    batch_size=32
)

y_true, y_pred = [], []

with torch.no_grad():
    for images, labels in val_loader:
        images = images.to(device)
        outputs = model(images)
        preds = outputs.argmax(dim=1).cpu().numpy()

        y_pred.extend(preds)
        y_true.extend(labels.numpy())

cm = confusion_matrix(y_true, y_pred)

plt.figure(figsize=(14, 12))
sns.heatmap(cm, cmap="Blues", xticklabels=False, yticklabels=False)
plt.title("Confusion Matrix")
plt.xlabel("Predicted")
plt.ylabel("True")
plt.show()

print("\nClassification Report:")
print(classification_report(y_true, y_pred))