"""
FontID — SigLIP2 NaFlex backbone + four linear-probe-style classification heads.

Self-contained: rebuilds the vision tower from the vision_config embedded in
config.json (no dependency on the base repo at load time) and loads
model.safetensors. Predicts font / language / color / style.

    from modeling_fontid import FontIDModel
    model = FontIDModel.from_pretrained("path/to/dir")   # or the HF repo dir
"""
import os, json
import torch
import torch.nn as nn
from safetensors.torch import load_file
from transformers.models.siglip2.modeling_siglip2 import (
    Siglip2VisionTransformer, Siglip2VisionConfig,
)

HEADS = ["font", "lang", "color", "style"]


class FontIDModel(nn.Module):
    """Pooled feature from the SigLIP2 attention-pooling head (d=768) -> shared
    Dropout(0.2) -> per head: Linear(768->512) -> LayerNorm -> GELU ->
    Dropout(0.1) -> Linear(512->n)."""

    def __init__(self, vision_config: dict, num_classes: dict, heads=HEADS):
        super().__init__()
        self.heads = list(heads)
        cfg = Siglip2VisionConfig(**vision_config)
        # Rebuilt from a dict -> _attn_implementation is unset; pin it so the
        # attention dispatch resolves (matches the trained eager forward).
        cfg._attn_implementation = "eager"
        self.vision = Siglip2VisionTransformer(cfg)
        d = cfg.hidden_size                      # 768
        self.dropout = nn.Dropout(0.2)

        def make_head(n):
            return nn.Sequential(
                nn.Linear(d, 512),
                nn.LayerNorm(512),
                nn.GELU(),
                nn.Dropout(0.1),
                nn.Linear(512, n),
            )
        self.head_modules = nn.ModuleDict({h: make_head(num_classes[h]) for h in self.heads})

    @torch.no_grad()
    def forward(self, pixel_values, pixel_attention_mask, spatial_shapes):
        vis = self.vision(
            pixel_values=pixel_values,
            attention_mask=pixel_attention_mask,   # HF NaFlex arg name
            spatial_shapes=spatial_shapes,
        )
        feat = self.dropout(vis.pooler_output)     # d=768
        return {h: self.head_modules[h](feat) for h in self.heads}

    @classmethod
    def from_pretrained(cls, path, device="cpu"):
        with open(os.path.join(path, "config.json")) as f:
            cfg = json.load(f)
        model = cls(cfg["vision_config"], cfg["num_classes"], cfg.get("heads", HEADS))
        sd = load_file(os.path.join(path, "model.safetensors"))
        missing, unexpected = model.load_state_dict(sd, strict=True)
        model.eval().to(device)
        return model