"""Time each component of vision tower forward to find the actual bottleneck."""
import sys, time
sys.path.insert(0, ".")

import torch
import torch.nn.functional as F
from peft import PeftModel
from transformers import AutoModelForImageTextToText, AutoProcessor

from training.Policy.policy_dataset import PolicyDataset, _load_frames
from training.Policy import make_cot_belief_cache as M


def main():
    print("=" * 70)
    print("Per-component timing of vision tower forward")
    print("=" * 70)
    proc = AutoProcessor.from_pretrained(
        "checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best")
    ds = PolicyDataset(
        manifests=["data/policy_labels/val.json"],
        split="val", n_frames=8, sampling="last_biased", source_filter="all",
    )
    all_imgs = [
        _load_frames(ds.samples[i]["source_dir"],
                      ds.samples[i]["frame_indices"], n_frames=8)
        for i in range(8)
    ]

    print("\n[load]")
    model = AutoModelForImageTextToText.from_pretrained(
        "models/Qwen3-VL-4B-Instruct",
        dtype=torch.bfloat16,
        attn_implementation="sdpa",
    )
    model.resize_token_embeddings(151674)
    model = PeftModel.from_pretrained(
        model, "checkpoints/VLA/qwen3vl4b_cot_belief_perframe/best"
    ).merge_and_unload()
    model.cuda().eval()

    # ALL submodule devices
    print("\n[device check] ALL submodules of vision tower:")
    cpu_modules = []
    for name, mod in model.visual.named_modules():
        try:
            ps = list(mod.parameters(recurse=False))
            if not ps:
                continue
            d = ps[0].device
            t = ps[0].dtype
            if d.type != "cuda":
                cpu_modules.append((name, str(d), str(t)))
        except Exception:
            pass
    if cpu_modules:
        print(f"   ⚠️ {len(cpu_modules)} submodules NOT on cuda:")
        for n, d, t in cpu_modules[:10]:
            print(f"      {n}  {d}  {t}")
    else:
        print("   ✓ all on cuda")

    # benchmark vision tower with bs=1
    print("\n[prep inputs bs=1]")
    inputs = M._build_inputs(proc, [all_imgs[0]], [{}], resize_short=336)
    pv = inputs["pixel_values"].cuda().to(torch.bfloat16)
    grid_thw = inputs["image_grid_thw"].cuda()
    print(f"   pixel_values: {tuple(pv.shape)}")
    print(f"   grid_thw: {tuple(grid_thw.shape)}, values:\n{grid_thw}")

    vt = model.visual
    n_blocks = len(list(vt.blocks))
    print(f"   vision tower has {n_blocks} blocks")

    # ── component-wise timing ──
    with torch.no_grad():
        torch.cuda.synchronize(); t0 = time.time()
        h = vt.patch_embed(pv)
        torch.cuda.synchronize(); print(f"   patch_embed: {(time.time()-t0)*1000:.1f} ms,  shape={tuple(h.shape)}")

        t0 = time.time()
        pos_embeds = vt.fast_pos_embed_interpolate(grid_thw)
        torch.cuda.synchronize(); print(f"   pos_embed_interpolate: {(time.time()-t0)*1000:.1f} ms")
        h = h + pos_embeds

        t0 = time.time()
        rope = vt.rot_pos_emb(grid_thw)
        torch.cuda.synchronize(); print(f"   rot_pos_emb: {(time.time()-t0)*1000:.1f} ms")

        seq_len = h.size(0)
        h = h.reshape(seq_len, -1)
        rope = rope.reshape(seq_len, -1)
        emb = torch.cat((rope, rope), dim=-1)
        position_embeddings = (emb.cos(), emb.sin())

        cu_seqlens = torch.repeat_interleave(
            grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
        ).cumsum(dim=0, dtype=torch.int32)
        cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

        # time each block
        block_times = []
        for i, blk in enumerate(vt.blocks):
            torch.cuda.synchronize()
            t0 = time.time()
            h = blk(h, cu_seqlens=cu_seqlens,
                    position_embeddings=position_embeddings)
            torch.cuda.synchronize()
            t = (time.time() - t0) * 1000
            block_times.append(t)
            if i < 3 or i == n_blocks - 1:
                print(f"   block[{i}]: {t:.1f} ms")
        print(f"   block 0-2 mean: {sum(block_times[:3])/3:.1f} ms")
        print(f"   block ALL mean: {sum(block_times)/len(block_times):.1f} ms")
        print(f"   block ALL total: {sum(block_times):.1f} ms")

        torch.cuda.synchronize(); t0 = time.time()
        out = vt.merger(h)
        torch.cuda.synchronize(); print(f"   merger: {(time.time()-t0)*1000:.1f} ms")

    # also benchmark a single attn vs MLP within block 0
    print("\n[zoom: block[0] attn vs mlp]")
    with torch.no_grad():
        blk = vt.blocks[0]
        h_in = h.detach().clone().requires_grad_(False)
        torch.cuda.synchronize(); t0 = time.time()
        for _ in range(3):
            ho = blk.attn(blk.norm1(h_in), cu_seqlens=cu_seqlens,
                          position_embeddings=position_embeddings)
        torch.cuda.synchronize()
        print(f"   attn (3 reps): {(time.time()-t0)*1000:.1f} ms total = {(time.time()-t0)/3*1000:.1f} ms/call")

        torch.cuda.synchronize(); t0 = time.time()
        for _ in range(3):
            mo = blk.mlp(blk.norm2(h_in))
        torch.cuda.synchronize()
        print(f"   mlp  (3 reps): {(time.time()-t0)*1000:.1f} ms total = {(time.time()-t0)/3*1000:.1f} ms/call")


if __name__ == "__main__":
    main()