"""Standalone loader for clip-vitb-mini-distilled (no repo imports).

Usage:
    from loader import load_model, embed_images
    model = load_model("student.pt")            # or a hf_hub_download path
    emb = embed_images(model, batch01)          # (B,3,H,W) in [0,1] -> (B,512)

The student is an 8.66M-parameter ViT (d=240, depth 12, heads 4, patch 16,
img 160) distilled from CLIP ViT-B/16 (LAION-2B) pooled image features on
COCO-2017 train (118,287 images, ~17 epochs). Outputs live in the teacher's
512-d projection space: compatible with the teacher's TEXT tower for
zero-shot/retrieval. See the model card for measured capability and caveats.
"""
import torch
import torch.nn as nn
import torch.nn.functional as F

CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
CLIP_STD = (0.26862954, 0.26130258, 0.27577711)


class Block(nn.Module):
    def __init__(self, d, heads):
        super().__init__()
        self.n1 = nn.LayerNorm(d)
        self.qkv = nn.Linear(d, 3 * d)
        self.proj = nn.Linear(d, d)
        self.n2 = nn.LayerNorm(d)
        self.fc1 = nn.Linear(d, 4 * d)
        self.fc2 = nn.Linear(4 * d, d)
        self.heads = heads

    def forward(self, x):
        B, N, C = x.shape
        q, k, v = (self.qkv(self.n1(x))
                   .reshape(B, N, 3, self.heads, C // self.heads)
                   .permute(2, 0, 3, 1, 4))
        a = F.scaled_dot_product_attention(q, k, v)
        x = x + self.proj(a.transpose(1, 2).reshape(B, N, C))
        return x + self.fc2(F.gelu(self.fc1(self.n2(x))))


class Student(nn.Module):
    """CLS-token readout, Linear head to the teacher's 512-d space."""
    def __init__(self, out_dim=512, d=240, depth=12, heads=4, patch=16,
                 img=160):
        super().__init__()
        self.patch = nn.Conv2d(3, d, patch, patch)
        self.cls = nn.Parameter(torch.zeros(1, 1, d))
        self.pos = nn.Parameter(torch.zeros(1, (img // patch) ** 2 + 1, d))
        self.blocks = nn.ModuleList(Block(d, heads) for _ in range(depth))
        self.norm = nn.LayerNorm(d)
        self.head = nn.Linear(d, out_dim)

    def forward_features(self, x):
        x = self.patch(x).flatten(2).transpose(1, 2)
        x = torch.cat([self.cls.expand(x.shape[0], -1, -1), x], 1) + self.pos
        for b in self.blocks:
            x = b(x)
        return self.norm(x)[:, 0]

    def forward(self, x):
        return self.head(self.forward_features(x))


def load_model(path, device=None):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")
    ck = torch.load(path, map_location="cpu", weights_only=True)
    sd = ck["state_dict"] if "state_dict" in ck else ck
    model = Student(out_dim=sd["head.weight"].shape[0])
    model.load_state_dict(sd, strict=True)
    return model.to(device).eval()


def load_rotation(path):
    """Load a frame rotation shipped beside a consensus-distilled student.
    The rotation maps the student's consensus frame into the CLIP-B/16
    (LAION-2B) deployment frame; it was fitted once, offline, on 2,500
    COCO-val pairs (fp64 orthogonal Procrustes) and is FROZEN."""
    d = torch.load(path, map_location="cpu", weights_only=True)
    return d["R"] if isinstance(d, dict) else d


@torch.no_grad()
def embed_images(model, x01, batch=256, rotation=None):
    """x01: (B,3,H,W) float in [0,1]. Returns L2-normalized (B,512).
    Pass `rotation` (from load_rotation) with the consensus-distilled
    champion to place outputs in the LAION text tower's frame — required
    for zero-shot/retrieval against that tower, and how every headline
    number on the card was measured."""
    dev = next(model.parameters()).device
    mean = torch.tensor(CLIP_MEAN).view(1, 3, 1, 1).to(dev)
    std = torch.tensor(CLIP_STD).view(1, 3, 1, 1).to(dev)
    out = []
    for i in range(0, len(x01), batch):
        x = x01[i:i + batch].to(dev)
        x = F.interpolate(x, size=(160, 160), mode="bicubic",
                          align_corners=False).clamp(0, 1)
        z = F.normalize(model((x - mean) / std), dim=-1)
        if rotation is not None:
            z = F.normalize(z.double() @ rotation.to(dev).double(),
                            dim=-1).float()
        out.append(z.cpu())
    return torch.cat(out)