import json
import torch
import torch.nn as nn
from PIL import Image, ImageOps
from torchvision import models, transforms


def build_model():
    model = models.efficientnet_b0(weights=None)
    model.classifier = nn.Sequential(
        nn.Dropout(p=0.2),
        nn.Linear(in_features=1280, out_features=1),
    )
    return model


def load_model(weights_path, device="cpu"):
    model = build_model()
    model.load_state_dict(torch.load(weights_path, map_location=device))
    model.to(device)
    model.eval()
    return model


def predict(model, image_path, config_path, device="cpu"):
    with open(config_path) as f:
        cfg = json.load(f)

    transform = transforms.Compose([
        transforms.Resize(cfg["resize_before_crop"]),
        transforms.CenterCrop(cfg["image_size"]),
        transforms.ToTensor(),
        transforms.Normalize(cfg["normalize_mean"], cfg["normalize_std"]),
    ])

    img = ImageOps.exif_transpose(Image.open(image_path)).convert("RGB")
    x = transform(img).unsqueeze(0).to(device)

    with torch.no_grad():
        prob_class1 = torch.sigmoid(model(x)).item()

    idx = int(prob_class1 > 0.5)
    confidence = prob_class1 if idx == 1 else 1 - prob_class1
    return cfg["class_names"][idx], confidence