"""Load and run the exported ConvNeXt Small model with PyTorch and Torchvision."""

from __future__ import annotations

import json
from pathlib import Path

import torch
from PIL import Image
from torchvision import models, transforms


def load_model(model_dir: str | Path, device: str = "auto"):
    """Load model weights and the preprocessing parameters from an export folder."""
    model_dir = Path(model_dir)
    config = json.loads((model_dir / "config.json").read_text(encoding="utf-8"))
    if config["model_name"] != "convnext_small":
        raise ValueError(f"Unsupported model architecture: {config['model_name']}")

    if device == "auto":
        if torch.cuda.is_available():
            device = "cuda"
        elif torch.backends.mps.is_available():
            device = "mps"
        else:
            device = "cpu"
    target = torch.device(device)

    model = models.convnext_small(weights=None)
    model.classifier[-1] = torch.nn.Linear(
        model.classifier[-1].in_features, config["num_labels"],
    )
    state_dict = torch.load(
        model_dir / config["weights_file"], map_location="cpu", weights_only=True,
    )
    model.load_state_dict(state_dict)
    model.to(target).eval()

    preprocess = transforms.Compose([
        transforms.Resize(config["resize_size"]),
        transforms.CenterCrop(config["image_size"]),
        transforms.ToTensor(),
        transforms.Normalize(mean=config["image_mean"], std=config["image_std"]),
    ])
    return model, preprocess, config["class_names"], target


@torch.inference_mode()
def predict(image: Image.Image, model, preprocess, class_names: list[str], device) -> list[dict]:
    """Return all class scores, ordered from highest to lowest."""
    tensor = preprocess(image.convert("RGB")).unsqueeze(0).to(device)
    probabilities = torch.softmax(model(tensor), dim=1)[0].cpu().tolist()
    return sorted(
        ({"class_name": name, "probability": float(probabilities[index])}
         for index, name in enumerate(class_names)),
        key=lambda score: score["probability"], reverse=True,
    )