import torch
from collections import OrderedDict
import numpy as np
from PIL import Image
import torchvision.transforms as transforms
from model import get_model

_preprocess = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    )
])

def load_model(weights_path: str, device: torch.device):
    checkpoint = torch.load(weights_path, map_location=device, weights_only=False)
    if "model_state_dict" not in checkpoint:
        raise KeyError("model_state_dict not found in checkpoint")
    state_dict = checkpoint["model_state_dict"]
    new_state_dict = OrderedDict()
    for k, v in state_dict.items():
        new_state_dict[k.replace("module.", "")] = v
    model = get_model()
    model.load_state_dict(new_state_dict)
    model.to(device)
    model.eval()
    return model

def preprocess(image):
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    return _preprocess(image).unsqueeze(0)

def predict(image, model, device):
    model.eval()
    with torch.no_grad():
        tensor = preprocess(image).to(device)
        output = model(tensor)
        pred = torch.argmax(output, dim=1).squeeze(0).cpu()
    return pred