#!/usr/bin/env python3
"""
Test sharegpt4v_sft dataset initialization to find bottlenecks.

Usage:
    cd /lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow
    python3 shiyi_scripts/test_sharegpt4v_sft_init.py
"""
import os
import sys
import time
import json
from pathlib import Path

PROJECT_ROOT = Path.cwd()
sys.path.insert(0, str(PROJECT_ROOT / "VILA"))

print("=" * 80)
print("Testing sharegpt4v_sft Dataset Initialization")
print("=" * 80)

# Simulate VILA environment
os.environ["VILA_DATASETS"] = "cs-oci-ord"

data_path = "/lustre/fsw/portfolios/nvr/projects/nvr_elm_llm/dataset/vila-sft/ShareGPT4V/jason-filter-sharegpt4v_mix665k_cap23k_coco-ap9k_lcs3k_sam9k_div2k.json"

print(f"\nData path: {data_path}")
print(f"File exists: {os.path.exists(data_path)}")
if os.path.exists(data_path):
    file_size = os.path.getsize(data_path) / (1024**3)
    print(f"File size: {file_size:.2f} GB")

# Step 1: Load JSON
print("\n" + "-" * 80)
print("Step 1: Loading JSON file...")
start = time.time()
with open(data_path) as f:
    instances = json.load(f)
load_time = time.time() - start
print(f"✅ Loaded {len(instances):,} samples in {load_time:.2f}s")
print(f"   Throughput: {len(instances) / load_time:.0f} samples/sec")

# Step 2: Simulate padding logic
print("\n" + "-" * 80)
print("Step 2: Simulating padding logic...")

# Calculate global_batch_size (from training config)
per_device_batch = 1
num_gpus = 8
grad_acc = 8
global_batch_size = per_device_batch * num_gpus * grad_acc

print(f"global_batch_size: {global_batch_size}")
print(f"len(instances): {len(instances):,}")

residual = global_batch_size - len(instances) % global_batch_size
print(f"residual: {residual}")

if residual != global_batch_size:
    print(f"\n  Padding is needed...")
    
    # Check if need to repeat entire dataset
    repeat_times = global_batch_size // len(instances)
    print(f"  repeat_times: {repeat_times}")
    
    if repeat_times >= 2:
        print(f"  Will repeat entire dataset {repeat_times} times")
        start = time.time()
        instances_repeated = instances * repeat_times
        repeat_time = time.time() - start
        print(f"  ✅ Repeated in {repeat_time:.2f}s, new length: {len(instances_repeated):,}")
        instances = instances_repeated
        residual = global_batch_size - len(instances) % global_batch_size
        print(f"  New residual: {residual}")
    
    # Sample additional instances
    print(f"\n  Sampling {residual} additional samples...")
    start = time.time()
    import random
    selected_elements = random.sample(range(len(instances)), residual)
    sample_time = time.time() - start
    print(f"  ✅ Sampled in {sample_time:.2f}s")
    
    # Extend instances
    print(f"  Extending instances list...")
    start = time.time()
    additional_instance = [instances[i] for i in selected_elements]
    extend_time = time.time() - start
    print(f"  ✅ Created additional list in {extend_time:.2f}s")
    
    start = time.time()
    instances.extend(additional_instance)
    extend_time2 = time.time() - start
    print(f"  ✅ Extended in {extend_time2:.2f}s")
    
    print(f"  Final length: {len(instances):,}")
else:
    print(f"\n  No padding needed (already multiple of {global_batch_size})")

# Step 3: Check first few samples
print("\n" + "-" * 80)
print("Step 3: Checking sample structure...")

for i in range(min(3, len(instances))):
    sample = instances[i]
    print(f"\nSample {i}:")
    print(f"  Keys: {list(sample.keys())}")
    if "conversations" in sample:
        print(f"  conversations: {len(sample['conversations'])} turns")
    if "image" in sample:
        img_val = sample["image"]
        if isinstance(img_val, str):
            print(f"  image: {img_val}")
        else:
            print(f"  image: {type(img_val)}")

# Step 4: Estimate memory usage
print("\n" + "-" * 80)
print("Step 4: Estimating memory usage...")

import sys
mem_bytes = sys.getsizeof(instances)
# Rough estimate (doesn't include nested objects)
mem_mb = mem_bytes / (1024**2)
print(f"Instances list size: ~{mem_mb:.1f} MB (rough estimate)")
print(f"Per-rank memory (8 ranks): ~{mem_mb * 8:.1f} MB")

print("\n" + "=" * 80)
print("Summary:")
print(f"  JSON load time: {load_time:.2f}s")
print(f"  Total samples: {len(instances):,}")
print(f"  Estimated init time per rank: {load_time + (sample_time if residual != global_batch_size else 0):.2f}s")
print(f"  Estimated total init time (8 ranks parallel): ~{load_time + 5:.0f}s")
print("=" * 80)

if load_time > 30:
    print("\n⚠️  WARNING: JSON loading is slow (>30s)")
    print("   Consider using WebDataset format or local SSD")
else:
    print("\n✅ JSON loading is acceptable")
    print("   If training hangs, the issue is likely elsewhere (collator, model, etc.)")

