"""Imager-1-Tiny - Model architecture and inference."""
import os
import torch
import torch.nn as nn
from torchvision import models
from PIL import Image

from .preprocess import get_transform


class ImagerTiny(nn.Module):
    """ResNet-50 backbone with a single-logit binary head."""

    def __init__(self, pretrained: bool = False):
        super().__init__()
        self.backbone = models.resnet50(
            weights="IMAGENET1K_V1" if pretrained else None
        )
        in_features = self.backbone.fc.in_features
        self.backbone.fc = nn.Linear(in_features, 1)

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


def load_model(weights_path: str = None, device: str = "cpu") -> ImagerTiny:
    """Load the trained Imager-1-Tiny model."""
    if weights_path is None:
        weights_path = os.path.join(
            os.path.dirname(__file__), "..", "model", "model_imager-1-tiny.pth"
        )
    model = ImagerTiny(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: ImagerTiny = None, device: str = "cpu") -> dict:
    """
    Predict whether an image is Real or AI-Generated.

    Args:
        image: PIL.Image, str path, or file-like object.
        model: Pre-loaded ImagerTiny (optional).
        device: "cpu" or "cuda".

    Returns:
        dict with probability, label, confidence.
    """
    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-tiny",
    "architecture": "resnet50",
    "num_params": 25_000_000,
    "input_size": 224,
    "num_classes": 1,
    "threshold": 0.385,
    "labels": {0: "Real", 1: "AI-Generated"},
    "normalize": {
        "mean": [0.485, 0.456, 0.406],
        "std": [0.229, 0.224, 0.225],
    },
}