#!/usr/bin/env python3
"""
Profile VLM vs LLM training pipeline speed & VRAM.

Prompt: 之前的llm的script跑的飞快，j63_small_fastdllmv2慢到吐血，
profiling VLM vs LLM forward pass to identify bottlenecks.

Usage:
  python3 chengyue_script/dvlm/utils/profile_vlm_vs_llm.py
"""

import os
import sys
import time
import json
import torch
import gc

# Ensure project paths
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
sys.path.insert(0, os.path.join(PROJECT_ROOT, "src"))
sys.path.insert(0, os.path.join(PROJECT_ROOT, "third_party", "vila"))
sys.path.insert(0, os.path.join(PROJECT_ROOT, "examples"))

def get_gpu_mem():
    """Get current GPU memory usage in MB."""
    if torch.cuda.is_available():
        return torch.cuda.memory_allocated() / 1024**2, torch.cuda.max_memory_allocated() / 1024**2
    return 0, 0

def profile_llm_forward(model_path, bd_size=32, seq_len=2048, batch_size=1, warmup=2, runs=5):
    """Profile LLM-only forward pass (text-only path in modeling.py)."""
    from transformers import AutoTokenizer
    from lmflow.models.qwen2_5_vl_mdm.modeling import Fast_dLLM_Qwen2_5_VLForConditionalGeneration
    from lmflow.models.qwen2_5_vl_mdm.configuration import Fast_dLLM_Qwen2_5_VLConfig
    
    print(f"\n{'='*60}")
    print(f"Profiling LLM-only forward (no pixel_values)")
    print(f"  model: {model_path}")
    print(f"  bd_size={bd_size}, seq_len={seq_len}, batch={batch_size}")
    print(f"{'='*60}")
    
    torch.cuda.reset_peak_memory_stats()
    gc.collect()
    torch.cuda.empty_cache()
    
    mem_before = get_gpu_mem()
    print(f"  GPU mem before load: {mem_before[0]:.0f} MB allocated, {mem_before[1]:.0f} MB peak")
    
    config = Fast_dLLM_Qwen2_5_VLConfig.from_pretrained(model_path, trust_remote_code=True)
    config.bd_size = bd_size
    config.complementary_mask = True
    config.always_mask_im_end = True
    config.use_block_causal_mask = True
    config.block_causal_no_dynamic = True
    
    model = Fast_dLLM_Qwen2_5_VLForConditionalGeneration.from_pretrained(
        model_path, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True
    ).cuda()
    model.train()
    model.gradient_checkpointing_enable()
    
    mem_after_load = get_gpu_mem()
    print(f"  GPU mem after load: {mem_after_load[0]:.0f} MB allocated, {mem_after_load[1]:.0f} MB peak")
    
    # Pad seq_len to bd_size multiple
    if seq_len % bd_size != 0:
        seq_len = seq_len + (bd_size - seq_len % bd_size)
    
    # Create dummy text-only batch (no pixel_values → triggers LLM path)
    input_ids = torch.randint(0, 150000, (batch_size, seq_len), device="cuda")
    labels = input_ids.clone()
    # Mask first half as prompt
    labels[:, :seq_len//2] = -100
    attention_mask = torch.ones_like(input_ids)
    
    mask_id = 151665  # |<MASK>|
    
    times = []
    for i in range(warmup + runs):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        
        output = model(
            input_ids=input_ids,
            labels=labels,
            attention_mask=attention_mask,
            pixel_values=None,
            image_grid_thw=None,
            mask_id=mask_id,
            num_items_in_batch=batch_size,
        )
        loss = output.loss
        loss.backward()
        
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        
        if i >= warmup:
            times.append(t1 - t0)
            print(f"  Run {i-warmup+1}: {t1-t0:.3f}s, loss={loss.item():.4f}")
        else:
            print(f"  Warmup {i+1}: {t1-t0:.3f}s")
        
        model.zero_grad()
    
    mem_after_run = get_gpu_mem()
    avg_time = sum(times) / len(times) if times else 0
    
    print(f"\n  LLM-only results:")
    print(f"    Avg time: {avg_time:.3f}s")
    print(f"    GPU peak: {mem_after_run[1]:.0f} MB")
    print(f"    Tokens/sec: {batch_size * seq_len / avg_time:.0f}")
    
    del model, output
    gc.collect()
    torch.cuda.empty_cache()
    
    return {"avg_time": avg_time, "peak_mem_mb": mem_after_run[1], "tps": batch_size * seq_len / avg_time if avg_time > 0 else 0}


def profile_vlm_forward(model_path, bd_size=32, seq_len=2048, batch_size=1, warmup=2, runs=5):
    """Profile VLM forward pass (with pixel_values → VLM path in modeling.py)."""
    from transformers import AutoTokenizer, AutoProcessor
    from lmflow.models.qwen2_5_vl_mdm.modeling import Fast_dLLM_Qwen2_5_VLForConditionalGeneration
    from lmflow.models.qwen2_5_vl_mdm.configuration import Fast_dLLM_Qwen2_5_VLConfig
    
    print(f"\n{'='*60}")
    print(f"Profiling VLM forward (with pixel_values)")
    print(f"  model: {model_path}")
    print(f"  bd_size={bd_size}, seq_len={seq_len}, batch={batch_size}")
    print(f"{'='*60}")
    
    torch.cuda.reset_peak_memory_stats()
    gc.collect()
    torch.cuda.empty_cache()
    
    mem_before = get_gpu_mem()
    print(f"  GPU mem before load: {mem_before[0]:.0f} MB allocated, {mem_before[1]:.0f} MB peak")
    
    config = Fast_dLLM_Qwen2_5_VLConfig.from_pretrained(model_path, trust_remote_code=True)
    config.bd_size = bd_size
    config.complementary_mask = True
    config.always_mask_im_end = True
    config.use_block_causal_mask = True
    config.block_causal_no_dynamic = False  # VLM uses dynamic
    
    model = Fast_dLLM_Qwen2_5_VLForConditionalGeneration.from_pretrained(
        model_path, config=config, torch_dtype=torch.bfloat16, trust_remote_code=True
    ).cuda()
    model.train()
    model.gradient_checkpointing_enable()
    
    mem_after_load = get_gpu_mem()
    print(f"  GPU mem after load: {mem_after_load[0]:.0f} MB allocated, {mem_after_load[1]:.0f} MB peak")
    
    # Pad seq_len to bd_size multiple
    if seq_len % bd_size != 0:
        seq_len = seq_len + (bd_size - seq_len % bd_size)
    
    # Simulate a VLM batch: we need pixel_values and image_grid_thw
    # Use image_token_id (151655) for image placeholders
    image_token_id = 151655
    # Qwen2.5-VL: spatial_merge_size=2, so 14x14 grid → 7*7=49 merged tokens in LLM space
    # But ViT input needs full patches: 14*14 = 196 patches, each with 3*14*14=1176 dims (temporal=1)
    n_image_tokens = 49  # After merge: (14/2)*(14/2) = 49 tokens in the LLM sequence
    n_raw_patches = 196  # Raw ViT patches before merge: 14*14 = 196
    
    input_ids = torch.randint(0, 150000, (batch_size, seq_len), device="cuda")
    # Place image tokens in first part
    input_ids[:, 10:10+n_image_tokens] = image_token_id
    labels = input_ids.clone()
    labels[:, :seq_len//2] = -100
    attention_mask = torch.ones_like(input_ids)
    
    # Fake pixel_values: shape [total_patches, temporal*channel*patch_h*patch_w]
    # Qwen2.5-VL uses temporal=2, channel=3, patch=14x14 → 2*3*14*14 = 1176
    pixel_values = torch.randn(n_raw_patches, 1176, device="cuda", dtype=torch.bfloat16)
    image_grid_thw = torch.tensor([[1, 14, 14]], device="cuda", dtype=torch.long)
    
    mask_id = 151665
    
    times = []
    for i in range(warmup + runs):
        torch.cuda.synchronize()
        t0 = time.perf_counter()
        
        output = model(
            input_ids=input_ids,
            labels=labels,
            attention_mask=attention_mask,
            pixel_values=pixel_values,
            image_grid_thw=image_grid_thw,
            mask_id=mask_id,
            num_items_in_batch=batch_size,
        )
        loss = output.loss
        loss.backward()
        
        torch.cuda.synchronize()
        t1 = time.perf_counter()
        
        if i >= warmup:
            times.append(t1 - t0)
            print(f"  Run {i-warmup+1}: {t1-t0:.3f}s, loss={loss.item():.4f}")
        else:
            print(f"  Warmup {i+1}: {t1-t0:.3f}s")
        
        model.zero_grad()
    
    mem_after_run = get_gpu_mem()
    avg_time = sum(times) / len(times) if times else 0
    
    print(f"\n  VLM results:")
    print(f"    Avg time: {avg_time:.3f}s")
    print(f"    GPU peak: {mem_after_run[1]:.0f} MB")
    print(f"    Tokens/sec: {batch_size * seq_len / avg_time:.0f}")
    
    del model, output
    gc.collect()
    torch.cuda.empty_cache()
    
    return {"avg_time": avg_time, "peak_mem_mb": mem_after_run[1], "tps": batch_size * seq_len / avg_time if avg_time > 0 else 0}


def main():
    # Try multiple model paths
    model_candidates = [
        os.path.join(PROJECT_ROOT, "src/models/Qwen2.5-VL-3B-Instruct-mdm"),
        "/lustre/fs1/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/LMFlow/src/models/Qwen2.5-VL-3B-Instruct-mdm",
    ]
    
    model_path = None
    for p in model_candidates:
        if os.path.isdir(p):
            model_path = p
            break
    
    if model_path is None:
        print("ERROR: No MDM model checkpoint found. Tried:")
        for p in model_candidates:
            print(f"  {p}")
        sys.exit(1)
    
    print(f"Using model: {model_path}")
    print(f"CUDA: {torch.cuda.is_available()}, Device: {torch.cuda.get_device_name(0)}")
    
    # Profile LLM path
    llm_results = profile_llm_forward(model_path, bd_size=32, seq_len=2048, batch_size=1, warmup=1, runs=3)
    
    # Profile VLM path
    vlm_results = profile_vlm_forward(model_path, bd_size=32, seq_len=2048, batch_size=1, warmup=1, runs=3)
    
    # Summary
    print(f"\n{'='*60}")
    print(f"COMPARISON SUMMARY")
    print(f"{'='*60}")
    print(f"{'Metric':<25} {'LLM-only':<20} {'VLM':<20} {'Ratio':<10}")
    print(f"{'-'*75}")
    
    ratio_time = vlm_results['avg_time'] / llm_results['avg_time'] if llm_results['avg_time'] > 0 else 0
    ratio_mem = vlm_results['peak_mem_mb'] / llm_results['peak_mem_mb'] if llm_results['peak_mem_mb'] > 0 else 0
    
    print(f"{'Avg fwd+bwd time (s)':<25} {llm_results['avg_time']:<20.3f} {vlm_results['avg_time']:<20.3f} {ratio_time:<10.2f}x")
    print(f"{'Peak GPU mem (MB)':<25} {llm_results['peak_mem_mb']:<20.0f} {vlm_results['peak_mem_mb']:<20.0f} {ratio_mem:<10.2f}x")
    print(f"{'Tokens/sec':<25} {llm_results['tps']:<20.0f} {vlm_results['tps']:<20.0f}")
    print(f"{'='*60}")


if __name__ == "__main__":
    main()
