#!/usr/bin/env python3
"""
Test VILA mixture data loading speed and format.

Usage:
    cd /lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow
    OPENBLAS_NUM_THREADS=1 OMP_NUM_THREADS=1 python3 shiyi_scripts/test_vila_dataloader.py
"""
import os
import sys
import time
from pathlib import Path

# Setup paths
PROJECT_ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(PROJECT_ROOT / "src"))
sys.path.insert(0, str(PROJECT_ROOT / "VILA"))
sys.path.insert(0, str(PROJECT_ROOT / "examples"))

# Import patcher first
import vila_path_patcher

from lmflow.args import DatasetArguments
from lmflow.datasets.dataset import Dataset
from torch.utils.data import DataLoader

# Simulate training environment
os.environ["DATASET_BACKEND"] = "webdataset_multi_modal"
os.environ["VILA_DATASETS"] = "cs-oci-ord"
os.environ["RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
os.environ["LOCAL_RANK"] = "0"

print("=" * 80)
print("VILA Mixture Data Loading Test")
print("=" * 80)

# Create args
mixture_name = os.environ.get("TEST_MIXTURE", "j63_recipe")
print(f"\nTesting mixture: {mixture_name}")

data_args = DatasetArguments(dataset_path=mixture_name)

# Load tokenizer
print("\nLoading tokenizer...")
from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained(
    "Qwen/Qwen2.5-VL-3B-Instruct",
    trust_remote_code=True,
    use_fast=False
)
data_args.tokenizer = tokenizer
print(f"✅ Tokenizer loaded: {type(tokenizer).__name__}")

# Create dataset
print("\n" + "=" * 80)
print("Creating dataset...")
start = time.time()
dataset = Dataset(data_args, backend="webdataset_multi_modal")
elapsed = time.time() - start
print(f"✅ Dataset created in {elapsed:.2f}s")
print(f"   Dataset type: {type(dataset.backend_dataset).__name__}")
print(f"   Dataset length: {len(dataset)}")

# Check if it's a ConcatDataset
from torch.utils.data import ConcatDataset
if isinstance(dataset.backend_dataset, ConcatDataset):
    print(f"   Sub-datasets: {len(dataset.backend_dataset.datasets)}")
    for i, sub_ds in enumerate(dataset.backend_dataset.datasets):
        print(f"     [{i}] {type(sub_ds).__name__}: {len(sub_ds)} samples")

# Test loading individual samples
print("\n" + "=" * 80)
print("Testing individual sample loading...")
num_samples = min(10, len(dataset))
total_time = 0

for i in range(num_samples):
    start = time.time()
    try:
        sample = dataset.backend_dataset[i]
        elapsed = time.time() - start
        total_time += elapsed
        
        # Check sample structure
        has_messages = "messages" in sample
        num_messages = len(sample.get("messages", [])) if has_messages else 0
        
        status = "✅" if has_messages else "⚠️"
        print(f"{status} Sample {i}: {elapsed:.3f}s, has_messages={has_messages}, num_messages={num_messages}")
        
        # Check first message details
        if has_messages and num_messages > 0:
            first_msg = sample["messages"][0]
            content = first_msg.get("content", [])
            has_image = any(item.get("type") == "image" for item in content)
            has_text = any(item.get("type") == "text" for item in content)
            
            if has_image:
                img_item = next((item for item in content if item.get("type") == "image"), None)
                if img_item and "image" in img_item:
                    img = img_item["image"]
                    print(f"     Image: {type(img).__name__} size={img.size if hasattr(img, 'size') else 'N/A'}")
            
            if has_text:
                text_item = next((item for item in content if item.get("type") == "text"), None)
                if text_item:
                    text = text_item.get("text", "")
                    preview = text[:50] + "..." if len(text) > 50 else text
                    print(f"     Text: {preview}")
            
    except Exception as e:
        elapsed = time.time() - start
        print(f"❌ Sample {i}: FAILED after {elapsed:.3f}s - {e}")
        import traceback
        traceback.print_exc()

avg_time = total_time / num_samples if num_samples > 0 else 0
print(f"\nAverage sample load time: {avg_time:.3f}s")

# Test DataLoader with collator
print("\n" + "=" * 80)
print("Testing DataLoader (with collator)...")

try:
    from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL
    from transformers import AutoProcessor
    
    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", trust_remote_code=True)
    collator = DataCollatorForQwenVL(
        processor=processor,
        tokenizer=tokenizer,
        bd_size=None,  # No MDM for testing
    )
    
    loader = DataLoader(
        dataset.backend_dataset,
        batch_size=2,
        num_workers=0,
        shuffle=False,
        collate_fn=collator
    )
    
    print(f"DataLoader created: batch_size=2, num_workers=0")
    print("Loading first 5 batches...")
    
    for i, batch in enumerate(loader):
        elapsed_batch = time.time()
        print(f"\nBatch {i}:")
        print(f"  Keys: {list(batch.keys())}")
        for key, val in batch.items():
            if hasattr(val, "shape"):
                print(f"    {key}: shape={val.shape}, dtype={val.dtype}")
            else:
                print(f"    {key}: type={type(val)}")
        
        batch_time = time.time() - elapsed_batch
        print(f"  Batch processing time: {batch_time:.3f}s")
        
        if i >= 4:  # Only test first 5 batches
            break
    
    print("\n✅ DataLoader test complete!")
    
except Exception as e:
    print(f"\n❌ DataLoader test failed: {e}")
    import traceback
    traceback.print_exc()

print("\n" + "=" * 80)
print("Test Summary:")
print(f"  Dataset: {mixture_name}")
print(f"  Total samples: {len(dataset)}")
print(f"  Avg load time: {avg_time:.3f}s/sample")
print(f"  Estimated time for 1 epoch (single process): {avg_time * len(dataset) / 3600:.1f} hours")
print("=" * 80)

