"""
Hugging Face Space — GPT2VL Stackformer V2 — Image Captioning Only
===================================================================
Model trained in float32 with torch.amp.autocast (fp16 matmuls, fp32 accumulators).
Inference mirrors training exactly:
  • float32 model weights
  • torch.amp.autocast for GPU execution
  • greedy argmax decoding (no sampling)
  • ViT stays float32; patch tokens cast to resampler dtype inside encode_image

Checkpoint layout on the Hub:
  config.json                 — architecture hyperparameters
  model_trainable.safetensors — adapter weights (resampler.* + cross_blocks.*)
"""

import os, json
import torch
import torch.nn as nn
import torch.nn.functional as F
import gradio as gr
from PIL import Image

# ── ZeroGPU shim ─────────────────────────────────────────────────────────────
try:
    import spaces
except ImportError:
    class spaces:                          # noqa: E302
        @staticmethod
        def GPU(duration=None):
            def decorator(fn): return fn
            return decorator

from torchvision.models import vit_b_16, ViT_B_16_Weights
from torchvision import transforms
from transformers import GPT2TokenizerFast, GPT2LMHeadModel
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file

REPO_ID     = os.environ.get("MODEL_REPO_ID", "gurumurthy3/gpt2vl-stackformer-v2")
device      = torch.device("cuda" if torch.cuda.is_available() else "cpu")
MODEL_DTYPE = torch.float32   # model saved and loaded in fp32 — matches training


# ── Stackformer components (self-contained, FP32-stable LayerNorm) ────────────
class LayerNormalization(nn.Module):
    """Always upcasts to float32 for mean/var math — safe under torch.amp.autocast."""
    def __init__(self, embed_dim, eps=1e-5, device=None, dtype=None):
        super().__init__()
        self.eps = eps
        fk = {"device": device, "dtype": dtype}
        self.weight = nn.Parameter(torch.ones(embed_dim, **fk))
        self.bias   = nn.Parameter(torch.zeros(embed_dim, **fk))

    def forward(self, x):
        orig = x.dtype
        x32  = x.float()
        mean = x32.mean(-1, keepdim=True)
        var  = x32.var(-1, keepdim=True, unbiased=False)
        out  = self.weight.float() * (x32 - mean) / (var + self.eps).sqrt() + self.bias.float()
        return out.to(orig)


class FF_GELU(nn.Module):
    """Feed-forward with GELU, matching stackformer FF_GELU structure."""
    def __init__(self, embed_dim, hidden_dim, dropout=0.0, device=None, dtype=None):
        super().__init__()
        kw = {"device": device, "dtype": dtype}
        # NOTE: stackformer FF_GELU stores as self.gelu = nn.Sequential(...)
        # We expose the same attribute so weight-copy code can use gelu[0]/gelu[3]
        self.gelu = nn.Sequential(
            nn.Linear(embed_dim, hidden_dim, **kw),  # index 0
            nn.GELU(),                                # index 1
            nn.Dropout(dropout),                      # index 2
            nn.Linear(hidden_dim, embed_dim, **kw),  # index 3
            nn.Dropout(dropout),                      # index 4
        )

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


class AbsolutePositionEmbedding(nn.Module):
    def __init__(self, seq_len, embed_dim, device=None, dtype=None):
        super().__init__()
        self.embedding = nn.Embedding(seq_len, embed_dim, device=device, dtype=dtype)

    def forward(self, x):
        B, T = x.shape[:2]
        pos = torch.arange(T, device=self.embedding.weight.device, dtype=torch.long)
        return self.embedding(pos).unsqueeze(0).expand(B, -1, -1)


class Multi_Head_Attention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.0, qkv_bias=True, device=None, dtype=None):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim  = embed_dim // num_heads
        self.dropout_p = dropout
        kw = {"device": device, "dtype": dtype}
        self.qkv_proj  = nn.Linear(embed_dim, embed_dim * 3, bias=qkv_bias, **kw)
        self.out_proj  = nn.Linear(embed_dim, embed_dim,     bias=qkv_bias, **kw)

    def forward(self, x, mask=True):
        B, T, C = x.shape
        q, k, v = self.qkv_proj(x).split(self.embed_dim, dim=-1)
        def rs(t): return t.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        dp = self.dropout_p if self.training else 0.0
        out = F.scaled_dot_product_attention(rs(q), rs(k), rs(v), dropout_p=dp, is_causal=mask)
        return self.out_proj(out.transpose(1, 2).contiguous().view(B, T, C))


class Cross_MultiHead_Attention(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout=0.0, qkv_bias=True, device=None, dtype=None):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.head_dim  = embed_dim // num_heads
        self.dropout_p = dropout
        kw = {"device": device, "dtype": dtype}
        self.q_proj  = nn.Linear(embed_dim, embed_dim,     bias=qkv_bias, **kw)
        self.kv_proj = nn.Linear(embed_dim, embed_dim * 2, bias=qkv_bias, **kw)
        self.out_proj = nn.Linear(embed_dim, embed_dim,    bias=qkv_bias, **kw)

    def forward(self, x, context, mask=False, attn_mask=None):
        B, T, C = x.shape
        S = context.size(1)
        q = self.q_proj(x)
        k, v = self.kv_proj(context).split(self.embed_dim, dim=-1)
        def rsq(t): return t.view(B, T, self.num_heads, self.head_dim).transpose(1, 2)
        def rsc(t): return t.view(B, S, self.num_heads, self.head_dim).transpose(1, 2)
        dp = self.dropout_p if self.training else 0.0
        out = F.scaled_dot_product_attention(rsq(q), rsc(k), rsc(v), dropout_p=dp, is_causal=False)
        return self.out_proj(out.transpose(1, 2).contiguous().view(B, T, C))


class EncoderBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, hidden_dim, dropout=0.0, qkv_bias=True, device=None, dtype=None):
        super().__init__()
        kw = dict(device=device, dtype=dtype)
        self.self_attn = Multi_Head_Attention(embed_dim, num_heads, dropout=dropout, qkv_bias=qkv_bias, **kw)
        self.ffn   = FF_GELU(embed_dim, hidden_dim, dropout=dropout, **kw)
        self.norm1 = LayerNormalization(embed_dim, **kw)
        self.norm2 = LayerNormalization(embed_dim, **kw)

    # expose aliases matching stackformer naming so weight-transfer code works
    @property
    def attention(self): return self.self_attn
    @property
    def ff(self): return self.ffn

    def forward(self, x, mask=False):
        x = x + self.self_attn(self.norm1(x), mask=mask)
        x = x + self.ffn(self.norm2(x))
        return x


class TransformerEncoder(nn.Module):
    def __init__(self, embed_dim, num_heads, hidden_dim, num_layers, dropout=0.0, qkv_bias=True, device=None, dtype=None):
        super().__init__()
        kw = dict(embed_dim=embed_dim, num_heads=num_heads, hidden_dim=hidden_dim,
                  dropout=dropout, qkv_bias=qkv_bias, device=device, dtype=dtype)
        self.layers     = nn.ModuleList([EncoderBlock(**kw) for _ in range(num_layers)])
        self.final_norm = LayerNormalization(embed_dim, device=device, dtype=dtype)

    def forward(self, x, mask=False):
        for layer in self.layers:
            x = layer(x, mask=mask)
        return self.final_norm(x)


class GPT_2(nn.Module):
    def __init__(self, vocab_size, num_layers, embed_dim, num_heads, seq_len,
                 dropout=0.1, hidden_dim=0, qkv_bias=True, eps=1e-5, device="cpu", dtype=None):
        super().__init__()
        hidden_dim = hidden_dim or 4 * embed_dim
        kw = dict(device=device, dtype=dtype)
        self.embedding          = nn.Embedding(vocab_size, embed_dim, **kw)
        self.position_embedding = AbsolutePositionEmbedding(seq_len, embed_dim, **kw)
        self.backbone           = TransformerEncoder(embed_dim, num_heads, hidden_dim, num_layers,
                                                     dropout=dropout, qkv_bias=qkv_bias, **kw)
        self.final_norm         = LayerNormalization(embed_dim, eps=eps, **kw)
        self.lm_head            = nn.Linear(embed_dim, vocab_size, bias=False, **kw)

    def forward(self, x):
        x = self.embedding(x) + self.position_embedding(x)
        x = self.backbone(x, mask=True)
        return self.lm_head(x)


# ── Multimodal components ──────────────────────────────────────────────────────
class GatedSparseCrossAttnBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout, qkv_bias=True, device="cpu", dtype=None):
        super().__init__()
        kw = dict(device=device, dtype=dtype)
        self.norm       = LayerNormalization(embed_dim, **kw)
        self.cross_attn = Cross_MultiHead_Attention(embed_dim, num_heads, dropout=dropout, qkv_bias=qkv_bias, **kw)
        self.drop       = nn.Dropout(dropout)
        self.alpha      = nn.Parameter(torch.zeros(1, **kw))

    def forward(self, x, context):
        residual = x
        x_norm   = self.norm(x)
        attn_out = self.cross_attn(x_norm, context, mask=False)
        attn_out = self.drop(attn_out)
        return residual + self.alpha * attn_out


class PerceiverResamplerSF(nn.Module):
    def __init__(self, embed_dim, num_latents, depth, num_heads, dropout=0.0,
                 hidden_dim=None, device="cpu", dtype=None):
        super().__init__()
        hidden_dim = hidden_dim or embed_dim * 2
        kw = dict(dropout=dropout, qkv_bias=True, device=device, dtype=dtype)
        self.latents      = nn.Parameter(torch.randn(num_latents, embed_dim, device=device, dtype=dtype) * 0.02)
        self.cross_layers = nn.ModuleList([Cross_MultiHead_Attention(embed_dim, num_heads, **kw) for _ in range(depth)])
        self.norm_latent  = nn.ModuleList([LayerNormalization(embed_dim, device=device, dtype=dtype) for _ in range(depth)])
        self.norm_media   = nn.ModuleList([LayerNormalization(embed_dim, device=device, dtype=dtype) for _ in range(depth)])
        self.ffns         = nn.ModuleList([FF_GELU(embed_dim, hidden_dim, dropout, device=device, dtype=dtype) for _ in range(depth)])
        self.ffn_norms    = nn.ModuleList([LayerNormalization(embed_dim, device=device, dtype=dtype) for _ in range(depth)])
        self.depth        = depth

    def forward(self, media_seq):
        b = media_seq.shape[0]
        x = self.latents.unsqueeze(0).expand(b, -1, -1)
        for i in range(self.depth):
            xn  = self.norm_latent[i](x)
            ctx = self.norm_media[i](media_seq)
            x   = x + self.cross_layers[i](xn, ctx, mask=False)
            x   = x + self.ffns[i](self.ffn_norms[i](x))
        return x


class TorchvisionViTEncoder(nn.Module):
    """Frozen ViT-B/16. Kept in float32 (same as notebook — NO dtype cast at init).
    Patch tokens are cast to the resampler's dtype inside GPT2VL.encode_image."""
    def __init__(self, pretrained=True, freeze=True):
        super().__init__()
        weights    = ViT_B_16_Weights.IMAGENET1K_V1 if pretrained else None
        self.model = vit_b_16(weights=weights)   # stays float32
        self.hidden_dim = self.model.hidden_dim
        if freeze:
            for p in self.parameters():
                p.requires_grad = False

    @torch.no_grad()
    def forward(self, images):
        f   = self.model._process_input(images)
        b   = f.shape[0]
        x   = torch.cat((self.model.class_token.expand(b, -1, -1), f), dim=1)
        x   = self.model.encoder(x)
        return x   # (B, 197, 768) — includes CLS token


class GPT2VL(nn.Module):
    def __init__(self, cfg, device="cpu", dtype=None):
        super().__init__()
        self.cfg = cfg
        kw = dict(device=device, dtype=dtype)
        self.gpt2 = GPT_2(
            vocab_size=cfg["vocab_size"], num_layers=cfg["num_layers"],
            embed_dim=cfg["embed_dim"],   num_heads=cfg["num_heads"],
            seq_len=cfg["context_length"], dropout=cfg["dropout"],
            hidden_dim=cfg["hidden_dim"], qkv_bias=cfg["qkv_bias"], **kw,
        )
        self.cross_attention_pos = set(cfg["cross_attention_pos"])
        self.cross_blocks = nn.ModuleDict({
            str(i): GatedSparseCrossAttnBlock(cfg["embed_dim"], cfg["num_heads"],
                                               cfg["dropout"], cfg["qkv_bias"], **kw)
            for i in cfg["cross_attention_pos"]
        })
        # ViT stays float32 (no dtype arg) — matches notebook Cell 13
        self.vision_encoder = TorchvisionViTEncoder(pretrained=True, freeze=True)
        # No vision_project needed: vision_dim == embed_dim == 768
        self.resampler = PerceiverResamplerSF(
            cfg["embed_dim"], cfg["num_visual_tokens"], cfg["perceiver_depth"],
            cfg["perceiver_heads"], cfg["dropout"], device=device, dtype=dtype,
        )

    def encode_image(self, images):
        with torch.no_grad():
            patch_tokens = self.vision_encoder(images)   # (B, 197, 768) float32
        # Cast to resampler's dtype (float32 normally; fp16 under autocast)
        patch_tokens = patch_tokens.to(dtype=self.resampler.latents.dtype)
        return self.resampler(patch_tokens[:, 1:, :])    # drop CLS → (B, 196, 768)

    def forward(self, input_ids, images=None, visual_context=None):
        if visual_context is None and images is not None:
            visual_context = self.encode_image(images)

        x = self.gpt2.embedding(input_ids) + self.gpt2.position_embedding(input_ids)
        backbone = self.gpt2.backbone
        for i, layer in enumerate(backbone.layers):
            x = layer(x, mask=True)
            if i in self.cross_attention_pos and visual_context is not None:
                x = self.cross_blocks[str(i)](x, visual_context)
        return self.gpt2.lm_head(backbone.final_norm(x))

    def freeze_text_backbone(self):
        for p in self.gpt2.parameters():
            p.requires_grad = False
        for p in self.vision_encoder.parameters():
            p.requires_grad = False


# ── Tokenizer ─────────────────────────────────────────────────────────────────
tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({"pad_token": "<|pad|>"})
CFG_VOCAB_SIZE = len(tokenizer)   # 50258 base + 1 pad = 50259

bos_id = tokenizer.bos_token_id or tokenizer.eos_token_id
eos_id = tokenizer.eos_token_id

# ── Download & assemble model ─────────────────────────────────────────────────
print(f"[startup] downloading checkpoint from {REPO_ID} …")
config_path  = hf_hub_download(REPO_ID, "config.json")
weights_path = hf_hub_download(REPO_ID, "model_trainable.safetensors")

with open(config_path) as f:
    CFG = json.load(f)
CFG["vocab_size"] = CFG_VOCAB_SIZE   # sync with tokenizer (notebook Cell 23 line 1)
CONTEXT_LENGTH = CFG["context_length"]

print("[startup] building model in float32 …")
model = GPT2VL(CFG, device=str(device), dtype=MODEL_DTYPE).to(device=device, dtype=MODEL_DTYPE)
model.vision_encoder.to(device)   # ViT always on the right device

# ── GPT-2 weight transfer (mirrors notebook Cell 23 exactly) ─────────────────
print("[startup] transferring pretrained GPT-2 weights …")
hf_gpt2     = GPT2LMHeadModel.from_pretrained("gpt2")
hf_state    = hf_gpt2.state_dict()
hf_vocab_sz = hf_state["transformer.wte.weight"].shape[0]

layers = model.gpt2.backbone.layers

with torch.no_grad():
    # ── Token embedding ──────────────────────────────────────────────────────
    new_emb = model.gpt2.embedding.weight.data
    new_emb[:hf_vocab_sz] = hf_state["transformer.wte.weight"].to(MODEL_DTYPE)
    if CFG["vocab_size"] > hf_vocab_sz:
        nn.init.normal_(new_emb[hf_vocab_sz:], mean=0.0, std=0.02)

    # ── Position embedding ───────────────────────────────────────────────────
    model.gpt2.position_embedding.embedding.weight.copy_(
        hf_state["transformer.wpe.weight"][: CFG["context_length"]].to(MODEL_DTYPE)
    )

    # ── Transformer layers ───────────────────────────────────────────────────
    for i, block in enumerate(layers):
        p = f"transformer.h.{i}."

        # LayerNorm
        block.norm1.weight.copy_(hf_state[p + "ln_1.weight"].to(MODEL_DTYPE))
        block.norm1.bias.copy_(  hf_state[p + "ln_1.bias"].to(MODEL_DTYPE))
        block.norm2.weight.copy_(hf_state[p + "ln_2.weight"].to(MODEL_DTYPE))
        block.norm2.bias.copy_(  hf_state[p + "ln_2.bias"].to(MODEL_DTYPE))

        # Attention — HF Conv1D weight shape is (C, 3C); split BEFORE transposing
        # (matches notebook Cell 23 exactly):
        #   w_q, w_k, w_v = w_qkv.split(768, dim=1)   # each (768, 768)
        #   W_fused = cat([w_q.T, w_k.T, w_v.T], dim=0)  # (2304, 768)
        w_qkv = hf_state[p + "attn.c_attn.weight"]   # (768, 2304)
        b_qkv = hf_state[p + "attn.c_attn.bias"]     # (2304,)
        w_q, w_k, w_v = w_qkv.split(CFG["embed_dim"], dim=1)
        b_q, b_k, b_v = b_qkv.split(CFG["embed_dim"], dim=0)
        W_fused = torch.cat([w_q.T, w_k.T, w_v.T], dim=0).to(MODEL_DTYPE)  # (2304, 768)
        b_fused = torch.cat([b_q, b_k, b_v], dim=0).to(MODEL_DTYPE)         # (2304,)
        block.self_attn.qkv_proj.weight.copy_(W_fused)
        block.self_attn.qkv_proj.bias.copy_(b_fused)
        block.self_attn.out_proj.weight.copy_(hf_state[p + "attn.c_proj.weight"].T.to(MODEL_DTYPE))
        block.self_attn.out_proj.bias.copy_(  hf_state[p + "attn.c_proj.bias"].to(MODEL_DTYPE))

        # FFN — stackformer FF_GELU stores layers as self.gelu (nn.Sequential)
        # indices: 0=fc1, 1=GELU, 2=Dropout, 3=fc2, 4=Dropout
        block.ffn.gelu[0].weight.copy_(hf_state[p + "mlp.c_fc.weight"].T.to(MODEL_DTYPE))
        block.ffn.gelu[0].bias.copy_(  hf_state[p + "mlp.c_fc.bias"].to(MODEL_DTYPE))
        block.ffn.gelu[3].weight.copy_(hf_state[p + "mlp.c_proj.weight"].T.to(MODEL_DTYPE))
        block.ffn.gelu[3].bias.copy_(  hf_state[p + "mlp.c_proj.bias"].to(MODEL_DTYPE))

    # ── Final LayerNorm ──────────────────────────────────────────────────────
    model.gpt2.backbone.final_norm.weight.copy_(hf_state["transformer.ln_f.weight"].to(MODEL_DTYPE))
    model.gpt2.backbone.final_norm.bias.copy_(  hf_state["transformer.ln_f.bias"].to(MODEL_DTYPE))

    # ── LM head weight tying (notebook Cell 23 line 996) ────────────────────
    # GPT-2 ties lm_head.weight == wte; copy the (possibly extended) embedding
    model.gpt2.lm_head.weight.copy_(new_emb)

del hf_gpt2, hf_state

# ── Load trained adapter weights ─────────────────────────────────────────────
print("[startup] loading trained adapter weights …")
trained_state = load_file(weights_path)
# Weights were saved from a float32 model → load as-is, no dtype conversion needed
missing, unexpected = model.load_state_dict(trained_state, strict=False)
n_loaded = sum(1 for k in trained_state if k not in unexpected)
print(f"[startup] loaded {n_loaded} adapter tensors | unexpected: {len(unexpected)}")
if missing:
    print(f"[startup] WARNING missing keys: {missing[:5]} …")

model.eval()
print("[startup] model ready.")

# ── Image preprocessing (same as training dataset transform) ─────────────────
image_tx = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.Lambda(lambda im: im.convert("RGB")),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


# ── Local Sample Image Preparation (6 Examples from user repo) ───────────────
import time
import urllib.request
import ssl
from PIL import Image, ImageDraw

SAMPLES_DIR = os.path.join(os.path.dirname(__file__), "samples")
os.makedirs(SAMPLES_DIR, exist_ok=True)
SAMPLE_FILES = []

def _prepare_sample_images():
    ssl_ctx = ssl.create_default_context()
    ssl_ctx.check_hostname = False
    ssl_ctx.verify_mode = ssl.CERT_NONE

    base_url = "https://raw.githubusercontent.com/Gurumurthy30/multimodal-gpt2-demo/main/v1/examples/"
    colors = [(180, 140, 100), (140, 160, 200), (100, 140, 180), (180, 180, 180), (200, 150, 120), (160, 200, 140)]

    for i in range(1, 7):
        filename = f"example{i}.png"
        url = f"{base_url}{filename}"
        fallback_color = colors[i - 1]
        local_path = os.path.join(SAMPLES_DIR, filename)
        if not os.path.exists(local_path) or os.path.getsize(local_path) == 0:
            try:
                req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
                with urllib.request.urlopen(req, context=ssl_ctx, timeout=8) as resp, open(local_path, "wb") as f:
                    f.write(resp.read())
            except Exception as e:
                print(f"[startup] Notice: Sample download skipped ({filename}): {e}, generating fallback image...")
                try:
                    img = Image.new("RGB", (320, 240), color=fallback_color)
                    draw = ImageDraw.Draw(img)
                    draw.rectangle([20, 20, 300, 220], outline=(255, 255, 255), width=2)
                    img.save(local_path)
                except Exception as fe:
                    print(f"[startup] Fallback generation error: {fe}")
        if os.path.exists(local_path) and os.path.getsize(local_path) > 0:
            SAMPLE_FILES.append([local_path])

_prepare_sample_images()


def _gpu_duration(pil_image, max_new_tokens, temperature, top_k, top_p):
    return min(120, 15 + int(max_new_tokens) * 0.5)


def _sample_next_token(logits, temperature=0.7, top_k=40, top_p=0.9):
    """Samples next token using temperature scaling, top-k filtering, and top-p (nucleus) filtering."""
    if temperature <= 1e-4:
        probs = F.softmax(logits, dim=-1)
        top_prob, top_idx = torch.max(probs, dim=-1)
        return top_idx.unsqueeze(-1), top_prob.item()

    logits_scaled = logits / temperature

    if top_k > 0:
        top_k = min(top_k, logits_scaled.size(-1))
        v, _ = torch.topk(logits_scaled, top_k)
        min_topk = v[:, -1:]
        logits_scaled = torch.where(logits_scaled < min_topk, torch.full_like(logits_scaled, -float("Inf")), logits_scaled)

    if top_p < 1.0:
        sorted_logits, sorted_indices = torch.sort(logits_scaled, descending=True, dim=-1)
        sorted_probs = F.softmax(sorted_logits, dim=-1)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

        sorted_indices_to_remove = cumulative_probs > top_p
        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
        sorted_indices_to_remove[..., 0] = 0

        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
        logits_scaled = torch.where(indices_to_remove, torch.full_like(logits_scaled, -float("Inf")), logits_scaled)

    probs = F.softmax(logits_scaled, dim=-1)
    next_tok = torch.multinomial(probs, num_samples=1)
    tok_prob = probs[0, next_tok.item()].item()
    return next_tok, tok_prob


@spaces.GPU(duration=_gpu_duration)
@torch.no_grad()
def generate_caption(
    pil_image: Image.Image,
    max_new_tokens: int = 40,
    temperature: float = 0.7,
    top_k: int = 40,
    top_p: float = 0.9,
):
    """
    Caption generation for image using visual context + autoregressive sampling.
    Returns: caption, latency badge HTML, metrics grid HTML.
    """
    if pil_image is None:
        empty_metrics = _build_metrics_html(0, 0, int(max_new_tokens), int(top_k), 0.0)
        return "⚠️ Please upload an image first.", "⚡ 0 ms", empty_metrics

    t0 = time.perf_counter()
    img_t = image_tx(pil_image).unsqueeze(0).to(device)

    amp_ctx = (
        torch.amp.autocast(device_type="cuda", dtype=torch.float16)
        if device.type == "cuda"
        else torch.amp.autocast(device_type="cpu", enabled=False)
    )

    step_probs = []
    with amp_ctx:
        visual_ctx = model.encode_image(img_t)
        gen_ids = torch.full((1, 1), bos_id, dtype=torch.long, device=device)

        for _ in range(int(max_new_tokens)):
            logits = model(gen_ids, visual_context=visual_ctx)  # (1, T, V)
            last_logits = logits[0, -1, :].unsqueeze(0).float() # (1, V)
            next_tok, prob = _sample_next_token(last_logits, temperature=temperature, top_k=int(top_k), top_p=top_p)
            step_probs.append(prob)

            gen_ids = torch.cat([gen_ids, next_tok], dim=1)
            if next_tok.item() == eos_id:
                break
            if gen_ids.shape[1] >= CONTEXT_LENGTH:
                break

    latency_ms = (time.perf_counter() - t0) * 1000.0

    ids = gen_ids[0, 1:].tolist()
    if eos_id in ids:
        ids = ids[: ids.index(eos_id)]
        step_probs = step_probs[: len(ids)]

    caption = tokenizer.decode(ids, skip_special_tokens=True).strip() or "…"
    mean_conf = (sum(step_probs) / max(len(step_probs), 1)) * 100.0 if step_probs else 0.0

    latency_html = f"⚡ {latency_ms:.0f} ms"
    metrics_html = _build_metrics_html(latency_ms, len(ids), int(max_new_tokens), int(top_k), mean_conf)

    return caption, latency_html, metrics_html


def _build_metrics_html(latency_ms, token_count, max_tokens, top_k, confidence):
    conf_pct = min(100.0, max(0.0, confidence))
    return f"""
    <div class="metrics-grid">
      <div class="metric-card">
        <div class="metric-label">Confidence</div>
        <div class="metric-value">{conf_pct:.1f}%</div>
        <div class="metric-bar-bg">
          <div class="metric-bar-fill" style="width: {conf_pct:.1f}%;"></div>
        </div>
      </div>
      <div class="metric-card">
        <div class="metric-label">Top-K</div>
        <div class="metric-value">{top_k}</div>
      </div>
      <div class="metric-card">
        <div class="metric-label">Tokens</div>
        <div class="metric-value">{token_count} / {max_tokens}</div>
      </div>
      <div class="metric-card">
        <div class="metric-label">Inference Time</div>
        <div class="metric-value">{latency_ms:.0f} ms</div>
      </div>
    </div>
    """


# ── Gradio UI ─────────────────────────────────────────────────────────────────
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&family=JetBrains+Mono:wght@400;500&display=swap');

:root {
  --bg-dark: #0b0d17;
  --panel-bg: #121526;
  --panel-border: #1e243b;
  --card-bg: #181c30;
  --accent-purple: #6c5ce7;
  --accent-purple-hover: #5b4cc4;
  --accent-glow: rgba(108, 92, 231, 0.4);
  --text-main: #f1f3f9;
  --text-muted: #8c96b5;
  --badge-bg: #1e243b;
}

body, html {
  background-color: var(--bg-dark) !important;
  color: var(--text-main) !important;
  font-family: 'Inter', system-ui, -apple-system, sans-serif !important;
}

.gradio-container {
  max-width: 1240px !important;
  margin: 0 auto !important;
  padding: 20px !important;
  background: transparent !important;
}

/* ── Top Header ── */
.header-container {
  display: flex;
  align-items: center;
  justify-content: space-between;
  padding-bottom: 20px;
  margin-bottom: 20px;
  border-bottom: 1px solid var(--panel-border);
}

.header-left {
  display: flex;
  align-items: center;
  gap: 14px;
}

.logo-icon {
  width: 44px;
  height: 44px;
  border-radius: 12px;
  background: linear-gradient(135deg, #6c5ce7, #a29bfe);
  display: flex;
  align-items: center;
  justify-content: center;
  font-size: 22px;
  box-shadow: 0 4px 16px rgba(108, 92, 231, 0.3);
}

.header-title-text {
  font-size: 1.5rem;
  font-weight: 700;
  color: #ffffff;
  line-height: 1.2;
}

.header-subtitle-text {
  font-size: 0.82rem;
  color: var(--text-muted);
  margin-top: 2px;
}

.header-badges {
  display: flex;
  align-items: center;
  gap: 10px;
}

.badge-pill {
  font-family: 'JetBrains Mono', monospace;
  font-size: 0.72rem;
  padding: 5px 12px;
  border-radius: 20px;
  background: var(--badge-bg);
  border: 1px solid var(--panel-border);
  color: var(--text-main);
  display: flex;
  align-items: center;
  gap: 6px;
}

.badge-ready {
  background: rgba(34, 197, 94, 0.12);
  border-color: rgba(34, 197, 94, 0.3);
  color: #4ade80;
}

.dot-online {
  width: 7px;
  height: 7px;
  border-radius: 50%;
  background-color: #22c55e;
  box-shadow: 0 0 8px #22c55e;
}

/* ── Panel Cards ── */
.dashboard-panel {
  background: var(--panel-bg);
  border: 1px solid var(--panel-border);
  border-radius: 16px;
  padding: 18px;
  margin-bottom: 16px;
}

.panel-header {
  display: flex;
  align-items: center;
  justify-content: space-between;
  font-size: 0.95rem;
  font-weight: 600;
  color: #ffffff;
  margin-bottom: 14px;
}

.panel-header-left {
  display: flex;
  align-items: center;
  gap: 8px;
}

/* ── Image Workspace (Left) ── */
.workspace-tabs {
  display: flex;
  gap: 8px;
  margin-bottom: 14px;
}

.tab-btn-active {
  background: var(--accent-purple) !important;
  color: #ffffff !important;
  font-size: 0.8rem !important;
  font-weight: 600 !important;
  padding: 6px 14px !important;
  border-radius: 8px !important;
  border: none !important;
}

.tab-btn-inactive {
  background: transparent !important;
  color: var(--text-muted) !important;
  font-size: 0.8rem !important;
  padding: 6px 14px !important;
  border-radius: 8px !important;
  border: 1px solid transparent !important;
}

#image-uploader {
  background: var(--card-bg) !important;
  border: 1.5px dashed var(--panel-border) !important;
  border-radius: 12px !important;
  min-height: 280px !important;
  overflow: hidden !important;
}

#image-uploader img {
  max-height: 300px !important;
  object-fit: contain !important;
}

#btn-generate {
  background: linear-gradient(135deg, #6c5ce7 0%, #5b4cc4 100%) !important;
  color: #ffffff !important;
  font-size: 1rem !important;
  font-weight: 600 !important;
  border: none !important;
  border-radius: 12px !important;
  padding: 14px !important;
  margin-top: 14px !important;
  box-shadow: 0 4px 20px var(--accent-glow) !important;
  cursor: pointer !important;
  width: 100% !important;
}

#btn-generate:hover {
  background: linear-gradient(135deg, #7d6df3 0%, #6c5ce7 100%) !important;
  box-shadow: 0 6px 24px var(--accent-glow) !important;
}

/* ── Results Panel (Right) ── */
.caption-box textarea {
  background: var(--card-bg) !important;
  border: 1px solid var(--panel-border) !important;
  border-radius: 12px !important;
  color: #ffffff !important;
  font-size: 1.15rem !important;
  line-height: 1.6 !important;
  padding: 16px !important;
  min-height: 100px !important;
}

.latency-badge {
  font-family: 'JetBrains Mono', monospace;
  font-size: 0.75rem;
  background: rgba(108, 92, 231, 0.2);
  border: 1px solid rgba(108, 92, 231, 0.4);
  color: #a29bfe;
  padding: 3px 10px;
  border-radius: 12px;
}

/* Actions Row */
.caption-actions {
  display: flex;
  gap: 10px;
  margin-top: 12px;
}

.action-btn-primary {
  background: var(--accent-purple) !important;
  color: #ffffff !important;
  border: none !important;
  border-radius: 8px !important;
  padding: 8px 16px !important;
  font-size: 0.82rem !important;
  font-weight: 500 !important;
}

.action-btn-outline {
  background: transparent !important;
  color: var(--text-main) !important;
  border: 1px solid var(--panel-border) !important;
  border-radius: 8px !important;
  padding: 8px 16px !important;
  font-size: 0.82rem !important;
}

/* ── Metrics Grid ── */
.metrics-grid {
  display: grid;
  grid-template-columns: repeat(4, 1fr);
  gap: 12px;
}

.metric-card {
  background: var(--card-bg);
  border: 1px solid var(--panel-border);
  border-radius: 12px;
  padding: 14px;
}

.metric-label {
  font-size: 0.72rem;
  color: var(--text-muted);
  margin-bottom: 6px;
}

.metric-value {
  font-size: 1.2rem;
  font-weight: 700;
  color: #ffffff;
  font-family: 'Inter', sans-serif;
}

.metric-bar-bg {
  width: 100%;
  height: 4px;
  background: var(--panel-border);
  border-radius: 2px;
  margin-top: 8px;
  overflow: hidden;
}

.metric-bar-fill {
  height: 100%;
  background: linear-gradient(90deg, #6c5ce7, #a29bfe);
  border-radius: 2px;
}

/* ── Quick Examples Bar ── */
.examples-bar {
  background: var(--panel-bg);
  border: 1px solid var(--panel-border);
  border-radius: 16px;
  padding: 16px 20px;
  margin-top: 10px;
}

/* Gradio Overrides */
footer { display: none !important; }
.gradio-container .block { background: transparent !important; border: none !important; }
"""

JS = "() => document.documentElement.classList.add('dark')"
theme = gr.themes.Base(primary_hue="indigo", secondary_hue="purple", neutral_hue="slate")

num_cross = str(sorted(CFG["cross_attention_pos"]))

with gr.Blocks(theme=theme, css=CSS, js=JS, title="GPT2VL — Image Captioning") as demo:

    # ── Top Header ────────────────────────────────────────────────────────────
    gr.HTML(f"""
    <div class="header-container">
      <div class="header-left">
        <div class="logo-icon">🌌</div>
        <div>
          <div class="header-title-text">GPT2VL Image Captioning</div>
          <div class="header-subtitle-text">Vision-Language Model (GPT-2 + ViT-B/16 + Perceiver Resampler)</div>
        </div>
      </div>
      <div class="header-badges">
        <div class="badge-pill badge-ready">
          <span class="dot-online"></span> Model Ready
        </div>
        <div class="badge-pill">{CFG['num_visual_tokens']} visual tokens</div>
        <div class="badge-pill">FP32 • {str(device).upper()}</div>
        <a href="https://github.com/stackformer-labs/Stackformer" target="_blank" class="badge-pill" style="text-decoration:none;">
          <svg height="14" width="14" viewBox="0 0 16 16" fill="currentColor"><path d="M8 0C3.58 0 0 3.58 0 8c0 3.54 2.29 6.53 5.47 7.59.4.07.55-.17.55-.38 0-.19-.01-.82-.01-1.49-2.01.37-2.53-.49-2.69-.94-.09-.23-.48-.94-.82-1.13-.28-.15-.68-.52-.01-.53.63-.01 1.08.58 1.23.82.72 1.21 1.87.87 2.33.66.07-.52.28-.87.51-1.07-1.78-.2-3.64-.89-3.64-3.95 0-.87.31-1.59.82-2.15-.08-.2-.36-1.02.08-2.12 0 0 .67-.21 2.2.82.64-.18 1.32-.27 2-.27.68 0 1.36.09 2 .27 1.53-1.04 2.2-.82 2.2-.82.44 1.1.16 1.92.08 2.12.51.56.82 1.28.82 2.15 0 3.07-1.87 3.75-3.65 3.95.29.25.54.73.54 1.48 0 1.07-.01 1.93-.01 2.2 0 .21.15.46.55.38A8.013 8.013 0 0016 8c0-4.42-3.58-8-8-8z"/></svg>
        </a>
      </div>
    </div>
    """)

    with gr.Row(equal_height=False):
        # ── Left Column: Image Workspace ──────────────────────────────────────
        with gr.Column(scale=5):
            with gr.Group(elem_classes=["dashboard-panel"]):
                # Workspace Header
                gr.HTML("""
                <div class="workspace-tabs">
                  <button class="tab-btn-active">📷 Upload Image</button>
                  <button class="tab-btn-inactive">🖼 Quick Examples</button>
                </div>
                """)

                img_in = gr.Image(
                    elem_id="image-uploader",
                    type="pil",
                    label="Drag & drop an image or click to upload",
                    show_label=True,
                    sources=["upload", "clipboard"],
                    height=290,
                )

                gen_btn = gr.Button("✨ Generate Caption", elem_id="btn-generate")

        # ── Right Column: Model Results ───────────────────────────────────────
        with gr.Column(scale=7):
            # Card 1: Generated Caption
            with gr.Group(elem_classes=["dashboard-panel"]):
                with gr.Row():
                    gr.HTML('<div class="panel-header"><div class="panel-header-left">✨ Generated Caption</div></div>')
                    latency_badge = gr.HTML('<div class="latency-badge">⚡ 0 ms</div>')

                caption_out = gr.Textbox(
                    elem_classes=["caption-box"],
                    show_label=False,
                    placeholder="Your generated caption will appear here after clicking Generate...",
                    interactive=False,
                    lines=3,
                )

                gr.HTML("""
                <div class="caption-actions">
                  <button class="action-btn-primary">📋 Copy</button>
                  <button class="action-btn-outline">💾 Download</button>
                  <button class="action-btn-outline">🔗 Share</button>
                </div>
                """)

            # Card 2: Model Metrics
            with gr.Group(elem_classes=["dashboard-panel"]):
                gr.HTML('<div class="panel-header"><div class="panel-header-left">⚙️ Model Metrics</div></div>')
                metrics_out = gr.HTML(
                    _build_metrics_html(0, 0, 40, 40, 0.0)
                )

            # Card 3: Generation Parameters
            with gr.Group(elem_classes=["dashboard-panel"]):
                gr.HTML('<div class="panel-header"><div class="panel-header-left">🎛️ Generation Parameters</div></div>')

                with gr.Row():
                    sl_temp = gr.Slider(
                        minimum=0.0,
                        maximum=1.5,
                        value=0.7,
                        step=0.05,
                        label="Temperature",
                        info="0.0 = Greedy argmax",
                    )
                    sl_top_k = gr.Slider(
                        minimum=1,
                        maximum=100,
                        value=40,
                        step=1,
                        label="Top-K",
                    )
                    sl_top_p = gr.Slider(
                        minimum=0.0,
                        maximum=1.0,
                        value=0.9,
                        step=0.05,
                        label="Top-P",
                    )

                with gr.Accordion("Show Advanced Settings", open=False):
                    sl_max = gr.Slider(
                        minimum=5,
                        maximum=CONTEXT_LENGTH - 1,
                        value=40,
                        step=1,
                        label="Max New Tokens",
                    )

    # ── Bottom Row: Quick Examples ────────────────────────────────────────────
    if SAMPLE_FILES:
        with gr.Row():
            with gr.Group(elem_classes=["examples-bar"]):
                gr.HTML('<div class="panel-header"><div class="panel-header-left">💡 Quick Examples</div></div>')
                gr.Examples(
                    examples=SAMPLE_FILES,
                    inputs=[img_in],
                    label=None,
                )

    # ── Event Trigger ─────────────────────────────────────────────────────────
    gen_btn.click(
        fn=generate_caption,
        inputs=[img_in, sl_max, sl_temp, sl_top_k, sl_top_p],
        outputs=[caption_out, latency_badge, metrics_out],
        api_name="caption",
    )

if __name__ == "__main__":
    demo.queue().launch(server_name="0.0.0.0", server_port=7860)