import torch
import torch.nn as nn
from torchvision import transforms
from torchvision.models import efficientnet_v2_s
from PIL import Image
from collections import OrderedDict

DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

MODEL_PATH = "app/model/model.pth"
IMG_SIZE = 224

checkpoint = torch.load(
    MODEL_PATH,
    map_location=DEVICE,
    weights_only=False
)

class_names = checkpoint["class_names"]
num_classes = len(class_names)

model = efficientnet_v2_s(weights=None)

in_features = model.classifier[1].in_features

model.classifier = nn.Sequential(
    nn.Dropout(p=0.3, inplace=True),
    nn.Linear(in_features, num_classes),
)

state_dict = checkpoint["model_state_dict"]

new_state_dict = OrderedDict()

for k, v in state_dict.items():
    name = k.replace("module.", "")
    new_state_dict[name] = v

if "n_averaged" in new_state_dict:
    del new_state_dict["n_averaged"]

model.load_state_dict(new_state_dict)

model.to(DEVICE)
model.eval()

transform = transforms.Compose([
    transforms.Resize((IMG_SIZE, IMG_SIZE)),
    transforms.ToTensor(),
    transforms.Normalize(
        [0.485, 0.456, 0.406],
        [0.229, 0.224, 0.225]
    ),
])

def predict_image(image):

    image = image.convert("RGB")

    image_tensor = transform(image).unsqueeze(0).to(DEVICE)

    with torch.no_grad():

        outputs = model(image_tensor)

        probabilities = torch.softmax(outputs, dim=1)

        confidence, predicted = torch.max(probabilities, 1)

    return {
        "prediction": class_names[predicted.item()],
        "confidence": round(confidence.item() * 100, 2)
    }