"""
Hugging Face Space — gurumurthy3/gpt2-stackformer-vision_V2

ChatGPT-style UI using gr.MultimodalTextbox (built-in image upload).
Two generation bugs fixed:
  1. EOS / pad token banned from sampling (randomly-init rows, never trained).
  2. Only *new* tokens decoded — prompt is never echoed back.
"""

import os
import tempfile

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

from torchvision.models import vit_b_16, ViT_B_16_Weights
from torchvision import transforms
from transformers import GPT2TokenizerFast
from huggingface_hub import snapshot_download

import stackformer.modules.Attention as sf_attn
from stackformer.modules.Attention import Cross_MultiHead_Attention
from stackformer.modules.Normalization import LayerNormalization
from stackformer.modules.Feed_forward import FF_GELU
from stackformer.models.OpenAI import GPT_2

REPO_ID = os.environ.get("MODEL_REPO_ID", "gurumurthy3/gpt2-stackformer-vision_V2")
device  = torch.device("cuda" if torch.cuda.is_available() else "cpu")


# ── stackformer patch 1: inverted causal mask ─────────────────────────────────
def _causal_mask_is_correct():
    torch.manual_seed(0)
    attn = sf_attn.Multi_Head_Attention(embed_dim=8, num_heads=2, qkv_bias=True)
    attn.eval()
    x = torch.randn(1, 4, 8)
    with torch.no_grad():
        out1 = attn(x, mask=True)
        x2 = x.clone(); x2[0, -1] += 100.0
        out2 = attn(x2, mask=True)
    return torch.allclose(out1[0, 0], out2[0, 0], atol=1e-4)

def _patch_causal_masking():
    if _causal_mask_is_correct():
        return False
    def _fixed(q, k, v, attn_mask, dropout_p):
        if attn_mask is not None and attn_mask.dtype == torch.bool:
            attn_mask = ~attn_mask
        return F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask,
                                              dropout_p=dropout_p, is_causal=False)
    sf_attn._run_sdpa = _fixed
    return True

# ── stackformer patch 2: attention dropout ignores .eval() ───────────────────
def _patch_dropout():
    names = ["Self_Attention","Multi_Head_Attention","Multi_Head_Attention_With_RoPE",
             "Cross_MultiHead_Attention","Multi_query_Attention",
             "Multi_query_Attention_With_RoPE","Group_query_Attention",
             "Group_query_Attention_With_RoPE"]
    done = []
    for name in names:
        cls = getattr(sf_attn, name, None)
        if cls is None: continue
        orig = cls.forward
        def wrap(orig):
            def fwd(self, *a, **kw):
                old = self.dropout_p
                if not self.training: self.dropout_p = 0.0
                try:    return orig(self, *a, **kw)
                finally: self.dropout_p = old
            return fwd
        cls.forward = wrap(orig)
        done.append(name)
    return done

print("[startup] patching stackformer …")
_mp = _patch_causal_masking()
_dp = _patch_dropout()
print(f"[startup] causal-mask patched: {_mp} | dropout-eval patched: {len(_dp)} classes")


# ── model architecture (verbatim from training notebook) ─────────────────────
class SparseCrossAttnBlock(nn.Module):
    def __init__(self, embed_dim, num_heads, dropout, qkv_bias=True, device="cpu", dtype=torch.float32):
        super().__init__()
        self.norm      = LayerNormalization(embed_dim, device=device, dtype=dtype)
        self.cross_attn= Cross_MultiHead_Attention(embed_dim, num_heads, dropout=dropout,
                                                    qkv_bias=qkv_bias, device=device, dtype=dtype)
        self.drop      = nn.Dropout(dropout)
    def forward(self, x, context):
        return x + self.drop(self.cross_attn(self.norm(x), context, mask=False))

class PerceiverResamplerSF(nn.Module):
    def __init__(self, embed_dim, num_latents, depth, num_heads, dropout=0.0,
                 hidden_dim=None, device="cpu", dtype=torch.float32):
        super().__init__()
        hidden_dim = hidden_dim or embed_dim * 2
        self.latents    = nn.Parameter(torch.randn(num_latents, embed_dim) * 0.02)
        MK = dict(dropout=dropout, qkv_bias=True, device=device, dtype=dtype)
        self.cross_layers = nn.ModuleList([Cross_MultiHead_Attention(embed_dim, num_heads, **MK) 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):
            x = x + self.cross_layers[i](self.norm_latent[i](x), self.norm_media[i](media_seq), mask=False)
            x = x + self.ffns[i](self.ffn_norms[i](x))
        return x

class TorchvisionViTEncoder(nn.Module):
    def __init__(self, pretrained=False, freeze=True):
        super().__init__()
        weights = ViT_B_16_Weights.IMAGENET1K_V1 if pretrained else None
        self.model = vit_b_16(weights=weights)
        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)
        return self.model.encoder(x)[:, 1:, :]

class GPT2VL(nn.Module):
    def __init__(self, cfg, device="cpu", dtype=torch.float32):
        super().__init__()
        self.cfg = cfg
        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"],
                           device=device, dtype=dtype)
        self.cross_attention_pos = set(cfg["cross_attention_pos"])
        self.cross_blocks = nn.ModuleDict({
            str(i): SparseCrossAttnBlock(cfg["embed_dim"], cfg["num_heads"],
                                          cfg["dropout"], cfg["qkv_bias"], device, dtype)
            for i in cfg["cross_attention_pos"]})
        self.vision_encoder = TorchvisionViTEncoder(pretrained=False, freeze=True)
        self.vision_project = (nn.Identity() if cfg["vision_dim"] == cfg["embed_dim"]
                               else nn.Linear(cfg["vision_dim"], cfg["embed_dim"],
                                              device=device, dtype=dtype))
        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(): pts = self.vision_encoder(images)
        return self.resampler(self.vision_project(pts))
    def forward(self, input_ids, visual_context=None):
        x = self.gpt2.embedding(input_ids) + self.gpt2.position_embedding(input_ids)
        for i, layer in enumerate(self.gpt2.decoder.layers):
            x = layer(x)
            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(self.gpt2.final_norm(x))


# ── load checkpoint ───────────────────────────────────────────────────────────
print(f"[startup] downloading {REPO_ID} …")
local_dir = snapshot_download(REPO_ID)
ckpt      = torch.load(os.path.join(local_dir, "model_checkpoint.pth"), map_location=device)
CFG       = ckpt["config"]
model     = GPT2VL(CFG, device=str(device), dtype=torch.float32).to(device)
model.load_state_dict(ckpt["model_state_dict"], strict=True)
model.eval()
tokenizer = GPT2TokenizerFast.from_pretrained(os.path.join(local_dir, "tokenizer"))
print("[startup] model ready.")

bos_id         = tokenizer.bos_token_id or tokenizer.eos_token_id
eos_id         = tokenizer.eos_token_id
pad_id         = tokenizer.pad_token_id
CONTEXT_LENGTH = CFG["context_length"]
# ids that must never be sampled — untrained / special tokens
BANNED = {i for i in (eos_id, pad_id) if i is not None}

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]),
])


# ── generation ────────────────────────────────────────────────────────────────
def top_k_top_p(logits, top_k=0, top_p=1.0):
    if top_k > 0:
        vals, _ = torch.topk(logits, min(top_k, logits.size(-1)))
        logits  = logits.masked_fill(logits < vals[..., -1, None], float("-inf"))
    if top_p < 1.0:
        sl, si  = torch.sort(logits, descending=True)
        cp      = torch.softmax(sl, dim=-1).cumsum(dim=-1)
        mask    = cp > top_p
        mask[..., 1:] = mask[..., :-1].clone(); mask[..., 0] = False
        logits  = logits.scatter(1, si, sl.masked_fill(mask, float("-inf")))
    return logits

@torch.no_grad()
def generate(prompt_text, pil_image, max_new_tokens, temperature, top_p, top_k):
    model.eval()
    vc = None
    if pil_image is not None:
        img_t = image_tx(pil_image).unsqueeze(0).to(device, non_blocking=True)
        vc    = model.encode_image(img_t)

    prompt_text = (prompt_text or "").strip()
    ids = ([bos_id] + tokenizer.encode(prompt_text, add_special_tokens=False)
           if prompt_text else [bos_id])
    ids        = ids[:max(1, CONTEXT_LENGTH - 1)]
    prompt_len = len(ids)
    seq        = torch.tensor([ids], dtype=torch.long, device=device)

    for _ in range(int(max_new_tokens)):
        if seq.shape[1] >= CONTEXT_LENGTH:
            break
        logits = model(seq, visual_context=vc)[:, -1, :]
        # ── fix: ban EOS / pad — their lm_head rows were never trained ────────
        for bid in BANNED:
            if 0 <= bid < logits.shape[-1]:
                logits[:, bid] = float("-inf")
        logits  = logits / max(1e-8, float(temperature))
        logits  = top_k_top_p(logits, top_k=int(top_k), top_p=float(top_p))
        nxt     = torch.multinomial(F.softmax(logits, dim=-1), num_samples=1)
        seq     = torch.cat([seq, nxt], dim=1)

    # ── fix: return only the *new* tokens, never the prompt echo ──────────────
    new_text = tokenizer.decode(seq[0, prompt_len:].tolist(),
                                skip_special_tokens=True).strip()
    return new_text or "…"


# ── chat handler ──────────────────────────────────────────────────────────────
def chat(message, history, max_new_tokens, temperature, top_p, top_k):
    """
    message from gr.MultimodalTextbox: {"text": str, "files": [{"path":...}, ...]}
    history for gr.Chatbot type="messages": [{"role": ..., "content": ...}, ...]
    """
    user_text = (message.get("text") or "").strip()
    files     = message.get("files") or []

    pil_image = None
    img_path  = None

    if files:
        f = files[0]
        # Gradio 5: files are dicts {"path": "...", "url": "...", ...}
        img_path = f if isinstance(f, str) else f.get("path") or f.get("url", "")
        try:
            pil_image = Image.open(img_path).convert("RGB")
        except Exception:
            pil_image = None

    if not user_text and pil_image is None:
        return history

    # ── build user bubble(s) ─────────────────────────────────────────────────
    # Gradio 5 Chatbot content must be: str | {"path": str} | gr.Component
    new_msgs = []
    if pil_image is not None:
        new_msgs.append({"role": "user", "content": {"path": img_path}})
    if user_text:
        new_msgs.append({"role": "user", "content": user_text})
    elif pil_image is not None and not user_text:
        pass  # image-only is fine, no text bubble needed

    history = history + new_msgs

    # ── generate ─────────────────────────────────────────────────────────────
    reply   = generate(user_text, pil_image, max_new_tokens, temperature, top_p, top_k)
    history = history + [{"role": "assistant", "content": reply}]
    return history


# ── UI ────────────────────────────────────────────────────────────────────────
CSS = """
/* ── layout ── */
html, body, .gradio-container { height: 100% !important; }
.gradio-container { max-width: 860px !important; margin: 0 auto !important; }

/* ── header card ── */
#header-card {
    background: linear-gradient(135deg, #1e3a5f 0%, #0f2744 100%);
    border-radius: 16px;
    padding: 20px 28px 16px;
    margin-bottom: 12px;
    color: white;
}
#header-card h1 { margin: 0 0 4px; font-size: 1.45rem; font-weight: 700; }
#header-card p  { margin: 0; font-size: 0.82rem; opacity: 0.75; }
#header-card a  { color: #7ec8e3; text-decoration: none; }
#header-card a:hover { text-decoration: underline; }

/* ── chat window ── */
#chatwindow {
    border: 1.5px solid #e0e4ea !important;
    border-radius: 14px !important;
    background: #fafbfc !important;
}

/* ── input box ── */
#input-bar { margin-top: 8px; }
#input-bar .gr-form { border-radius: 14px !important; }

/* ── settings accordion ── */
#settings { margin-top: 6px; border-radius: 12px !important; }

/* ── clear button ── */
#clear-btn { margin-top: 6px; border-radius: 8px !important; }

/* ── footer note ── */
#footer-note { font-size: 0.76rem; color: #888; text-align: center; margin-top: 8px; }
footer { display: none !important; }
"""

with gr.Blocks(css=CSS, title="GPT-2 Vision Chat", fill_height=True) as demo:

    # ── header ───────────────────────────────────────────────────────────────
    gr.HTML(f"""
    <div id="header-card">
      <h1>🖼️ GPT-2 + Vision Chat</h1>
      <p>
        Model: <a href="https://huggingface.co/{REPO_ID}" target="_blank">{REPO_ID}</a>
        &nbsp;·&nbsp; GPT-2 small backbone + frozen ViT-B/16
        &nbsp;·&nbsp; {CFG['num_visual_tokens']} visual tokens
        &nbsp;·&nbsp; fine-tuned on Flickr8k
        &nbsp;·&nbsp; built on <a href="https://pypi.org/project/stackformer/" target="_blank">stackformer</a>
        &nbsp;·&nbsp; running on <b>{str(device).upper()}</b>
      </p>
    </div>
    """)

    # ── conversation state ────────────────────────────────────────────────────
    state = gr.State([])

    # ── chat window ───────────────────────────────────────────────────────────
    chatbot = gr.Chatbot(
        elem_id="chatwindow",
        label="",
        type="messages",
        height=480,
        show_copy_button=True,
        layout="bubble",
        placeholder=(
            "### 👋 Welcome!\n"
            "Type a message below, or attach an image (📎) to get a caption.\n\n"
            "_GPT-2 backbone is frozen — only the vision adapter was trained "
            "on Flickr8k. Expect short, simple captions._"
        ),
        avatar_images=(
            None,
            "https://huggingface.co/front/assets/huggingface_logo-noborder.svg",
        ),
    )

    # ── multimodal input bar ──────────────────────────────────────────────────
    msg_box = gr.MultimodalTextbox(
        elem_id="input-bar",
        placeholder="Message… attach 📎 an image to describe it",
        file_types=["image"],
        file_count="single",
        lines=1,
        max_lines=6,
        submit_btn=True,
        autofocus=True,
        show_label=False,
    )

    # ── generation settings ───────────────────────────────────────────────────
    with gr.Accordion("⚙️ Generation settings", open=False, elem_id="settings"):
        with gr.Row():
            sl_max  = gr.Slider(4, CONTEXT_LENGTH-1, value=40, step=1,
                                label="Max new tokens",
                                info="How many tokens to generate")
            sl_temp = gr.Slider(0.1, 1.5, value=0.85, step=0.05,
                                label="Temperature",
                                info="Higher = more creative / random")
        with gr.Row():
            sl_topp = gr.Slider(0.5, 1.0, value=0.9, step=0.05,
                                label="Top-p (nucleus)",
                                info="Probability mass to sample from")
            sl_topk = gr.Slider(0, 200, value=50, step=5,
                                label="Top-k",
                                info="0 = disabled")

    # ── clear ─────────────────────────────────────────────────────────────────
    gr.Button("🗑️  Clear conversation", variant="secondary",
              size="sm", elem_id="clear-btn").click(
        lambda: ([], []), outputs=[chatbot, state])

    # ── footer note ───────────────────────────────────────────────────────────
    gr.HTML('<p id="footer-note">GPT-2 backbone is frozen — only ~16.6M vision-adapter '
            'params were trained (5 h, single T4). Results are short and simple by design.</p>')

    # ── event: submit ─────────────────────────────────────────────────────────
    def _submit(message, history, max_tok, temp, topp, topk):
        new_history = chat(message, history, max_tok, temp, topp, topk)
        return new_history, new_history, {"text": "", "files": []}

    submit_args = dict(
        fn      = _submit,
        inputs  = [msg_box, state, sl_max, sl_temp, sl_topp, sl_topk],
        outputs = [chatbot, state, msg_box],
    )
    msg_box.submit(**submit_args)

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