import os
import cv2
import random
import lmdb
import json
import numpy as np
import torch
import torch.nn.functional as F
from pycocotools import mask
from transformers import CLIPImageProcessor
# from FCLM.model.language_model.segment_anything.utils.transforms import ResizeLongestSide
# from tools.utils import DEFAULT_IMAGE_TOKEN
# from dataset.utils.utils import ANSWER_LIST, SEG_QUESTIONS
DEFAULT_IMAGE_TOKEN = "<image>"

from copy import deepcopy
from typing import Tuple

import numpy as np
import torch
from torch.nn import functional as F
from torchvision.transforms.functional import resize  # type: ignore
from torchvision.transforms.functional import to_pil_image


class ResizeLongestSide:
    """
    Resizes images to the longest side 'target_length', as well as provides
    methods for resizing coordinates and boxes. Provides methods for
    transforming both numpy array and batched torch tensors.
    """

    def __init__(self, target_length: int) -> None:
        self.target_length = target_length

    def apply_image(self, image: np.ndarray) -> np.ndarray:
        """
        Expects a numpy array with shape HxWxC in uint8 format.
        """
        target_size = self.get_preprocess_shape(
            image.shape[0], image.shape[1], self.target_length
        )
        return np.array(resize(to_pil_image(image), target_size))

    def apply_coords(
        self, coords: np.ndarray, original_size: Tuple[int, ...]
    ) -> np.ndarray:
        """
        Expects a numpy array of length 2 in the final dimension. Requires the
        original image size in (H, W) format.
        """
        old_h, old_w = original_size
        new_h, new_w = self.get_preprocess_shape(
            original_size[0], original_size[1], self.target_length
        )
        coords = deepcopy(coords).astype(float)
        coords[..., 0] = coords[..., 0] * (new_w / old_w)
        coords[..., 1] = coords[..., 1] * (new_h / old_h)
        return coords

    def apply_boxes(
        self, boxes: np.ndarray, original_size: Tuple[int, ...]
    ) -> np.ndarray:
        """
        Expects a numpy array shape Bx4. Requires the original image size
        in (H, W) format.
        """
        boxes = self.apply_coords(boxes.reshape(-1, 2, 2), original_size)
        return boxes.reshape(-1, 4)

    def apply_image_torch(self, image: torch.Tensor) -> torch.Tensor:
        """
        Expects batched images with shape BxCxHxW and float format. This
        transformation may not exactly match apply_image. apply_image is
        the transformation expected by the model.
        """
        # Expects an image in BCHW format. May not exactly match apply_image.
        target_size = self.get_preprocess_shape(
            image.shape[0], image.shape[1], self.target_length
        )
        return F.interpolate(
            image, target_size, mode="bilinear", align_corners=False, antialias=True
        )

    def apply_coords_torch(
        self, coords: torch.Tensor, original_size: Tuple[int, ...]
    ) -> torch.Tensor:
        """
        Expects a torch tensor with length 2 in the last dimension. Requires the
        original image size in (H, W) format.
        """
        old_h, old_w = original_size
        new_h, new_w = self.get_preprocess_shape(
            original_size[0], original_size[1], self.target_length
        )
        coords = deepcopy(coords).to(torch.float)
        coords[..., 0] = coords[..., 0] * (new_w / old_w)
        coords[..., 1] = coords[..., 1] * (new_h / old_h)
        return coords

    def apply_boxes_torch(
        self, boxes: torch.Tensor, original_size: Tuple[int, ...]
    ) -> torch.Tensor:
        """
        Expects a torch tensor with shape Bx4. Requires the original image
        size in (H, W) format.
        """
        boxes = self.apply_coords_torch(boxes.reshape(-1, 2, 2), original_size)
        return boxes.reshape(-1, 4)

    @staticmethod
    def get_preprocess_shape(
        oldh: int, oldw: int, long_side_length: int
    ) -> Tuple[int, int]:
        """
        Compute the output size given input size and target long side length.
        """
        scale = long_side_length * 1.0 / max(oldh, oldw)
        newh, neww = oldh * scale, oldw * scale
        neww = int(neww + 0.5)
        newh = int(newh + 0.5)
        return (newh, neww)

class GrandReferSegmDataset(torch.utils.data.Dataset):
    CLASSES = ('object',)
    IMG_MEAN = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
    IMG_STD = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
    IMG_SIZE = 1024
    IGNORE_LABEL = 255
    # ./data/GranD/object_lvl_lmdb
    # ./data/GranD/object_lvl_lmdb/txt_files
    # ./data/GranD/SAM/SA1B_img
    def __init__(self, epoch_samples=500 * 8 * 2 * 10,
                 precision: str = "fp32", image_size: int = 1024,
                 validation=False, split='train', random_sampling=True, inference=False,
                 base_dir='./data/GranD'):
        self.epoch_samples = epoch_samples
        self.image_size = image_size
        self.precision = precision
        self.transform = ResizeLongestSide(image_size)
        self.begin_str = f"""The {DEFAULT_IMAGE_TOKEN} provides an overview of the picture.\n"""
        self.validation = validation
        self.random_sampling = random_sampling
        # Defining paths
        self.base_dir = os.path.join(base_dir, 'object_lvl_lmdb')
        # self.base_dir = os.path.join(base_dir, 'caption_grounding_lmdb')
        self.image_folder = os.path.join(base_dir, 'SAM/SA1B_img')
        ann_file_name = "Grand_Referring_Expression_lmdb"
        self.annos = lmdb.open(self.base_dir, readonly=True, max_readers=1, lock=False, readahead=False, meminit=False)
        mode = "Val" if validation else "Train"
        self.data_infos = self._load_annotations(
            os.path.join(base_dir, 'object_lvl_lmdb/txt_files')
            )
        print('\033[92m' + "----SEGM-{}: GranD Referring Segm dataset initialized----".format(mode) + '\033[0m')

    def _load_annotations(self, ann_file):
        with open(ann_file, 'r') as f:
            data_infos = [line.strip() for line in f if line.strip()]
        data_infos = data_infos[0: 1000] if self.validation else data_infos
        return data_infos

    def __len__(self):
        return len(self.data_infos)

    def grounding_enc_processor(self, x: torch.Tensor) -> torch.Tensor:
        x = (x - self.IMG_MEAN) / self.IMG_STD
        h, w = x.shape[-2:]
        x = F.pad(x, (0, self.IMG_SIZE - w, 0, self.IMG_SIZE - h))
        return x

    def create_conversations(self, labels, questions):
        questions = []
        answers = []
        for i, label in enumerate(labels):
            question = random.choice(questions)
            questions.append(question)
            answers.append(label)

        conversations = []
        conv = conversation_lib.default_conversation.copy()
        conv.messages = []
        for i, (question, answer) in enumerate(zip(questions, answers)):
            if i == 0:
                question = self.begin_str + question
            conv.append_message(conv.roles[0], question)
            conv.append_message(conv.roles[1], answer)
        conversations.append(conv.get_prompt())
        return questions, conversations

    def _parse_annotations(self,ann_info):
        annotations = {'masks': [], 'labels': [] , 'attributes': [] , 'score': []}
        
        for ann in ann_info:
            if 'score' in ann.keys():
                rle = ann.get("segmentation")
                if rle:
                    m = mask.decode(rle)
                    m = m.astype(np.uint8)
                    annotations['masks'].append(m)
                    annotations['labels'].append(ann['labels'])
                    annotations['attributes'].append(ann['attributes'])
                    annotations['score'].append(ann['score'])

        return annotations

    def __getitem__(self, idx):
        image_name = self.data_infos[idx] if (self.validation or not self.random_sampling) else self.data_infos[
            random.randint(0, len(self.data_infos) - 1)]
        image_path = os.path.join(self.image_folder, image_name)
        # Get the annotation from lmdb
        with self.annos.begin() as txn:
            json_contents = txn.get(image_name.encode())
        json_contents = json.loads(json_contents.decode('utf-8'))
        ann_info = json_contents[image_name]
        ann = self._parse_annotations(ann_info)
        
        valid_score = [i for i, v in enumerate(ann['score']) if v >0.5]
        new_valid_score = []
        for i in valid_score:
            mask=ann['masks'][i]
            if len(mask[mask==1])>100:
                # print(len(mask[mask==1]))
                new_valid_score.append(i)
            else:
                print('no mask',i)
        valid_score=new_valid_score
        data_item = {"image_path": image_path,
                     "filename": image_name,
                     "labels": [ann['labels'][i] for i in valid_score],
                     "masks": [ann['masks'][i] for i in valid_score],
                     "attributes": [ann['attributes'][i] for i in valid_score],
                     "score": [ann['score'][i] for i in valid_score]
                      }
        
        data_item=self.process_data(data_item)
        return data_item

    def process_data(self, data_item):
        
        data_labels = data_item['labels']
        data_masks = data_item['masks']
        image_path = data_item['image_path']
        data_attributes = data_item['attributes']
        data_score = data_item['score']
        image = cv2.imread(image_path)

        # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        image = self.transform.apply_image(image)
        #####################vis##################
        # for i,mask in enumerate(data_masks):
        #     mask[mask==1]=255
        #     # mask = self.transform.apply_image(mask)
        #     mask = cv2.resize(mask,(image.shape[1],image.shape[0]))
        #     mask_color = cv2.cvtColor(mask.astype(np.uint8), cv2.COLOR_GRAY2BGR)
        #     color = np.array([255, 0, 0], dtype=np.uint8)  # 红色
        #     mask_color[mask > 0] = color
        #     blended = cv2.addWeighted(image, 0.5, mask_color, 0.5, 0)
        #     cv2.imwrite(f'overlay_{i}.jpg', blended)
        #     print(f'save:overlay_{i}')
        # for i,mask in enumerate(data_masks):
        #     mask[mask==1]=255
        #     mask = self.transform.apply_image(mask)
        #     cv2.imwrite(f'mask_{i}.jpg', mask)
        #     print(f'save:mask_{i}')
        ###########################################
        ############ anno for category. ############
        valid_idx = [i for i, v in enumerate(data_labels) if v is not None]
        if len(valid_idx)>3:
            sampled_idx_label = random.sample(valid_idx, 3)
        else:
            sampled_idx_label=valid_idx
        sampled_mask_labels = [data_masks[i] for i in sampled_idx_label]
        # sampled_labels = [random.choice(data_labels[i]) for i in sampled_idx_label]
        sampled_labels = [" ".join(data_labels[i]) for i in sampled_idx_label]
        cleaned_items = []
        for item in sampled_labels:
            words = item.split()
            unique_words = list(dict.fromkeys(words))
            cleaned_items.append(" ".join(unique_words))
        print(cleaned_items)
        sampled_score = [data_score[i] for i in sampled_idx_label]
        for i,mask in enumerate(sampled_mask_labels):
            mask[mask==1]=255
            mask = cv2.resize(mask,(image.shape[1],image.shape[0]))
            mask_color = cv2.cvtColor(mask.astype(np.uint8), cv2.COLOR_GRAY2BGR)
            color = np.array([255, 0, 0], dtype=np.uint8)  # 红色
            mask_color[mask > 0] = color
            blended = cv2.addWeighted(image, 0.5, mask_color, 0.5, 0)
            # cv2.imwrite(f'labels_{i}_{sampled_labels[i]}.jpg', blended)
            cv2.imwrite(f'labels_{sampled_score[i]}_{sampled_labels[i]}.jpg', blended)
            print(f'save:labels_{i}')
        ############ anno for attributes. ############
        valid_idx = [i for i, v in enumerate(data_attributes) if v is not None]
        if len(valid_idx)>3:
            sampled_idx_attr = random.sample(valid_idx, 3)
        else:
            sampled_idx_attr=valid_idx
        sampled_mask_attri = [data_masks[i] for i in sampled_idx_attr]
        sampled_attributes = [random.choice(data_attributes[i]) for i in sampled_idx_attr]
        sampled_score = [data_score[i] for i in sampled_idx_attr]
        for i,mask in enumerate(sampled_mask_attri):
            mask[mask==1]=255
            mask = cv2.resize(mask,(image.shape[1],image.shape[0]))
            mask_color = cv2.cvtColor(mask.astype(np.uint8), cv2.COLOR_GRAY2BGR)
            color = np.array([255, 0, 0], dtype=np.uint8)  # 红色
            mask_color[mask > 0] = color
            blended = cv2.addWeighted(image, 0.5, mask_color, 0.5, 0)
            # cv2.imwrite(f'attributes_{i}_{sampled_attributes[i]}.jpg', blended)
            cv2.imwrite(f'attributes_{sampled_score[i]}_{sampled_attributes[i]}.jpg', blended)
            print(f'save:attributes_{i}')
        str_region = ""
        region_num=len(sampled_attributes)+len(sampled_labels)
        for i in range(region_num):
            mask = ann['annotation'][i]['segmentation']
            masks.append(mask)
            if i>0:
                str_region+= ','
            str_region += "region"+str(i+1)+"<mask><pos>"

        for i in range(len(ann['conversations'])//2):
            if i==0:
                if region_num==1:
                    mid_str = "Ther are 1 part region in the picture: "+str_region+'. '
                else:
                    mid_str = "Ther are {} part regions in the picture: ".format(str(region_num))+str_region+'. '

        # for ann in ann_list:

        #     metadata_fn_path = ann['metadata_fn']
        #     img_fn = ann['img_fn']
        #     img_path = os.path.join(self.img_prefix,img_fn)
        #     annotations = json.load(open(os.path.join(self.img_prefix, metadata_fn_path)))
        #     masks = annotations['segms']
        #     bboxes = np.array(annotations['boxes'])

        #     class_names = ann['objects']
        #     num_objects = len(class_names)
        #     ref_string = ''
        #     for i in range(num_objects):
        #         ref_string = ref_string +  f'region{i+1} <mask><pos>' + ','
        #     ref_string = ref_string[:-1]
        #     ref_prefix = random.choice(Ref_WAY)

        #     begion_string = ref_prefix.replace('<region>', ref_string)
        #     qa_s = []

        #     q = ann['question_orig']
        #     q = replace_numbers_with_tags(q, class_names)
        #     a = ann['answer_orig']
        #     a = replace_numbers_with_tags(a, class_names)
        #     why = ann['rationale_orig']
        #     why = replace_numbers_with_tags(why, class_names)
        #     if (q is None) or (a is None) or (why) is None:
        #         continue


        #     qa_s.append({'from': 'human', 'value': begion_string + q})
        #     qa_s.append({'from': 'gpt', 'value': a})
        #     qa_s.append({'from': 'human', 'value': random.choice(WHY_QUESTIONS)})
        #     qa_s.append({'from': 'gpt', 'value': why})
        #     data_infos.append(dict(
        #         img_path = img_path,
        #         bboxes = bboxes,
        #         masks = masks,
        #         labels= class_names,
        #         qas = qa_s)
        #     )


        
if __name__ == "__main__":
    dataset=GrandReferSegmDataset()
    dataset[0]
    dataset[1]