"""
COCO dataset which returns image_id for evaluation.

Mostly copy-paste from https://github.com/pytorch/vision/blob/13b35ff/references/detection/coco_utils.py
"""

import torch
import json
from PIL import Image, ImageDraw

from .modulated_coco import ConvertCocoPolysToMask
from .tsv import ODTSVDataset
from pycocotools.coco import COCO
from maskrcnn_benchmark.structures.bounding_box import BoxList
import random
from .od_to_grounding import convert_object_detection_to_grounding_optimized_for_od, check_for_positive_overflow, sanity_check_target_after_processing


class CocoDetectionTSV(ODTSVDataset):
    def __init__(self,
                 name,
                 yaml_file,
                 transforms,
                 return_tokens,
                 tokenizer,
                 extra_fields,
                 random_sample_negative=-1,
                 add_detection_prompt=False,
                 add_detection_prompt_advanced=False,
                 use_od_data_aug=False,
                 control_probabilities={},
                 disable_shuffle=False,
                 prompt_engineer_version="v2",
                 prompt_limit_negative=-1,
                 positive_question_probability=0.6,
                 negative_question_probability=0.8,
                 full_question_probability=0.5,
                 disable_clip_to_image=False,
                 separation_tokens=" ",
                 no_mask_for_od=False,
                 max_num_labels=-1,
                 max_query_len=256,
                 **kwargs
                 ):
        super(CocoDetectionTSV, self).__init__(yaml_file, extra_fields, **kwargs)

        self._transforms = transforms
        self.name = name
        self.max_query_len = max_query_len
        self.prepare = ConvertCocoPolysToMask(
            return_masks=False,
            return_tokens=return_tokens,
            tokenizer=tokenizer,
            max_query_len=max_query_len
        )
        self.tokenizer = tokenizer

        self.control_probabilities = control_probabilities
        self.random_sample_negative = random_sample_negative
        self.add_detection_prompt = add_detection_prompt
        self.add_detection_prompt_advanced = add_detection_prompt_advanced
        self.use_od_data_aug = use_od_data_aug

        self.prompt_engineer_version = prompt_engineer_version
        self.prompt_limit_negative = prompt_limit_negative
        self.positive_question_probability = positive_question_probability
        self.negative_question_probability = negative_question_probability
        self.full_question_probability = full_question_probability
        self.separation_tokens = separation_tokens
        self.disable_clip_to_image = disable_clip_to_image
        self.disable_shuffle = disable_shuffle
        self.no_mask_for_od = no_mask_for_od
        self.max_num_labels = max_num_labels

    def __len__(self):
        return super(CocoDetectionTSV, self).__len__()

    def categories(self, no_background=True):
        categories = self.coco.dataset["categories"]
        label_list = {}
        for index, i in enumerate(categories):
            # assert(index + 1 == i["id"])
            if not no_background or (i["name"] != "__background__" and i['id'] != 0):
                label_list[i["id"]] = i["name"]
        return label_list

    def __getitem__(self, idx):
        # tgt is a BoxList
        img, target, _, scale = super(CocoDetectionTSV, self).__getitem__(idx)
        image_id = self.get_img_id(idx)
        restricted_negative_list = None

        if not self.disable_clip_to_image:
            target = target.clip_to_image(remove_empty=True)

        original_box_num = len(target)

        target, positive_caption_length = check_for_positive_overflow(target, self.ind_to_class, self.tokenizer, self.max_query_len-2) # leave some space for the special tokens

        if len(target) < original_box_num:
            print("WARNING: removed {} boxes due to positive caption overflow".format(original_box_num - len(target)))

        annotations, caption, greenlight_span_for_masked_lm_objective, label_to_positions = convert_object_detection_to_grounding_optimized_for_od(
            target=target,
            image_id=image_id,
            ind_to_class=self.ind_to_class,
            disable_shuffle=self.disable_shuffle,
            add_detection_prompt=self.add_detection_prompt,
            add_detection_prompt_advanced=self.add_detection_prompt_advanced,
            random_sample_negative=self.random_sample_negative,
            control_probabilities=self.control_probabilities,
            restricted_negative_list=restricted_negative_list,
            separation_tokens=self.separation_tokens,
            max_num_labels=self.max_num_labels,
            positive_caption_length=positive_caption_length,
            tokenizer=self.tokenizer,
            max_seq_length=self.max_query_len-2
        )

        # assert(len(self.tokenizer.tokenize(caption)) <= self.max_query_len-2)

        # print(caption)
        anno = {"image_id": image_id, "annotations": annotations, "caption": caption, "label_to_positions": label_to_positions}
        anno["greenlight_span_for_masked_lm_objective"] = greenlight_span_for_masked_lm_objective

        if self.no_mask_for_od:
            anno["greenlight_span_for_masked_lm_objective"].append((-1, -1, -1))

        img, anno = self.prepare(img, anno, box_format="xyxy")

        if self._transforms is not None:
            img, target = self._transforms(img, target)
        
        # add additional property
        for ann in anno:
            target.add_field(ann, anno[ann])

        sanity_check_target_after_processing(target)

        return img, target, idx

    def get_raw_image(self, idx):
        image, *_ = super(CocoDetectionTSV, self).__getitem__(idx)
        return image

    def get_img_id(self, idx):
        line_no = self.get_line_no(idx)
        if self.label_tsv is not None:
            row = self.label_tsv.seek(line_no)
            img_id = row[0]
            try:
                return int(img_id)
            except:
                return idx