"""TinyViT for surface-defect image classification (standalone, no project deps)."""
import torch
import torch.nn as nn


class TinyViT(nn.Module):
    """Image-classification ViT: img_size x img_size -> patch x patch patches."""

    def __init__(self, img_size: int = 64, patch: int = 8, d_model: int = 64,
                 n_layers: int = 2, n_heads: int = 4, n_classes: int = 6,
                 in_ch: int = 1, dropout: float = 0.1):
        super().__init__()
        assert img_size % patch == 0, "img_size must be divisible by patch"
        self.img_size = img_size
        self.patch = patch
        self.n_patches = (img_size // patch) ** 2
        self.patch_embed = nn.Conv2d(in_ch, d_model,
                                     kernel_size=patch, stride=patch)
        self.cls = nn.Parameter(torch.zeros(1, 1, d_model))
        self.pos = nn.Parameter(
            torch.zeros(1, self.n_patches + 1, d_model))
        layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=n_heads,
            dim_feedforward=4 * d_model, dropout=dropout,
            batch_first=True, activation="gelu")
        self.encoder = nn.TransformerEncoder(layer, num_layers=n_layers)
        self.norm = nn.LayerNorm(d_model)
        self.head = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(d_model, d_model),
            nn.GELU(),
            nn.Linear(d_model, n_classes))

    def forward(self, x):
        B = x.shape[0]
        t = self.patch_embed(x).flatten(2).transpose(1, 2)  # (B, P, d)
        t = torch.cat([self.cls.expand(B, -1, -1), t], dim=1) + self.pos
        t = self.encoder(t)
        return self.head(self.norm(t)[:, 0])


CLASS_NAMES = ["crazing", "inclusion", "patches",
               "pitted_surface", "rolled-in_scale", "scratches"]

# Normalisation statistics from training (uint8/255 scale)
MEAN = 0.5013
STD = 0.2075


def load_model(weights_path: str, device: str = "cpu") -> TinyViT:
    model = TinyViT()
    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state)
    model.eval()
    return model


def predict_image(model: TinyViT, image) -> tuple[int, "np.ndarray"]:
    """Classify a PIL image or numpy array (grayscale or RGB).

    Returns (predicted_class_id, class_probabilities[6]).
    """
    import numpy as np
    from PIL import Image
    if isinstance(image, np.ndarray):
        image = Image.fromarray(image)
    image = image.convert("L").resize((64, 64), Image.BILINEAR)
    x = np.asarray(image, dtype=np.float32) / 255.0
    x = (x - MEAN) / STD
    xt = torch.from_numpy(x)[None, None].to(next(model.parameters()).device)
    with torch.no_grad():
        probs = torch.softmax(model(xt), dim=1).cpu().numpy()[0]
    return int(probs.argmax()), probs