"""Imager-1-Pro - Model architecture and inference (ConvNeXt-Base)."""
import os
import torch
import torch.nn as nn
from torchvision import models
from PIL import Image

from .preprocess import get_transform


class ImagerPro(nn.Module):
    """ConvNeXt-Base backbone with a single-logit binary head."""

    def __init__(self, pretrained: bool = False):
        super().__init__()
        self.backbone = models.convnext_base(
            weights="IMAGENET1K_V1" if pretrained else None
        )
        # ConvNeXt's classifier is a Sequential: [LayerNorm, Flatten, Linear]
        in_features = self.backbone.classifier[2].in_features
        self.backbone.classifier[2] = nn.Linear(in_features, 1)

    def forward(self, x):
        return self.backbone(x)


def load_model(weights_path: str = None, device: str = "cpu") -> ImagerPro:
    """Load the trained Imager-1-Pro model."""
    if weights_path is None:
        weights_path = os.path.join(
            os.path.dirname(__file__), "..", "model", "model_imager-1-pro.pth"
        )
    model = ImagerPro(pretrained=False)
    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state)
    model.to(device)
    model.eval()
    return model


def predict(image, model: ImagerPro = None, device: str = "cpu") -> dict:
    """Predict whether an image is Real or AI-Generated."""
    if model is None:
        model = load_model(device=device)

    transform = get_transform(train=False)

    if isinstance(image, str):
        image = Image.open(image)
    elif hasattr(image, "read"):
        image = Image.open(image)
    image = image.convert("RGB")

    tensor = transform(image).unsqueeze(0).to(device)

    with torch.no_grad():
        prob = torch.sigmoid(model(tensor)).item()

    label = "AI-Generated" if prob > MODEL_CONFIG["threshold"] else "Real"
    return {
        "probability": round(prob, 4),
        "label": label,
        "confidence": round(prob if label == "AI-Generated" else 1 - prob, 4),
    }


MODEL_CONFIG = {
    "name": "imager-1-pro",
    "architecture": "convnext_base",
    "num_params": 89_000_000,
    "input_size": 224,
    "num_classes": 1,
    "threshold": 0.5,
    "labels": {0: "Real", 1: "AI-Generated"},
    "normalize": {
        "mean": [0.485, 0.456, 0.406],
        "std": [0.229, 0.224, 0.225],
    },
}