"""Self-contained model definition for the encoder-free VLM demo.

This mirrors the inference-relevant subset of the training repo's `vlm.py`
(no wandb / datasets / matplotlib imports), so the module is light enough for a
Space while producing exactly the same state-dict keys as the checkpoint.

Architecture: a learned patch embedder (no pretrained vision encoder) projects
512x512 images into Qwen3-1.7B's hidden space; the projected patch embeddings
are spliced into the <|image|> token positions and the pretrained decoder
generates from inputs_embeds.
"""
from dataclasses import dataclass, field

import torch
import torch.nn as nn
from torchvision import transforms
from transformers import AutoTokenizer, AutoModelForCausalLM, AutoConfig

if torch.cuda.is_available():
    device = "cuda"
elif torch.backends.mps.is_available():
    device = "mps"
else:
    device = "cpu"

# Standardize arbitrary input images to the fixed 512x512 the embedder expects.
transform = transforms.Compose([
    transforms.Resize(512),
    transforms.CenterCrop(512),
    transforms.ToTensor(),
])


@dataclass
class VisionEmbedderConfig:
    img_size:            int   = 512
    patch_size:          int   = 32
    n_in_channels:       int   = 3
    out_dim:             int   = 2048
    max_patch_positions: int   = 16
    layernorm_eps:       float = 1e-6
    flat_patch_dim:      int   = field(init=False)
    num_patches:         int   = field(init=False)

    def __post_init__(self):
        self.flat_patch_dim = self.patch_size ** 2 * self.n_in_channels
        self.num_patches    = (self.img_size // self.patch_size) ** 2


@dataclass
class DecoderConfig:
    checkpoint:  str = "Qwen/Qwen3-1.7B"
    image_token: str = "<|image|>"
    hidden_size: int = field(init=False)

    def __post_init__(self):
        cfg = AutoConfig.from_pretrained(self.checkpoint)
        self.hidden_size = cfg.hidden_size


@dataclass
class VLMConfig:
    vision:  VisionEmbedderConfig = field(default_factory=VisionEmbedderConfig)
    decoder: DecoderConfig        = field(default_factory=DecoderConfig)


def extract_flattened_patches(x: torch.Tensor, patch_size: int) -> torch.Tensor:
    """Split (B, C, H, W) into non-overlapping P x P patches: (B, N, C*P*P)."""
    B, C, H, W = x.shape
    P = patch_size
    nh, nw = H // P, W // P
    x = x.reshape(B, C, nh, P, nw, P)
    x = x.permute(0, 2, 4, 1, 3, 5)
    x = x.reshape(B, nh * nw, C * P * P)
    return x


class VisionEmbedder(nn.Module):
    def __init__(self, cfg: VisionEmbedderConfig):
        super().__init__()
        self.cfg = cfg
        self.ln1 = nn.LayerNorm(cfg.flat_patch_dim, eps=cfg.layernorm_eps)
        self.fc  = nn.Linear(cfg.flat_patch_dim, cfg.out_dim)
        self.ln2 = nn.LayerNorm(cfg.out_dim, eps=cfg.layernorm_eps)
        self.x_pos_emb = nn.Parameter(torch.randn(1, cfg.max_patch_positions, cfg.out_dim))
        self.y_pos_emb = nn.Parameter(torch.randn(1, cfg.max_patch_positions, cfg.out_dim))
        self.ln3 = nn.LayerNorm(cfg.out_dim, eps=cfg.layernorm_eps)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        B, C, H, W = x.shape
        P  = self.cfg.patch_size
        nh, nw = H // P, W // P
        assert nh <= self.cfg.max_patch_positions and nw <= self.cfg.max_patch_positions, \
            f"Patch grid {nh}x{nw} exceeds max_patch_positions={self.cfg.max_patch_positions}"

        x = extract_flattened_patches(x, P)
        x = self.ln1(x)
        x = self.fc(x)
        x = self.ln2(x)

        row_emb = self.y_pos_emb[:, :nh, :]
        col_emb = self.x_pos_emb[:, :nw, :]
        pos = (row_emb.unsqueeze(2) + col_emb.unsqueeze(1)).reshape(1, nh * nw, -1)
        x = x + pos

        x = self.ln3(x)
        return x


def replace_img_tokens_with_embd(input_ids, combined_embd, image_embd, image_token_id):
    """Overwrite <|image|> token positions with projected patch embeddings."""
    mask   = (input_ids == image_token_id)
    result = combined_embd.clone()
    result[mask] = image_embd.reshape(-1, image_embd.shape[-1]).to(result.dtype)
    return result


class VLM(nn.Module):
    def __init__(self, cfg: VLMConfig):
        super().__init__()
        self.cfg = cfg
        self.vision_embedder = VisionEmbedder(cfg.vision)
        self.connector = nn.Linear(cfg.vision.out_dim, cfg.decoder.hidden_size, bias=False)
        self.tokenizer = AutoTokenizer.from_pretrained(
            cfg.decoder.checkpoint,
            extra_special_tokens={"image_token": cfg.decoder.image_token},
        )
        self.decoder = AutoModelForCausalLM.from_pretrained(cfg.decoder.checkpoint)
        self.decoder.resize_token_embeddings(len(self.tokenizer))

    @property
    def image_token_id(self):
        return self.tokenizer.image_token_id

    def _replace_img_tokens_with_embd(self, input_ids, combined_embd, image_embd):
        return replace_img_tokens_with_embd(input_ids, combined_embd, image_embd, self.image_token_id)

    def forward(self, input_ids: torch.Tensor, image: torch.Tensor) -> torch.Tensor:
        image_embd           = self.vision_embedder(image)
        projected_image_embd = self.connector(image_embd)
        token_embd = self.decoder.get_input_embeddings()(input_ids)
        combined_embd = self._replace_img_tokens_with_embd(
            input_ids, token_embd, projected_image_embd
        )
        return self.decoder(inputs_embeds=combined_embd).logits