#!/usr/bin/env python3
"""
Debug j63_recipe mixture to find which sub-dataset causes issues.

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/debug_j63_datasets.py
"""
import os
import sys
from pathlib import Path

PROJECT_ROOT = Path.cwd()
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 vila_path_patcher

os.environ["VILA_DATASETS"] = "cs-oci-ord"
os.environ["RANK"] = "0"
os.environ["WORLD_SIZE"] = "1"
os.environ["LOCAL_RANK"] = "0"

from lmflow.args import DatasetArguments
from lmflow.datasets.dataset import Dataset
from transformers import AutoTokenizer

print("=" * 80)
print("Debugging j63_recipe Mixture")
print("=" * 80)

print("\nLoading tokenizer...")
tokenizer = AutoTokenizer.from_pretrained(
    "Qwen/Qwen2.5-VL-3B-Instruct",
    trust_remote_code=True,
    use_fast=False
)
print(f"✅ Tokenizer loaded")

data_args = DatasetArguments(dataset_path="j63_recipe")
data_args.tokenizer = tokenizer

print("\nCreating dataset...")
dataset = Dataset(data_args, backend="webdataset_multi_modal")
print(f"✅ Dataset created")

from torch.utils.data import ConcatDataset

# Get the wrapped VILA dataset
adapter = dataset.backend_dataset
vila_dataset = adapter.vila_dataset

print(f"\nAdapter type: {type(adapter).__name__}")
print(f"VILA dataset type: {type(vila_dataset).__name__}")
print(f"Is ConcatDataset: {isinstance(vila_dataset, ConcatDataset)}")

if isinstance(vila_dataset, ConcatDataset):
    print(f"\n{'=' * 80}")
    print(f"ConcatDataset Analysis")
    print(f"{'=' * 80}")
    print(f"Number of sub-datasets: {len(vila_dataset.datasets)}")
    print(f"Cumulative sizes: {vila_dataset.cumulative_sizes}")
    print(f"Total samples: {len(vila_dataset)}")
    
    for i, sub_ds in enumerate(vila_dataset.datasets):
        print(f"\n{'─' * 80}")
        print(f"Sub-dataset {i}")
        print(f"{'─' * 80}")
        print(f"  Type: {type(sub_ds).__name__}")
        print(f"  Module: {type(sub_ds).__module__}")
        print(f"  Length: {len(sub_ds)}")
        
        # Check attributes
        attrs = {
            "instances": hasattr(sub_ds, "instances"),
            "list_data_dict": hasattr(sub_ds, "list_data_dict"),
            "media_dir": hasattr(sub_ds, "media_dir"),
            "image_folder": hasattr(sub_ds, "image_folder"),
            "data_path": hasattr(sub_ds, "data_path"),
        }
        print(f"  Attributes: {attrs}")
        
        if hasattr(sub_ds, 'media_dir'):
            media_dir = sub_ds.media_dir
            exists = os.path.exists(media_dir) if media_dir else False
            print(f"  media_dir: {media_dir} (exists={exists})")
        
        if hasattr(sub_ds, 'image_folder'):
            image_folder = sub_ds.image_folder
            exists = os.path.exists(image_folder) if image_folder else False
            print(f"  image_folder: {image_folder} (exists={exists})")
        
        if hasattr(sub_ds, 'data_path'):
            data_path = sub_ds.data_path
            exists = os.path.exists(data_path) if data_path else False
            print(f"  data_path: {data_path} (exists={exists})")
        
        # Try to access raw data
        print(f"  Testing raw data access...")
        try:
            if hasattr(sub_ds, 'instances'):
                raw_sample = sub_ds.instances[0]
                source = "instances"
            elif hasattr(sub_ds, 'list_data_dict'):
                raw_sample = sub_ds.list_data_dict[0]
                source = "list_data_dict"
            else:
                print(f"    ⚠️  No known raw data attribute found")
                raw_sample = None
                source = None
            
            if raw_sample:
                print(f"    ✅ Accessed via '{source}'")
                print(f"    Sample keys: {list(raw_sample.keys())}")
                print(f"    Has 'conversations': {'conversations' in raw_sample}")
                print(f"    Has 'image': {'image' in raw_sample}")
                print(f"    Has 'images': {'images' in raw_sample}")
                
                if 'image' in raw_sample:
                    img_val = raw_sample['image']
                    print(f"    image type: {type(img_val)}")
                    if isinstance(img_val, str):
                        print(f"    image value: {img_val[:100]}")
        except Exception as e:
            print(f"    ❌ ERROR: {e}")

print(f"\n{'=' * 80}")
print("Testing Adapter on Samples from Each Sub-dataset")
print(f"{'=' * 80}")

# Test first sample from each sub-dataset
for i in range(len(vila_dataset.datasets)):
    # Calculate global index for first sample of this sub-dataset
    if i == 0:
        global_idx = 0
    else:
        global_idx = vila_dataset.cumulative_sizes[i-1]
    
    print(f"\n--- Testing sub-dataset {i}, global_idx={global_idx} ---")
    try:
        adapted_sample = adapter[global_idx]
        
        has_messages = "messages" in adapted_sample
        num_messages = len(adapted_sample.get("messages", [])) if has_messages else 0
        
        print(f"  ✅ Adapter success")
        print(f"     has_messages: {has_messages}")
        print(f"     num_messages: {num_messages}")
        
        if has_messages and num_messages > 0:
            first_msg = adapted_sample["messages"][0]
            print(f"     first_msg role: {first_msg.get('role')}")
            content = first_msg.get("content", [])
            print(f"     first_msg content items: {len(content)}")
            for item in content:
                item_type = item.get("type", "unknown")
                if item_type == "image":
                    img = item.get("image")
                    if img:
                        print(f"       - image: {type(img).__name__} size={img.size if hasattr(img, 'size') else 'N/A'}")
                    else:
                        print(f"       - image: None")
                elif item_type == "text":
                    text = item.get("text", "")
                    preview = text[:50] + "..." if len(text) > 50 else text
                    print(f"       - text: {preview}")
        
    except Exception as e:
        print(f"  ❌ Adapter FAILED: {e}")
        import traceback
        traceback.print_exc()

print(f"\n{'=' * 80}")
print("Testing Random Samples Across All Sub-datasets")
print(f"{'=' * 80}")

# Test a few random samples
import random
random.seed(42)
test_indices = random.sample(range(len(adapter)), min(20, len(adapter)))

success = 0
failed = 0

for idx in test_indices:
    try:
        sample = adapter[idx]
        if "messages" in sample and len(sample["messages"]) > 0:
            success += 1
            if success <= 3:  # Show first 3 successes
                print(f"✅ Sample {idx}: {len(sample['messages'])} messages")
        else:
            failed += 1
            print(f"⚠️  Sample {idx}: missing or empty messages")
    except Exception as e:
        failed += 1
        print(f"❌ Sample {idx}: {e}")

print(f"\n{'=' * 80}")
print(f"Summary:")
print(f"  Total samples tested: {len(test_indices)}")
print(f"  Success: {success}")
print(f"  Failed: {failed}")
print(f"  Success rate: {success / len(test_indices) * 100:.1f}%")
print(f"{'=' * 80}")

if failed > 0:
    print("\n⚠️  WARNING: Some samples failed to load!")
    print("   This may cause training to hang when these samples are encountered.")
    print("   Check the error messages above to identify problematic sub-datasets.")
else:
    print("\n✅ All samples loaded successfully!")

