import json

from ..root import DATASETS, IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER, POINTS_PLACEHOLDER
from ..utils import MInstrDataset


@DATASETS.register_module()
class ClevrDataset(MInstrDataset):
    def __init__(self, *args, scene_graph_file, version, **kwargs):
        super().__init__(*args, **kwargs, placeholders=(IMAGE_PLACEHOLDER, QUESTION_PLACEHOLDER))
        self.scene_graph_file = scene_graph_file
        self.version = version
        qtype, atype = version.split('-')
        assert qtype in ['q']
        assert atype in ['a', 's', 'bs']
        self.qtype = qtype
        self.atype = atype

        if scene_graph_file is None:
            self.scene_graph = None
        else:
            self.scene_graph = [line for line in open(scene_graph_file, 'r', encoding='utf8')]

    def get_raw_item(self, index):
        question = json.loads(self.data[index])
        if self.scene_graph is None:
            scene = None
        else:
            scene = json.loads(self.scene_graph[question['image_index']])
        return question, scene

    def __getitem__(self, index):
        question, scene = self.get_raw_item(index)
        img_path = question['image_filename']
        image = self.get_image(img_path)

        if self.atype == 'a':
            boxes = []
            answer = f"The answer is {question['answer']}."
            answer_boxes_seq = []
        elif self.atype == 's':
            answer, boxes, answer_boxes_seq = clevr_ss_cot(obj=question, scene=scene, add_ref=False)
            answer += f" The answer is {question['answer']}."
        elif self.atype == 'bs':
            answer, boxes, answer_boxes_seq = clevr_ss_cot(obj=question, scene=scene, add_ref=True)
            answer += f" The answer is {question['answer']}."
        else:
            assert False

        if self.qtype == 'q':
            query_boxes_seq = []
            final_query = self.get_template().replace(QUESTION_PLACEHOLDER, question['question'])
        else:
            assert False

        ret = {
            'image': image,
            'target': {'points': boxes},
            'conversations': [
                {
                    'from': 'human',
                    'value': final_query,
                    'points_seq': query_boxes_seq,
                },
                {
                    'from': 'gpt',
                    'value': answer,
                    'points_seq': answer_boxes_seq,
                }
            ]
        }
        return ret


def get_boxes_idx(boxes_list, refs):
    def get_idx(boxes_list, box):
        if box in boxes_list:
            return boxes_list.index(box)
        else:
            boxes_list.append(box)
            return len(boxes_list) - 1

    idx = [get_idx(boxes_list, box) for box in refs]
    return idx


def clevr_ss_cot(obj, scene, add_ref=False):
    cot = []
    boxes = []
    seq = []

    def can_add_ref():
        if p['function'] in ['unique', 'union', 'intersect', 'relate', 'same_size', 'same_shape', 'same_material', 'same_color']:
            return True
        if p['function'] in ['scene', 'filter_color', 'filter_material', 'filter_shape', 'filter_size']:
            if idx + 1 < len(obj['program']) and obj['program'][idx + 1]['function'] in ['exist', 'count']:
                return True
        return False

    for idx, p in enumerate(obj['program']):
        func = f"{p['function']}:{p['value_inputs'][0]}" if 'value_inputs' in p and p['value_inputs'] else p['function']
        inputs = f"[{','.join(map(str, p['inputs']))}]" if p['inputs'] else ""

        if add_ref and can_add_ref():
            if p['ans']:
                objs = POINTS_PLACEHOLDER
                idx = get_boxes_idx(boxes_list=boxes, refs=[scene['objects'][_]['pixel_coords'][:2] for _ in p['ans']])
                seq.append(idx)
            else:
                objs = f" Found no object."
        else:
            objs = ""
        cot.append(f"{func}{inputs}{objs}")

    ret = " -> ".join(cot)
    return ret, boxes, seq