# backbones.py — image/text encoder adapters (OpenAI CLIP + SigLIP2) for the gradstep engine.
import torch
from torch import nn

import generate_fast as gf
from generate_fast import clip

OPENAI_MODELS = ("ViT-B/16", "ViT-B/32", "ViT-L/14")
SIGLIP_NORM = ([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])


class SiglipVisual(nn.Module):
    """(B,3,S,S) pixel values -> pooled image embedding; mimics the CLIP visual interface."""

    def __init__(self, vision_model, input_resolution, output_dim):
        super().__init__()
        self.vision_model = vision_model
        self.input_resolution = input_resolution
        self.output_dim = output_dim

    def forward(self, x):
        return self.vision_model(x).pooler_output


def load_backbone(name):
    """-> dict(kind, visual, cut_size, output_dim, encode_text(txt, device), text_module)."""
    if name in OPENAI_MODELS:
        p = clip.load(name, device="cpu", jit=False)[0].eval().requires_grad_(False)
        # text tower in fp32 (GradStep casts the visual to bf16; CLIP's encode_text would
        # otherwise cast text inputs by the visual dtype and mismatch the text weights)
        p.transformer.float()
        p.token_embedding.float()
        p.ln_final.float()
        p.positional_embedding.data = p.positional_embedding.data.float()
        p.text_projection.data = p.text_projection.data.float()

        def encode_text(txt, device):
            t = clip.tokenize(txt).to(device)
            x = p.token_embedding(t) + p.positional_embedding
            x = p.transformer(x.permute(1, 0, 2)).permute(1, 0, 2)
            x = p.ln_final(x)
            return x[torch.arange(x.shape[0]), t.argmax(dim=-1)] @ p.text_projection

        return {"kind": "openai", "visual": p.visual, "cut_size": p.visual.input_resolution,
                "output_dim": p.visual.output_dim, "encode_text": encode_text, "text_module": p}

    from transformers import AutoModel, AutoTokenizer
    m = AutoModel.from_pretrained(name).eval().requires_grad_(False)
    tok = AutoTokenizer.from_pretrained(name)
    res = m.config.vision_config.image_size
    dim = m.config.vision_config.hidden_size
    visual = SiglipVisual(m.vision_model, res, dim)

    def encode_text(txt, device):
        ids = tok([txt], padding="max_length", max_length=64,
                  truncation=True, return_tensors="pt").input_ids.to(device)
        return m.text_model(input_ids=ids).pooler_output  # transformers-5-safe

    return {"kind": "siglip", "visual": visual, "cut_size": res, "output_dim": dim,
            "encode_text": encode_text, "text_module": m.text_model}


class _FP32LayerNorm(nn.LayerNorm):
    """CLIP-style LN: compute in fp32, cast back — cuts bf16 gradient noise."""

    def forward(self, x):
        return nn.functional.layer_norm(
            x.float(), self.normalized_shape, self.weight, self.bias, self.eps).to(x.dtype)


def _fp32_layernorms(module):
    for name, child in module.named_children():
        if isinstance(child, nn.LayerNorm):
            ln = _FP32LayerNorm(child.normalized_shape, eps=child.eps,
                                elementwise_affine=child.elementwise_affine)
            if child.elementwise_affine:
                ln.weight.data = child.weight.data.float()
                if child.bias is not None:
                    ln.bias.data = child.bias.data.float()
            ln.requires_grad_(False)
            setattr(module, name, ln)
        else:
            _fp32_layernorms(child)


def _siglip_fixups(sm, backbone):
    if backbone["kind"] == "siglip":
        # HF siglip: bf16 everywhere except LNs in fp32 (mirrors the CLIP visual policy)
        sm.visual.to(torch.bfloat16)
        _fp32_layernorms(sm.visual)
        mean, std = SIGLIP_NORM
        sm.norm_mean.copy_(torch.tensor(mean).view(1, 3, 1, 1))
        sm.norm_std.copy_(torch.tensor(std).view(1, 3, 1, 1))
    return sm


def build_gradstep(vq, backbone, cutn, sideY, sideX):
    sm = gf.GradStep(vq, backbone["visual"], cutn, backbone["cut_size"], sideY, sideX, 0.0).eval()
    return _siglip_fixups(sm, backbone)


def build_fusedstep(unet, backbone, cutn, image_size):
    import diffusion_fast as D  # lazy: only the diffusion Spaces ship this module
    sm = D.FusedStep(unet, backbone["visual"], cutn, backbone["cut_size"],
                     image_size, torch.bfloat16).eval()
    return _siglip_fixups(sm, backbone)