import os
import json
import logging
from collections import OrderedDict, defaultdict

logger = logging.getLogger(__name__)

from detectron2.data import DatasetCatalog, MetadataCatalog
from detectron2.structures import Boxes, BoxMode, PolygonMasks, RotatedBoxes
from .VisualGenome import _get_vg_meta

__all__ = ["load_syn_json", "register_syn"]

def _get_metadata(categories):
    if len(categories) == 0:
        return {}
    id_to_name = {x["id"]: x["name"] for x in categories}
    thing_dataset_id_to_contiguous_id = {i + 1: i for i in range(len(categories))}
    thing_classes = [id_to_name[k] for k in sorted(id_to_name)]
    return {
        "thing_dataset_id_to_contiguous_id": thing_dataset_id_to_contiguous_id,
        "thing_classes": thing_classes,
    }

def load_syn_json(json_file, image_root):
    dataset_dicts = []

    with open(json_file, 'r') as file:
        data_json = json.load(file)

    # Map image IDs to annotations
    imgid_to_descrs = {annotation["image_id"]: annotation for annotation in data_json["annotations"]}

    data_id = 0

    for img_dict in data_json["images"]:
        img_id = img_dict["id"]
        
        annotation = imgid_to_descrs[img_id]
        annotation["bbox_mode"] = BoxMode.XYWH_ABS
        annotation["category_id"] = 0
        annotation["object_description"] = img_dict["caption"]
        #combined_captions = [img_dict["caption"]] + img_dict["neg_caption"][:3]
        #annotations = [annotation] + [{} for _ in range(3)],

        # Create a record for the positive caption
        #annotation["phrase"] = img_dict["caption"]
        #for ann in annotation:
        #annotation["positive_map"] = 1
        pos_record = {
            "file_name": os.path.join(image_root, img_dict["file_name"]),
            "height": img_dict["height"],
            "width": img_dict["width"],
            #"expressions": img_dict["caption"],
            "references": img_dict["neg_caption"],
            #"text_prompt": img_dict["caption"],
            #"phrase": img_dict["caption"],
            #"id": data_id,
            "annotations": [annotation,],
            "has_mask": False,
            "length" : 1,
            #"annotations": [annotation] + [{} for _ in range(3)],
            #"instances": [annotation] + [{} for _ in range(3)],
            "task": "vg_syn2",
        }
        dataset_dicts.append(pos_record)
        data_id += 1  
                
    return dataset_dicts


_PREDEFINED_SYN_SPLITS = {
    "syn_train": (
        "chatgpt_3.5_10_re/des_dir/image", 
        "chatgpt_3.5_10_re/des_dir/gen_gt_pixart_fiber_multi_grain_parse_noun_valid_sen.json"
    ),
}

def _get_builtin_syn_metadata(name, image_root, json_file, metadata):
    """
    Args:
        name (str): the name that identifies a dataset, e.g. "coco_2014_train".
        metadata (dict): extra metadata associated with this dataset.  You can
            leave it as an empty dict.
        json_file (str): path to the json instance annotation file.
        image_root (str or path-like): directory which contains all the images.
    """
    assert isinstance(name, str), name
    assert isinstance(json_file, (str, os.PathLike)), json_file
    assert isinstance(image_root, (str, os.PathLike)), image_root
    # 1. register a function which returns dicts 
    DatasetCatalog.register(name, lambda: load_syn_json(json_file, image_root))
    # 2. Optionally, add metadata about this dataset,
    # since they might be useful in evaluation, visualization or logging
    MetadataCatalog.get(name).set(
        json_file=json_file, image_root=image_root, evaluator_type="coco", **metadata,
    )
     
def register_syn(root):
    for key, (image_root, json_file) in _PREDEFINED_SYN_SPLITS.items():
        # Assume pre-defined datasets live in `./datasets`.
        image_root = os.path.join(root, image_root)
        json_file = os.path.join(root, json_file) if "://" not in json_file else json_file
        _get_builtin_syn_metadata(key, image_root, json_file, _get_vg_meta())