"""Evaluate the trained model on the test set and dump metrics.json.

Run after `train.py` has produced `model/saved/brain_tumor_model.pth`.

Outputs:
    model/saved/metrics.json       - test accuracy, per-class precision/recall/F1, confusion matrix
    model/saved/confusion_matrix.png - heatmap visualization (if matplotlib available)
"""

from __future__ import annotations

import json
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from sklearn.metrics import (
    classification_report,
    confusion_matrix,
)
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import ImageFolder

from architecture import CLASS_NAMES, IMG_SIZE, build_model

PROJECT_ROOT = Path(__file__).resolve().parent.parent
TEST_DIR = PROJECT_ROOT / "data" / "raw" / "Testing"
SAVE_DIR = PROJECT_ROOT / "model" / "saved"
CHECKPOINT = SAVE_DIR / "brain_tumor_model.pth"
METRICS_PATH = SAVE_DIR / "metrics.json"
CM_PLOT_PATH = SAVE_DIR / "confusion_matrix.png"

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
IMAGENET_MEAN = [0.485, 0.456, 0.406]
IMAGENET_STD = [0.229, 0.224, 0.225]


def _save_confusion_plot(cm: np.ndarray, labels: list[str]) -> None:
    try:
        import matplotlib.pyplot as plt
    except ImportError:
        return
    fig, ax = plt.subplots(figsize=(6, 5))
    im = ax.imshow(cm, cmap="Blues")
    ax.set_xticks(range(len(labels)))
    ax.set_yticks(range(len(labels)))
    ax.set_xticklabels(labels, rotation=45, ha="right")
    ax.set_yticklabels(labels)
    ax.set_xlabel("Predicted")
    ax.set_ylabel("True")
    ax.set_title("Confusion Matrix (test set)")
    for i in range(cm.shape[0]):
        for j in range(cm.shape[1]):
            ax.text(j, i, str(cm[i, j]), ha="center", va="center", color="black")
    fig.colorbar(im, ax=ax)
    fig.tight_layout()
    fig.savefig(CM_PLOT_PATH, dpi=150)
    plt.close(fig)
    print(f"Confusion matrix plot saved to {CM_PLOT_PATH}")


def main() -> None:
    if not CHECKPOINT.exists():
        raise SystemExit(f"Checkpoint not found at {CHECKPOINT}. Run train.py first.")

    # Match the eval transform used during training: Resize(IMG_SIZE+24) → CenterCrop(IMG_SIZE).
    eval_transform = transforms.Compose(
        [
            transforms.Resize((IMG_SIZE + 24, IMG_SIZE + 24)),
            transforms.CenterCrop(IMG_SIZE),
            transforms.ToTensor(),
            transforms.Normalize(IMAGENET_MEAN, IMAGENET_STD),
        ]
    )
    test_ds = ImageFolder(TEST_DIR, transform=eval_transform)
    if test_ds.classes != CLASS_NAMES:
        raise SystemExit(
            f"Class order mismatch: {test_ds.classes} vs expected {CLASS_NAMES}"
        )
    loader = DataLoader(
        test_ds, batch_size=32, shuffle=False, num_workers=4, pin_memory=True
    )

    print(f"Device: {DEVICE} | test set: {len(test_ds)} images")
    model = build_model(pretrained=False).to(DEVICE)
    ckpt = torch.load(CHECKPOINT, map_location=DEVICE)
    model.load_state_dict(ckpt["model_state_dict"])
    model.eval()

    all_preds: list[int] = []
    all_targets: list[int] = []
    # Match training criterion so the reported loss is comparable to train_v2.log.
    criterion = nn.CrossEntropyLoss(label_smoothing=0.05)
    loss_sum, total = 0.0, 0
    with torch.no_grad():
        for inputs, targets in loader:
            inputs = inputs.to(DEVICE, non_blocking=True)
            targets = targets.to(DEVICE, non_blocking=True)
            logits = model(inputs)
            loss_sum += criterion(logits, targets).item() * inputs.size(0)
            total += inputs.size(0)
            preds = logits.argmax(dim=1).cpu().numpy()
            all_preds.extend(preds.tolist())
            all_targets.extend(targets.cpu().numpy().tolist())

    acc = float(np.mean(np.array(all_preds) == np.array(all_targets)))
    test_loss = loss_sum / total
    cm = confusion_matrix(all_targets, all_preds).tolist()
    report = classification_report(
        all_targets,
        all_preds,
        target_names=CLASS_NAMES,
        digits=4,
        output_dict=True,
    )

    metrics = dict(
        test_accuracy=acc,
        test_loss=test_loss,
        num_samples=total,
        class_names=CLASS_NAMES,
        confusion_matrix=cm,
        classification_report=report,
        checkpoint=str(CHECKPOINT.relative_to(PROJECT_ROOT)),
    )
    METRICS_PATH.write_text(json.dumps(metrics, indent=2))
    print(f"Test accuracy: {acc*100:.2f}% | loss: {test_loss:.4f}")
    print(f"Per-class F1: " + ", ".join(
        f"{c}={report[c]['f1-score']:.3f}" for c in CLASS_NAMES
    ))
    print(f"Confusion matrix:\n{np.array(cm)}")
    print(f"Metrics written to {METRICS_PATH}")
    _save_confusion_plot(np.array(cm), CLASS_NAMES)


if __name__ == "__main__":
    main()