import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from torchvision import transforms
from torchvision.models import ResNet18_Weights, resnet18


MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]


def build_model(num_classes: int, pretrained: bool = False):
    model = resnet18(weights=ResNet18_Weights.DEFAULT if pretrained else None)
    model.fc = torch.nn.Linear(model.fc.in_features, num_classes)
    return model


def preprocessing(config: dict, training: bool = False):
    size = config["image_size"]
    operations = (
        [transforms.RandomResizedCrop(size, scale=(.8, 1.0)), transforms.RandomHorizontalFlip()]
        if training else
        [transforms.Resize(round(size * 256 / 224)), transforms.CenterCrop(size)]
    )
    return transforms.Compose(operations + [
        transforms.ToTensor(),
        transforms.Normalize(config["mean"], config["std"]),
    ])


def load_classifier(directory: str | Path):
    directory = Path(directory)
    config = json.loads((directory / "config.json").read_text())
    model = build_model(len(config["labels"]))
    model.load_state_dict(load_file(str(directory / "model.safetensors")))
    model.eval()
    return model, config


@torch.inference_mode()
def classify(model, config: dict, image, top_k: int):
    if image is None:
        raise ValueError("Choose an image")
    if not 1 <= top_k <= len(config["labels"]):
        raise ValueError("Invalid top_k")
    tensor = preprocessing(config)(image.convert("RGB")).unsqueeze(0)
    tensor = tensor.to(next(model.parameters()).device)
    probabilities = model(tensor).softmax(dim=1)[0]
    scores, indices = probabilities.topk(top_k)
    return [
        {"label": config["labels"][index], "score": float(score)}
        for score, index in zip(scores.tolist(), indices.tolist())
    ]