#!/usr/bin/env python3
"""
Profile data loading pipeline: VILA dataset → DataCollatorForQwenVL.

Prompt: 先搞一下unittest，保证都过了，然后去做profiling

Measures:
1. VILA dataset __getitem__ latency (text-only vs image)
2. DataCollatorForQwenVL batch collation latency
3. process_vision_info + processor call latency

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

import os
import sys
import time
import statistics

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"))

import torch
from transformers import AutoProcessor, AutoTokenizer


def profile_collator_text_only(processor, tokenizer, n_samples=50, batch_size=1, bd_size=32):
    """Profile DataCollatorForQwenVL with text-only instances (no images)."""
    from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL

    collator = DataCollatorForQwenVL(
        processor=processor,
        tokenizer=tokenizer,
        bd_size=bd_size,
        max_length=None,
    )

    # Create text-only instances (simulating fast_dllmv2 / shizhe data)
    instances = []
    for i in range(n_samples):
        text_len = 50 + (i * 7) % 200  # Varying lengths
        instances.append({
            "messages": [
                {"role": "user", "content": [{"type": "text", "text": f"Question {i}: " + "word " * text_len}]},
                {"role": "assistant", "content": [{"type": "text", "text": f"Answer {i}: " + "reply " * (text_len // 2)}]},
            ]
        })

    # Warmup
    _ = collator(instances[:batch_size])

    # Profile individual collate calls
    times = []
    for start in range(0, min(n_samples, 20), batch_size):
        batch_instances = instances[start:start + batch_size]
        t0 = time.perf_counter()
        batch = collator(batch_instances)
        t1 = time.perf_counter()
        times.append(t1 - t0)

    avg = statistics.mean(times)
    med = statistics.median(times)
    p95 = sorted(times)[int(len(times) * 0.95)] if len(times) > 1 else times[0]

    print(f"\n  Text-only collator (batch_size={batch_size}, bd_size={bd_size}):")
    print(f"    Avg: {avg*1000:.1f}ms  Median: {med*1000:.1f}ms  P95: {p95*1000:.1f}ms")
    print(f"    Batch shape: input_ids={batch['input_ids'].shape}")
    return {"avg_ms": avg * 1000, "median_ms": med * 1000, "p95_ms": p95 * 1000}


def profile_collator_with_image(processor, tokenizer, n_samples=10, batch_size=1, bd_size=32):
    """Profile DataCollatorForQwenVL with image instances."""
    from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL
    from PIL import Image
    import tempfile

    collator = DataCollatorForQwenVL(
        processor=processor,
        tokenizer=tokenizer,
        bd_size=bd_size,
        max_length=2500,
    )

    # Create a temp image file
    tmp_dir = tempfile.mkdtemp()
    img_path = os.path.join(tmp_dir, "test.jpg")
    img = Image.new("RGB", (448, 448), color=(128, 128, 128))
    img.save(img_path)

    instances = []
    for i in range(n_samples):
        instances.append({
            "messages": [
                {"role": "user", "content": [
                    {"type": "text", "text": f"Question {i}: Describe this image."},
                    {"type": "image", "image": f"file://{img_path}"},
                ]},
                {"role": "assistant", "content": [{"type": "text", "text": f"Answer {i}: This is a gray image. " + "detail " * 50}]},
            ]
        })

    # Warmup
    _ = collator(instances[:batch_size])

    # Profile
    times = []
    for start in range(0, min(n_samples, 10), batch_size):
        batch_instances = instances[start:start + batch_size]
        t0 = time.perf_counter()
        batch = collator(batch_instances)
        t1 = time.perf_counter()
        times.append(t1 - t0)

    avg = statistics.mean(times)
    med = statistics.median(times)
    p95 = sorted(times)[int(len(times) * 0.95)] if len(times) > 1 else times[0]

    print(f"\n  Image collator (batch_size={batch_size}, bd_size={bd_size}):")
    print(f"    Avg: {avg*1000:.1f}ms  Median: {med*1000:.1f}ms  P95: {p95*1000:.1f}ms")
    print(f"    Batch shape: input_ids={batch['input_ids'].shape}")
    if "pixel_values" in batch:
        print(f"    pixel_values shape: {batch['pixel_values'].shape}")

    # Cleanup
    os.unlink(img_path)
    os.rmdir(tmp_dir)

    return {"avg_ms": avg * 1000, "median_ms": med * 1000, "p95_ms": p95 * 1000}


def profile_processor_overhead(processor):
    """Profile the processor call itself (tokenize + image process)."""
    from PIL import Image

    # Text-only
    text = "<|im_start|>user\n" + "Hello world. " * 100 + "<|im_end|>\n<|im_start|>assistant\n" + "Response. " * 50 + "<|im_end|>\n"

    times_text = []
    for _ in range(20):
        t0 = time.perf_counter()
        _ = processor(text=[text], padding=False, return_tensors="pt")
        t1 = time.perf_counter()
        times_text.append(t1 - t0)

    # With image
    img = Image.new("RGB", (448, 448), color=(64, 64, 64))
    times_img = []
    for _ in range(20):
        t0 = time.perf_counter()
        _ = processor(text=[text], images=[[img]], padding=False, return_tensors="pt")
        t1 = time.perf_counter()
        times_img.append(t1 - t0)

    avg_text = statistics.mean(times_text) * 1000
    avg_img = statistics.mean(times_img) * 1000

    print(f"\n  Processor call overhead:")
    print(f"    Text-only: {avg_text:.1f}ms avg")
    print(f"    With image: {avg_img:.1f}ms avg")
    print(f"    Image overhead: +{avg_img - avg_text:.1f}ms ({avg_img/avg_text:.1f}x)")
    return {"text_ms": avg_text, "image_ms": avg_img}


def profile_process_vision_info():
    """Profile qwen_vl_utils.process_vision_info."""
    from qwen_vl_utils import process_vision_info
    from PIL import Image
    import tempfile

    # Create temp image
    tmp_dir = tempfile.mkdtemp()
    img_path = os.path.join(tmp_dir, "test.jpg")
    img = Image.new("RGB", (448, 448), color=(128, 128, 128))
    img.save(img_path)

    # Text-only messages
    text_msgs = [
        {"role": "user", "content": [{"type": "text", "text": "Hello"}]},
        {"role": "assistant", "content": [{"type": "text", "text": "Hi"}]},
    ]

    # Image messages
    img_msgs = [
        {"role": "user", "content": [
            {"type": "text", "text": "Describe"},
            {"type": "image", "image": f"file://{img_path}"},
        ]},
        {"role": "assistant", "content": [{"type": "text", "text": "A gray image."}]},
    ]

    # Profile text-only
    times_text = []
    for _ in range(50):
        t0 = time.perf_counter()
        _, _ = process_vision_info(text_msgs)
        t1 = time.perf_counter()
        times_text.append(t1 - t0)

    # Profile with image
    times_img = []
    for _ in range(20):
        t0 = time.perf_counter()
        _, _ = process_vision_info(img_msgs)
        t1 = time.perf_counter()
        times_img.append(t1 - t0)

    avg_text = statistics.mean(times_text) * 1000
    avg_img = statistics.mean(times_img) * 1000

    print(f"\n  process_vision_info overhead:")
    print(f"    Text-only: {avg_text:.2f}ms avg")
    print(f"    With image: {avg_img:.2f}ms avg")
    print(f"    Image overhead: +{avg_img - avg_text:.2f}ms ({avg_img/avg_text:.1f}x)")

    os.unlink(img_path)
    os.rmdir(tmp_dir)
    return {"text_ms": avg_text, "image_ms": avg_img}


def main():
    print("=" * 60)
    print("DATA LOADING PIPELINE PROFILING")
    print("=" * 60)

    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", trust_remote_code=True)
    tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", trust_remote_code=True)

    print("\n--- 1. process_vision_info ---")
    pvi_results = profile_process_vision_info()

    print("\n--- 2. Processor call ---")
    proc_results = profile_processor_overhead(processor)

    print("\n--- 3. DataCollatorForQwenVL (text-only) ---")
    text_results = profile_collator_text_only(processor, tokenizer, n_samples=20, batch_size=1)

    print("\n--- 4. DataCollatorForQwenVL (with image) ---")
    img_results = profile_collator_with_image(processor, tokenizer, n_samples=10, batch_size=1)

    # Summary
    print(f"\n{'=' * 60}")
    print("SUMMARY: Data Loading Bottleneck Analysis")
    print(f"{'=' * 60}")
    print(f"{'Component':<35} {'Text-only':<15} {'With Image':<15} {'Overhead':<10}")
    print(f"{'-' * 75}")
    print(f"{'process_vision_info':<35} {pvi_results['text_ms']:<15.2f} {pvi_results['image_ms']:<15.2f} {pvi_results['image_ms']/pvi_results['text_ms']:<10.1f}x")
    print(f"{'processor() call':<35} {proc_results['text_ms']:<15.1f} {proc_results['image_ms']:<15.1f} {proc_results['image_ms']/proc_results['text_ms']:<10.1f}x")
    print(f"{'Full collator (1 sample)':<35} {text_results['avg_ms']:<15.1f} {img_results['avg_ms']:<15.1f} {img_results['avg_ms']/text_results['avg_ms']:<10.1f}x")
    print(f"{'=' * 60}")

    # Key insight
    print(f"\nKey insight:")
    print(f"  Image data loading adds ~{img_results['avg_ms'] - text_results['avg_ms']:.0f}ms per sample")
    print(f"  At batch_size=1 with 8 workers, this translates to")
    print(f"  ~{(img_results['avg_ms'] - text_results['avg_ms'])/8:.0f}ms amortized overhead per GPU step")
    print(f"  vs model fwd+bwd time of ~2800ms (from GPU profiling)")
    print(f"  → Data loading is {'BOTTLENECK' if img_results['avg_ms'] > 500 else 'manageable'} with 8 workers")


if __name__ == "__main__":
    main()
