# shiyi_scripts/count_vila_mixture.py
import os
import sys
from pathlib import Path

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

os.environ["VILA_DATASETS"] = "cs-oci-ord"

from llava.data.builder import parse_mixture, DATASETS, MIXTURES
from llava.data import datasets_mixture

datasets_mixture.register_datasets_mixtures()

mixture_name = "j63_recipe"
print(f"Mixture: {mixture_name}")
print("=" * 80)

dataset_names = parse_mixture(mixture_name)
print(f"Total sub-datasets: {len(dataset_names)}\n")

for i, name in enumerate(dataset_names, 1):
    if name in DATASETS:
        ds_config = DATASETS[name]
        data_path = ds_config.get("data_path", "N/A")
        media_dir = ds_config.get("media_dir", "N/A")
        
        # Try to estimate size
        if os.path.exists(data_path):
            if data_path.endswith(".jsonl"):
                with open(data_path) as f:
                    count = sum(1 for _ in f)
            elif data_path.endswith(".json"):
                import json
                with open(data_path) as f:
                    data = json.load(f)
                    count = len(data) if isinstance(data, list) else "?"
            else:
                count = "?"
        else:
            count = "file not found"
        
        print(f"{i}. {name}")
        print(f"   Samples: {count}")
        print(f"   Data: {data_path}")
        print(f"   Media: {media_dir}")
        print()
    else:
        print(f"{i}. {name} - NOT FOUND IN REGISTRY")
        print()

print("=" * 80)
print(f"Total mixture samples: 1,267,819 (from test)")
