"""Load saved MLP checkpoint and print test accuracy, F1, AUC."""
import os
import sys

import numpy as np
import torch
from sklearn.metrics import f1_score, roc_auc_score

REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
if REPO_ROOT not in sys.path:
    sys.path.insert(0, REPO_ROOT)

from data_preparation.prepare_dataset import get_dataloaders
from models.mlp.train import BaseModel

CKPT_PATH = os.path.join(REPO_ROOT, "checkpoints", "mlp_best.pt")


def main():
    if not os.path.isfile(CKPT_PATH):
        print(f"No checkpoint at {CKPT_PATH}. Train first: python -m models.mlp.train")
        return

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    train_loader, val_loader, test_loader, num_features, num_classes, _ = get_dataloaders(
        model_name="face_orientation",
        batch_size=32,
        split_ratios=(0.7, 0.15, 0.15),
        seed=42,
    )

    model = BaseModel(num_features, num_classes).to(device)
    model.load_state_dict(torch.load(CKPT_PATH, map_location=device, weights_only=True))
    model.eval()

    criterion = torch.nn.CrossEntropyLoss()
    test_loss, test_acc, test_probs, test_preds, test_labels = model.test_step(
        test_loader, criterion, device
    )

    f1 = float(f1_score(test_labels, test_preds, average="weighted"))
    if num_classes > 2:
        auc = float(roc_auc_score(test_labels, test_probs, multi_class="ovr", average="weighted"))
    else:
        auc = float(roc_auc_score(test_labels, test_probs[:, 1]))

    print("MLP (face_orientation) — test set")
    print("  Accuracy: {:.2%}".format(test_acc))
    print("  F1:       {:.4f}".format(f1))
    print("  ROC-AUC:  {:.4f}".format(auc))


if __name__ == "__main__":
    main()