#!/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_mixture.py
    
    # Test specific mixture:
    TEST_MIXTURE=j63_recipe python3 shiyi_scripts/test_vila_mixture.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
try:
    import vila_path_patcher
    print("✅ Path patcher loaded")
except ImportError:
    print("⚠️  Path patcher not found, paths may not be remapped")

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...")
start_time = time.time()
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 in {time.time() - start_time:.2f}s: {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):
        sub_type = type(sub_ds).__name__
        sub_len = len(sub_ds)
        # Try to get dataset name
        dataset_name = getattr(sub_ds, "name", "unknown")
        print(f"     [{i}] {sub_type}: {sub_len:,} samples (name={dataset_name})")

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

for i in range(num_samples):
    start = time.time()
    try:
        sample = dataset.backend_dataset[i]
        elapsed = time.time() - start
        total_time += elapsed
        success_count += 1
        
        # Check sample structure
        has_messages = "messages" in sample
        num_messages = len(sample.get("messages", [])) if has_messages else 0
        
        status = "✅" if has_messages else "⚠️"
        
        if i < 5:  # Show details for first 5 samples
            print(f"\n{status} Sample {i}: {elapsed:.3f}s")
            print(f"   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)
                
                print(f"   First message: role={first_msg.get('role')}, has_image={has_image}, has_text={has_text}")
                
                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"]
                        img_info = f"{type(img).__name__}"
                        if hasattr(img, 'size'):
                            img_info += f" size={img.size}"
                        print(f"   Image: {img_info}")
                
                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[:60] + "..." if len(text) > 60 else text
                        print(f"   Text: {preview}")
        else:
            # Just show progress for remaining samples
            if i % 5 == 0:
                print(f"{status} Sample {i}: {elapsed:.3f}s")
            
    except Exception as e:
        elapsed = time.time() - start
        fail_count += 1
        print(f"\n❌ Sample {i}: FAILED after {elapsed:.3f}s")
        print(f"   Error: {e}")
        if fail_count <= 2:  # Show traceback for first 2 failures
            import traceback
            traceback.print_exc()

avg_time = total_time / success_count if success_count > 0 else 0
print(f"\n{'=' * 80}")
print(f"Sample Loading Summary:")
print(f"  Success: {success_count}/{num_samples}")
print(f"  Failed: {fail_count}/{num_samples}")
print(f"  Average load time: {avg_time:.3f}s/sample")
if avg_time > 0:
    print(f"  Estimated throughput: {1/avg_time:.1f} samples/sec")
print(f"{'=' * 80}")

# Test DataLoader with collator
print("\nTesting DataLoader (with collator)...")

try:
    from lmflow.datasets.multi_modal_dataset import DataCollatorForQwenVL
    from transformers import AutoProcessor
    
    print("Loading processor...")
    processor = AutoProcessor.from_pretrained("Qwen/Qwen2.5-VL-3B-Instruct", trust_remote_code=True)
    
    print("Creating collator...")
    collator = DataCollatorForQwenVL(
        processor=processor,
        tokenizer=tokenizer,
        bd_size=None,  # No MDM for testing
    )
    
    print("Creating DataLoader...")
    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("\nLoading first 3 batches...")
    
    batch_times = []
    for i, batch in enumerate(loader):
        batch_start = 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}")
            elif isinstance(val, (list, tuple)):
                print(f"    {key}: type={type(val)}, len={len(val)}")
            else:
                print(f"    {key}: type={type(val)}")
        
        batch_time = time.time() - batch_start
        batch_times.append(batch_time)
        print(f"  Batch processing time: {batch_time:.3f}s")
        
        if i >= 2:  # Only test first 3 batches
            break
    
    avg_batch_time = sum(batch_times) / len(batch_times) if batch_times else 0
    print(f"\n✅ DataLoader test complete!")
    print(f"   Average batch time: {avg_batch_time:.3f}s")
    
except Exception as e:
    print(f"\n❌ DataLoader test failed: {e}")
    import traceback
    traceback.print_exc()

print("\n" + "=" * 80)
print("Test Summary:")
print(f"  Mixture: {mixture_name}")
print(f"  Total samples: {len(dataset):,}")
print(f"  Sample load time: {avg_time:.3f}s/sample")
if avg_time > 0:
    print(f"  Estimated time for 1000 samples: {avg_time * 1000 / 60:.1f} minutes")
print("=" * 80)
print("\n✅ All tests passed! Ready for training.")

