import os

from detectron2.data import DatasetCatalog, MetadataCatalog
from .grefcoco import load_grefcoco_json

def register_grefcoco(root):
    image_root = os.path.join(root, "images", "train2014")
    dataset_info = [
        ('grefcoco', 'unc', ['train', 'val', 'testA', 'testB']),
    ]
    for name, splitby, splits in dataset_info:
        for split in splits:
            dataset_id = '_'.join([name, splitby, split])
            DatasetCatalog.register(
                dataset_id,
                lambda root=root, name=name, splitby=splitby, split=split, image_root=image_root: 
                    load_grefcoco_json(root, name, splitby, split, image_root)
            )
            MetadataCatalog.get(dataset_id).set(
                evaluator_type="refer",
                dataset_name=name,
                splitby=splitby,
                split=split,
                root=root,
                image_root=image_root,
            )

_root = os.getenv("DETECTRON2_DATASETS", "datasets")
register_grefcoco(_root)