#!/bin/bash
#SBATCH -A nvr_elm_llm
#SBATCH -p batch_block1,batch_block3,batch_block4
#SBATCH -t 01:00:00
#SBATCH -N 1
#SBATCH --ntasks-per-node=1
#SBATCH -J debug_dataloader_mminstruct_metamathqa
#SBATCH --gpus-per-node 1
#SBATCH --cpus-per-task 8
#SBATCH --mem=128G
#SBATCH -e workdir/debug_dataloader_mminstruct_metamathqa_error.log
#SBATCH -o workdir/debug_dataloader_mminstruct_metamathqa_out.log
#SBATCH --open-mode=append

# Resolve project root and change to it immediately
PROJECT_ROOT="/lustre/fs11/portfolios/nvr/projects/nvr_elm_llm/users/chengyuew/workspace/shiyi/LMFlow"
cd "${PROJECT_ROOT}"

# Activate environment
source "${PROJECT_ROOT}/scripts/cluster_env.sh"
lmflow_activate_conda

export PYTHONPATH="${PROJECT_ROOT}/src:${PROJECT_ROOT}/VILA:${PROJECT_ROOT}/examples:${PYTHONPATH:-}"
export DATASET_BACKEND="${DATASET_BACKEND:-webdataset_multi_modal}"
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS="${OMP_NUM_THREADS:-8}"
export DATASET_PATH="${DATASET_PATH:-mminstruct+metamathqa}"
export TOKENIZER_NAME="${TOKENIZER_NAME:-Qwen/Qwen2.5-VL-3B-Instruct}"

echo "=========================================="
echo "Debug VILA dataloader for mixture: ${DATASET_PATH}"
echo "VILA_DATASETS: ${VILA_DATASETS}"
echo "PROJECT_ROOT: ${PROJECT_ROOT}"
echo "=========================================="

srun --nodes=1 --ntasks=1 bash -c '
    python3 - <<'"'"'PY'"'"'
import os
import sys
from pathlib import Path

project_root = Path(os.environ.get("PROJECT_ROOT", "")) if os.environ.get("PROJECT_ROOT") else 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"))

try:
    import vila_path_patcher  # noqa: F401
except Exception as e:
    print(f"Warning: vila_path_patcher not loaded: {e}", flush=True)

os.environ.setdefault("VILA_DATASETS", "cs-oci-ord")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
os.environ.setdefault("LOCAL_RANK", "0")

mixture_name = os.environ.get("DATASET_PATH", "mminstruct+metamathqa")
tokenizer_name = os.environ.get("TOKENIZER_NAME", "Qwen/Qwen2.5-VL-3B-Instruct")

print("=" * 80)
print(f"Mixture: {mixture_name}")
print(f"VILA_DATASETS: {os.environ.get('VILA_DATASETS')}")
print("=" * 80, flush=True)

# Show registry entries for the two datasets
from llava.data.builder import register_datasets, parse_mixture

registry = register_datasets()
for name in ["metamathqa", "mminstruct"]:
    meta = registry.get(name)
    print(f"[Registry] {name}: {meta}", flush=True)

print("\nExpanded mixture:", parse_mixture(mixture_name), flush=True)

from transformers import AutoTokenizer, AutoProcessor
from lmflow.args import DatasetArguments
from lmflow.datasets.dataset import Dataset
from torch.utils.data import ConcatDataset

print("\nLoading tokenizer/processor...", flush=True)
tokenizer = AutoTokenizer.from_pretrained(tokenizer_name, trust_remote_code=True, use_fast=False)
processor = AutoProcessor.from_pretrained(tokenizer_name, trust_remote_code=True)
print("✅ Tokenizer/processor loaded", flush=True)

data_args = DatasetArguments(dataset_path=mixture_name)
data_args.tokenizer = tokenizer
data_args.image_processor = getattr(processor, "image_processor", None)

print("\nCreating LMFlow Dataset (webdataset_multi_modal)...", flush=True)
dataset = Dataset(data_args, backend="webdataset_multi_modal")
adapter = dataset.backend_dataset
vila_dataset = getattr(adapter, "vila_dataset", None)

print(f"Adapter: {type(adapter).__name__}", flush=True)
print(f"VILA dataset: {type(vila_dataset).__name__}", flush=True)

def _print_subdataset_info(sub_ds, idx):
    print(f"\nSub-dataset {idx}:", flush=True)
    print(f"  type: {type(sub_ds).__name__} ({type(sub_ds).__module__})", flush=True)
    if hasattr(sub_ds, 'data_path'):
        dp = sub_ds.data_path
        print(f"  data_path: {dp} (exists={os.path.exists(dp)})", flush=True)
    if hasattr(sub_ds, 'media_dir'):
        md = sub_ds.media_dir
        print(f"  media_dir: {md} (exists={os.path.exists(md) if md else False})", flush=True)
    print(f"  len: {len(sub_ds)}", flush=True)

if isinstance(vila_dataset, ConcatDataset):
    print(f"\nConcatDataset with {len(vila_dataset.datasets)} sub-datasets", flush=True)
    for i, sub_ds in enumerate(vila_dataset.datasets):
        _print_subdataset_info(sub_ds, i)
else:
    _print_subdataset_info(vila_dataset, 0)

print("\nTesting adapter access for first sample of each sub-dataset...", flush=True)
if isinstance(vila_dataset, ConcatDataset):
    for i in range(len(vila_dataset.datasets)):
        global_idx = 0 if i == 0 else vila_dataset.cumulative_sizes[i - 1]
        try:
            sample = adapter[global_idx]
            print(f"  ✅ sub-dataset {i} sample ok, keys={list(sample.keys())}", flush=True)
        except Exception as e:
            print(f"  ❌ sub-dataset {i} sample failed: {e}", flush=True)
else:
    try:
        sample = adapter[0]
        print(f"  ✅ sample ok, keys={list(sample.keys())}", flush=True)
    except Exception as e:
        print(f"  ❌ sample failed: {e}", flush=True)

print("\nDone.", flush=True)
PY
'
