import torch
import torch.nn as nn
from transformers import AutoImageProcessor, DetrConfig, DetrForObjectDetection


class DETRVisionTower_v2(nn.Module):
    def __init__(self, vision_tower, args, delay_load=False):
        super().__init__()

        self.is_loaded = False
        self.vision_tower_name = vision_tower # "facebook/detr-resnet-50"
        self.select_layer = args.mm_vision_select_layer
        self.select_feature = getattr(args, 'mm_vision_select_feature', 'patch')

        if not delay_load:
            self.load_model()
        else:
            self.cfg_only = DetrConfig()

    def load_model(self):
        # print(f"Loading DETR model from {self.vision_tower_name}")
        
        self.image_processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
        
        self.vision_tower = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
        self.vision_tower.config.output_hidden_states=True
        self.vision_tower.requires_grad_(False)

        self.is_loaded = True

    def feature_select(self, image_forward_outs):
        image_features = image_forward_outs.decoder_hidden_states[self.select_layer]
        if self.select_feature == 'patch':
            image_features = image_features[:, 1:]
        elif self.select_feature == 'cls_patch':
            image_features = image_features
        else:
            raise ValueError(f'Unexpected select feature: {self.select_feature}')
        return image_features
    
    def convert_to_coco_api_format(self, image, outputs):
        # convert outputs (bounding boxes and class logits) to COCO API
        target_sizes = torch.tensor([image.size[::-1]])
        results = self.image_processor.post_process_object_detection(outputs, threshold=0.9, target_sizes=target_sizes)[
            0
        ]

        for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
            box = [round(i, 2) for i in box.tolist()]
            print(
                f"Detected {self.vision_tower.config.id2label[label.item()]} with confidence "
                f"{round(score.item(), 3)} at location {box}"
            )

    @torch.no_grad()
    def forward(self, images):
        if type(images) is list or images.ndim == 5:
            image_features_ls = []
            for image in images:
                outputs = self.vision_tower(image.to(device=self.device, dtype=self.dtype).unsqueeze(0), output_hidden_states=True)
                image_feature = self.feature_select(outputs).to(image.dtype)
                image_feature = image_feature[0, filter_useless_channel(outputs.logits)[0], :]
                image_features_ls.append(image_feature)
        else:
            outputs = self.vision_tower(images.to(device=self.device, dtype=self.dtype), output_hidden_states=True)
            image_features = self.feature_select(outputs).to(images.dtype)

            image_features_ls = []
            keep = filter_useless_channel(outputs.logits)
            for i in range(len(images)):
                image_feature = image_features[i, keep[i], :]
                image_features_ls.append(image_feature)

        # exp3(1/3): whether pad to the same length here
        # max_len = max([image_feature.shape[0] for image_feature in image_features_ls])
        # for i in range(len(image_features_ls)):
        #     image_features_ls[i] = torch.cat([image_features_ls[i], torch.zeros(max_len - len(image_features_ls[i]), image_features_ls[i].shape[1], device=self.device, dtype=self.dtype)], dim=0)
        # image_features_ls = torch.stack(image_features_ls, dim=0)
        return image_features_ls

    @property
    def dummy_feature(self):
        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)

    @property
    def dtype(self):
        return self.vision_tower.dtype

    @property
    def device(self):
        return self.vision_tower.device

    @property
    def config(self):
        if self.is_loaded:
            return self.vision_tower.config
        else:
            return self.cfg_only

    @property
    def hidden_size(self):
        return self.config.hidden_size

    @property
    def num_patches(self):
        return (self.config.image_size // self.config.patch_size) ** 2


def filter_useless_channel(out_logits):
    # use confidence score to filter out useless channel
    prob = nn.functional.softmax(out_logits, -1) # ([1, 100, 92])
    scores, labels = prob[..., :-1].max(-1)
    #### exp1: change the threshold
    keep_ls = scores > 0.7 # [bs, 100] 

    # ## exp2: use the top k scores and write to file
    # file_name = "./llava/backbone/results/I4_obj_dectection.json"
    # file_name2 = "./llava/backbone/results/I4_obj_interest.json"
    obj_results = {"scores": scores.tolist(), "labels": labels.tolist(),"keep_ls":keep_ls.tolist()}
    # # "keep_labels", "keep_scores", "found_scores", "people_scores"
    
    from llava.backbone.utils.convert2labels import get_keep_labels_scores, found_scores_given_labels, get_interest_obj_ls
    keep_labels, keep_scores = get_keep_labels_scores(obj_results)
    # keep_dict = {"keep_labels": keep_labels, "keep_scores": keep_scores}
    print("keep:", keep_labels, keep_scores)

    # io_ls = get_interest_obj_ls()
    # found_scores, found_scores_given_obj = found_scores_given_labels(obj_results, io_ls)
    # found_dict = {"io_ls":io_ls, "found_scores": found_scores, "found_scores_given_obj": found_scores_given_obj}
    # print("io_ls",io_ls, "scores:",found_scores_given_obj)

    # import json
    # with open(file_name, 'a') as f:
    #     json.dump(keep_dict, f, indent=4)
    # with open(file_name2, 'a') as f:
    #     json.dump(found_dict, f, indent=4)
    ### end exp2
    return keep_ls[..., 1:] # BUG: select_feature == 'cls_patch' will cause error here