"""Run make_cot_belief_cache with patched Qwen3VLVisionPatchEmbed.

Replaces the Conv3d patch projection with an equivalent Linear layer (math
identical, but ~64× faster because of a cuDNN slow-path bug for tiny Conv3d
on bf16). Saves whole-cache time from ~6 days to ~2 hours.

Usage: identical to make_cot_belief_cache, just call this instead.

  python tools/run_qwen3_cache_fast.py \\
      --ckpt_dir checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best \\
      --base_model models/Qwen3-VL-4B-Instruct \\
      --split val \\
      --out data/belief_cache_perframe_qwen3vl4b/multisrc_val.pt \\
      --n_frames 8 --sampling last_biased --source_filter all \\
      --batch_size 8 --num_workers 4 --chunk_size 2000
"""
import sys
sys.path.insert(0, ".")

import torch
import torch.nn as nn
from transformers.models.qwen3_vl.modeling_qwen3_vl import Qwen3VLVisionPatchEmbed


# ─── Lazy-replacement: first forward call replaces Conv3d with Linear ─────

_PATCH_APPLIED = {}


def _fast_patch_embed_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
    """Mathematically equivalent to original Conv3d-based forward, but
    routes through nn.Linear (which avoids the cuDNN slow-path bug on tiny
    Conv3d inputs)."""
    target_dtype = self.proj.weight.dtype

    # First call on this instance: convert Conv3d → Linear in place.
    if isinstance(self.proj, nn.Conv3d):
        conv = self.proj
        out_dim = conv.out_channels
        in_dim = (conv.in_channels * conv.kernel_size[0]
                  * conv.kernel_size[1] * conv.kernel_size[2])
        # Conv3d weight: (out, in, k_t, k_h, k_w) → flatten last 4 dims
        w_flat = conv.weight.detach().reshape(out_dim, in_dim).contiguous()
        bias = conv.bias.detach().clone() if conv.bias is not None else None
        new_proj = nn.Linear(in_dim, out_dim, bias=bias is not None)
        new_proj.weight.data.copy_(w_flat)
        if bias is not None:
            new_proj.bias.data.copy_(bias)
        new_proj.to(device=conv.weight.device, dtype=conv.weight.dtype)
        self.proj = new_proj
        if id(self) not in _PATCH_APPLIED:
            _PATCH_APPLIED[id(self)] = True
            print(f"[fast_patch] patched Qwen3VLVisionPatchEmbed @ id={id(self)}: "
                   f"Conv3d({in_dim}→{out_dim}) → Linear({in_dim}→{out_dim})",
                   flush=True)

    # Now self.proj is nn.Linear. Input may be (N, 1536) flat or (N, 3, 2, 16, 16).
    if hidden_states.dim() > 2 or hidden_states.shape[-1] != self.proj.in_features:
        hidden_states = hidden_states.reshape(-1, self.proj.in_features)
    hidden_states = hidden_states.to(dtype=target_dtype)
    return self.proj(hidden_states)


# Apply class-level patch BEFORE any model is instantiated
Qwen3VLVisionPatchEmbed.forward = _fast_patch_embed_forward
print("[fast_patch] Qwen3VLVisionPatchEmbed.forward replaced "
       "(lazy Conv3d → Linear conversion)", flush=True)


# Hand off to the original cache builder ───────────────────────────────────
from training.Policy import make_cot_belief_cache  # noqa: E402

if __name__ == "__main__":
    sys.exit(make_cot_belief_cache.main())