"""Differentiable loader for the Qwen3-VL vision tower (Qwen/Qwen3-VL-8B-Instruct).

Qwen3-VL's vision model consumes pre-patchified `hidden_states` (seq, patch_dim) plus
a `grid_thw` tensor, not a plain image. We reproduce the Qwen image-processor patchify
in pure torch so gradients flow back to pixels, then mean-pool the last hidden state.

The full checkpoint stores the vision tower under the `model.visual.` prefix, so we
load the shards manually and strip that prefix into a standalone Qwen3VLVisionModel
(a plain from_pretrained(REPO) leaves every vision weight randomly initialized).
"""
from __future__ import annotations

import glob
import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from huggingface_hub import snapshot_download
from safetensors.torch import load_file
from transformers import AutoConfig, Qwen3VLVisionModel

REPO = "Qwen/Qwen3-VL-8B-Instruct"
PREFIX = "model.visual."
PATCH = 16
TEMPORAL = 2
MERGE = 2
DEFAULT = 384   # multiple of PATCH*MERGE = 32


def load_qwen3vl(dtype: torch.dtype = torch.float32) -> nn.Module:
    # fp16/bf16 backward through this tower produces NaN grads (attention/norm overflow),
    # so we pin fp32 regardless of the ensemble's default custom-tower dtype.
    dtype = torch.float32
    cfg = AutoConfig.from_pretrained(REPO)
    vcfg = cfg.vision_config
    # Real (not meta) init so rotary/position buffers get materialized; we only pay a
    # one-time random-init of a 0.6B tower, then overwrite params from the checkpoint.
    model = Qwen3VLVisionModel(vcfg).to(dtype)

    root = Path(snapshot_download(REPO, allow_patterns=["*.safetensors", "*.json"]))
    idx = root / "model.safetensors.index.json"
    shards = (set(json.load(open(idx))["weight_map"].values())
              if idx.exists() else {p.name for p in root.glob("*.safetensors")})
    state: dict[str, torch.Tensor] = {}
    for shard in shards:
        for k, v in load_file(root / shard).items():
            if k.startswith(PREFIX):
                state[k[len(PREFIX):]] = v.to(dtype)

    missing, unexpected = model.load_state_dict(state, strict=False, assign=False)
    real_missing = [m for m in missing if "rotary" not in m and "inv_freq" not in m]
    if real_missing:
        raise RuntimeError(f"Qwen3-VL vision still missing {len(real_missing)} weights: "
                           f"{real_missing[:6]}")
    model.eval()
    model.requires_grad_(False)
    model._veil_cfg = vcfg
    return model.to(device="cuda", dtype=dtype)


def _patchify(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """x: (1,3,H,W) [0,1] -> (seq, 3*TEMPORAL*PATCH*PATCH) flattened patches + grid_thw.

    Mirrors Qwen2/3-VL image-processor reshape so token order matches the ViT.
    """
    _, c, h, w = x.shape
    gh, gw = h // PATCH, w // PATCH
    xt = x[0].unsqueeze(0).repeat(TEMPORAL, 1, 1, 1)          # (TEMPORAL, C, H, W)
    p = xt.reshape(1, TEMPORAL, c,
                   gh // MERGE, MERGE, PATCH,
                   gw // MERGE, MERGE, PATCH)
    p = p.permute(0, 3, 6, 4, 7, 2, 1, 5, 8).contiguous()
    flat = p.reshape(gh * gw, c * TEMPORAL * PATCH * PATCH)
    grid = torch.tensor([[1, gh, gw]], device=x.device, dtype=torch.long)
    return flat, grid


def image_feat_qwen3vl(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
    """x: (1,3,H,W) in [0,1] on cuda, requires_grad=True. Returns (1, D) pooled feature."""
    assert x.shape[0] == 1 and x.dim() == 4
    if x.shape[-1] != DEFAULT or x.shape[-2] != DEFAULT:
        x = F.interpolate(x, size=(DEFAULT, DEFAULT), mode="bicubic", align_corners=False)
    dtype = next(model.parameters()).dtype
    flat, grid = _patchify(x)
    out = model(flat.to(dtype=dtype), grid_thw=grid)
    hs = out.last_hidden_state if hasattr(out, "last_hidden_state") else (
        out[0] if isinstance(out, (tuple, list)) else out)
    return hs.float().mean(dim=0, keepdim=True)