import copy
import random
import os
import numpy as np
import torch
import ast
from FCLM.train.train import preprocess, preprocess_multimodal
from PIL import Image
from torch.utils.data import Dataset, DataLoader
import pandas as pd
from torchvision import transforms
import glob
import json
import cv2
import threading
import torch.nn.functional as F
from FCLM.model.language_model.segment_anything.utils.transforms import ResizeLongestSide
cv2.setNumThreads(0)
DEFAULT_IMAGE_TOKEN = "<image>"
SHORT_QUESTION_LIST = [
    DEFAULT_IMAGE_TOKEN + "\n" + "Can you segment the {class_name} in this image?",
    DEFAULT_IMAGE_TOKEN + "\n" + "Please segment the {class_name} in this image.",
    DEFAULT_IMAGE_TOKEN
    + "\n"
    + "What is {class_name} in this image? Please respond with segmentation mask.",
    DEFAULT_IMAGE_TOKEN
    + "\n"
    + "What is {class_name} in this image? Please output segmentation mask.",
]

LONG_QUESTION_LIST = [
    DEFAULT_IMAGE_TOKEN + "\n" + "{sent} Please respond with segmentation mask.",
    DEFAULT_IMAGE_TOKEN + "\n" + "{sent} Please output segmentation mask.",
]

EXPLANATORY_QUESTION_LIST = [
    "Please output segmentation mask and explain why.",
    "Please output segmentation mask and explain the reason.",
    "Please output segmentation mask and give some explanation.",
]

ANSWER_LIST = [
    "It is [SEG].",
    "Sure, [SEG].",
    "Sure, it is [SEG].",
    "Sure, the segmentation result is [SEG].",
    "[SEG].",
]

def resize_longest_side(image, target_size):
    """
    调整图像大小，保持纵横比，将最长边缩放到 target_size
    
    参数:
        image: PIL.Image 对象
        target_size: 最长边的目标尺寸（整数）
    
    返回:
        调整后的 PIL.Image 对象
    """
    # 获取原始图像尺寸
    width, height = image.size
    
    # 计算缩放比例
    max_dim = max(width, height)
    scale = target_size / max_dim
    
    # 计算新尺寸（保持整数像素）
    new_width = int(round(width * scale))
    new_height = int(round(height * scale))
    
    # 调整图像大小（使用高质量的 Lanczos 滤镜进行缩放）
    resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    return resized_image
def resize_mask_longest_side(mask_tensor, target_size):
    """
    调整0-1 mask张量的大小，保持纵横比，使用最近邻插值（纯PyTorch实现）
    """
    # 处理维度：确保输入是 (B, C, H, W) 格式（interpolate需要批量维度）
    if len(mask_tensor.shape) == 2:  # (H, W) → (1, 1, H, W)
        mask_tensor = mask_tensor.unsqueeze(0).unsqueeze(0)
    elif len(mask_tensor.shape) == 3:  # (C, H, W) → (1, C, H, W)
        mask_tensor = mask_tensor.unsqueeze(0)
    
    _, _, height, width = mask_tensor.shape
    max_dim = max(height, width)
    scale = target_size / max_dim
    new_height, new_width = int(round(height * scale)), int(round(width * scale))
    
    # 使用 interpolate 进行缩放，mode指定为'nearest'
    resized_mask = F.interpolate(
        mask_tensor,
        size=(new_height, new_width),
        mode='nearest-exact'  # 精确最近邻，避免中间值
    )
    
    # 恢复原始维度（去除批量维度）
    if resized_mask.shape[0] == 1:
        resized_mask = resized_mask.squeeze(0)
    
    return resized_mask
def get_mask_from_json(json_path, img):
    try:
        with open(json_path, "r") as r:
            anno = json.loads(r.read())
    except:
        with open(json_path, "r", encoding="cp1252") as r:
            anno = json.loads(r.read())

    inform = anno["shapes"]
    comments = anno["text"]
    is_sentence = anno["is_sentence"]

    height, width = img.shape[:2]

    ### sort polies by area
    area_list = []
    valid_poly_list = []
    for i in inform:
        label_id = i["label"]
        points = i["points"]
        if "flag" == label_id.lower():  ## meaningless deprecated annotations
            continue

        tmp_mask = np.zeros((height, width), dtype=np.uint8)
        cv2.polylines(tmp_mask, np.array([points], dtype=np.int32), True, 1, 1)
        cv2.fillPoly(tmp_mask, np.array([points], dtype=np.int32), 1)
        tmp_area = tmp_mask.sum()

        area_list.append(tmp_area)
        valid_poly_list.append(i)

    ### ground-truth mask
    sort_index = np.argsort(area_list)[::-1].astype(np.int32)
    sort_index = list(sort_index)
    sort_inform = []
    for s_idx in sort_index:
        sort_inform.append(valid_poly_list[s_idx])

    mask = np.zeros((height, width), dtype=np.uint8)
    for i in sort_inform:
        label_id = i["label"]
        points = i["points"]

        if "ignore" in label_id.lower():
            label_value = 255  # ignored during evaluation
        else:
            label_value = 1  # target

        cv2.polylines(mask, np.array([points], dtype=np.int32), True, label_value, 1)
        cv2.fillPoly(mask, np.array([points], dtype=np.int32), label_value)

    return mask, comments, is_sentence
class REASON_SEG(Dataset):
    
    def __init__(self,
                tokenizer=None,
                data_args=None,
                base_image_dir=None,
                # samples_per_epoch=500 * 8 * 2 * 10,
                samples_per_epoch=500 * 8 * 1 * 10,
                precision: str = "fp32",
                # image_size: int = 224,
                num_classes_per_sample: int = 3,
                exclude_val=False,
                reason_seg_data="ReasonSeg|train",
                explanatory=0.1,
                 ):
        self.INPUT_PIXEL=1024
        self.pixel_mean = torch.Tensor([123.675, 116.28, 103.53]).view(-1, 1, 1)
        self.pixel_std = torch.Tensor([58.395, 57.12, 57.375]).view(-1, 1, 1)
        self.img_size = 1024
        self.data_args = data_args
        self.tokenizer = tokenizer
        self.exclude_val = exclude_val
        self.reason_seg_data = reason_seg_data
        self.samples_per_epoch = samples_per_epoch
        self.explanatory = explanatory
        self.num_classes_per_sample = num_classes_per_sample

        self.base_image_dir = base_image_dir
        self.precision = precision
        self.transform = ResizeLongestSide(self.img_size)
        self.short_question_list = SHORT_QUESTION_LIST
        self.long_question_list = LONG_QUESTION_LIST
        self.answer_list = ANSWER_LIST
        
        reason_seg_data, splits = reason_seg_data.split("|")
        splits = splits.split("_")
        images = []
        for split in splits:
            images_split = glob.glob(
                os.path.join(
                    base_image_dir, split, "*.jpg"
                )
            )
            images.extend(images_split)
        jsons = [path.replace(".jpg", ".json") for path in images]
        self.reason_seg_data = (images, jsons)
        print(f"\033[1;32mnumber of reason_seg samples: {len(images)} input res:{self.INPUT_PIXEL}\033[0m")
        
        if explanatory != -1:
            self.explanatory_question_list = EXPLANATORY_QUESTION_LIST
            self.img_to_explanation = {}
            with open(
                os.path.join(
                    base_image_dir,
                    "explanatory",
                    "train.json",
                )
            ) as f:
                items = json.load(f)
            
            for item in items:
                img_name = item["image"]
                self.img_to_explanation[img_name] = {
                    "query": item["query"],
                    "outputs": item["outputs"],
                }
        super().__init__()
    def __len__(self):
        return self.samples_per_epoch
    def __getitem__(self, idx):
        data_item = self.get_data_item(idx)
        return data_item
    def preprocess_pad(self, x: torch.Tensor,if_img=True) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        # Normalize colors
        if if_img:
            x = (x - self.pixel_mean) / self.pixel_std

        # Pad
        h, w = x.shape[-2:]
        padh = self.img_size - h
        padw = self.img_size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x
    def get_data_item(self, idx):
        data_dict={}

        images, jsons = self.reason_seg_data
        # idx = random.randint(0, len(images) - 1)
        idx = idx % len(images)
        img_path = images[idx]
        json_path = jsons[idx]
        image = Image.open(img_path).convert('RGB')
        
        mask, sents, is_sentence = get_mask_from_json(json_path, cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))
        normalize = transforms.Compose([
            transforms.ToTensor(),                      # 转为Tensor并自动归一化到[0,1]
            transforms.Normalize(
                mean=[0.485, 0.456, 0.406],            # ImageNet均值
                std=[0.229, 0.224, 0.225]              # ImageNet标准差
            )
        ])
        image = normalize(image) 
        data_dict['img'] = image 
        if len(sents) >= self.num_classes_per_sample:
            sampled_inds = np.random.choice(
                list(range(len(sents))), size=self.num_classes_per_sample, replace=False
            )
        else:
            sampled_inds = list(range(len(sents)))
        sampled_sents = np.vectorize(sents.__getitem__)(sampled_inds).tolist()
        sampled_masks = [
            (mask == 1).astype(np.float32) for _ in range(len(sampled_inds))
        ]
        image_name = img_path.split("/")[-1]
        # if self.explanatory != -1 and image_name in self.img_to_explanation:
        #     if random.random() < self.explanatory:
        #         choice = 2
        #     else:
        #         choice = random.randint(0, 1)
        # else:
        #     choice = 0
        # choice = random.randint(0, 1)
        if random.random() < self.explanatory:
            choice = 1
        else:
            choice = 0 
        questions = []
        answers = []
        for text in sampled_sents:
            if is_sentence:
                question_template = random.choice(self.long_question_list)
                questions.append(question_template.format(sent=text))
            else:
                question_template = random.choice(self.short_question_list)
                questions.append(question_template.format(class_name=text.lower()))

            # add explanation if applicable
            img_name = img_path.split("/")[-1]
            if self.explanatory != -1 and img_name in self.img_to_explanation:
                if choice == 0:  # [SEG] token
                    answers.append(random.choice(self.answer_list))
                    #['It is [SEG].', 'Sure, [SEG].', 'Sure, it is [SEG].', 'Sure, the segmentation result is [SEG].', '[SEG].']
                elif choice == 1:  # [SEG] token + text answer
                    image_name = img_path.split("/")[-1]
                    answer = self.img_to_explanation[image_name]["outputs"]
                    answer = random.choice(self.answer_list) + " {}".format(answer)
                    questions[-1] = (
                        DEFAULT_IMAGE_TOKEN
                        + "\n"
                        + text
                        + " {}".format(random.choice(self.explanatory_question_list))
                    )
                    #'<image>\nWhat is the object that is typically used during the hurdling event in track and field? Please output segmentation mask and give some explanation.'
                    answers.append(answer)
                    #"It is [SEG]. In the picture, there is a hurdle, which is typically used in the hurdling event of track and field competitions. The athlete in the image is jumping over this hurdle during the competition. Hurdles are designed to challenge the athletes' agility, speed, and jumping ability as they attempt to clear them while running at high speeds."
                elif choice == 2:  # vanilla text answer
                    image_name = img_path.split("/")[-1]
                    answer = self.img_to_explanation[image_name]["outputs"]
                    questions[-1] = DEFAULT_IMAGE_TOKEN + "\n" + text
                    #'<image>\nWhat is the object that is typically used during the hurdling event in track and field?'
                    answers.append(answer)
                    #"In the picture, there is a hurdle, which is typically used in the hurdling event of track and field competitions. The athlete in the image is jumping over this hurdle during the competition. Hurdles are designed to challenge the athletes' agility, speed, and jumping ability as they attempt to clear them while running at high speeds."
                else:
                    raise ValueError("Not implemented yet.")
            else:
                answers.append(random.choice(self.answer_list))

            conversations = []
            sources = dict()
            sources['conversations'] = []
            for i in range(len(questions)):
                cur_sample=[]
                cur_sample.append(
                        {'from': 'human', 'value': questions[i]})
                cur_sample.append({'from': 'gpt', 'value': answers[i]})
                sources['conversations'].append(cur_sample)
            # print('sources:',sources)

        processor=self.data_args.image_processor
        qwen_img=Image.open(img_path).convert('RGB')

        data_dict_bs=dict()
        data_dict_bs['input_ids']=[]
        data_dict_bs['labels']=[]
        for idx in range(len(questions)):
            
            sources_each = preprocess_multimodal(
                copy.deepcopy([sources['conversations'][idx]]),
                self.data_args,
                0,if_qwen=True)
        
            text=sources_each[0][0]['value']
        
            W,H=qwen_img.size
            
            qwen_img_process=resize_longest_side(qwen_img,self.INPUT_PIXEL)
            inputs = processor(
                text=[text],
                images=qwen_img_process,
                videos=None,
                padding=True,
                return_tensors="pt",
            )
            
            tmp=inputs['input_ids']
            N=len(tmp[tmp==self.tokenizer('<|image_pad|>')['input_ids'][0]]) 
            
            data_dict = preprocess(
                sources_each,
                self.tokenizer,
                has_image=True
                )
            ####repeat <|image_pad|> * N
            old_data_dict=data_dict
            insert_position=torch.where(old_data_dict['input_ids']==self.tokenizer('<|image_pad|>')['input_ids'][0])[1]
            img_padding=torch.tensor([self.tokenizer('<|image_pad|>')['input_ids'][0]]*N).unsqueeze(0)
            target_padding=torch.tensor([-100]*N).unsqueeze(0)
            # (1,N) <|image_pad|>
            new_inputs=torch.cat((old_data_dict['input_ids'][:,:insert_position],img_padding,old_data_dict['input_ids'][:,insert_position+1:]),dim=1)
            # [1, 435] 345 个 image_pad
            new_target=torch.cat((old_data_dict['labels'][:,:insert_position],target_padding,old_data_dict['labels'][:,insert_position+1:]),dim=1)
            # data_dict['input_ids']=new_inputs #torch.Size([1, 1628])
            # data_dict['labels']=new_target #torch.Size([1, 1628])
            data_dict_bs['input_ids'].append(new_inputs.squeeze(0))
            data_dict_bs['labels'].append(new_target.squeeze(0))
            #############################
        # 
        
        
        sam_image = self.transform.apply_image(cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB))
        resize = sam_image.shape[:2]
        sam_image = self.preprocess_pad(torch.from_numpy(sam_image).permute(2, 0, 1).contiguous())
        # data_dict['image'] = sam_image #torch.Size([3, 1024, 1024])
        
        masks = np.stack(sampled_masks, axis=0)
        masks = torch.from_numpy(masks)
        resize_mask_longest_side(masks,self.INPUT_PIXEL)
        sample_num=len(questions)
        data_dict_bs['image'] = sam_image.unsqueeze(0).repeat(sample_num,1,1,1)
        data_dict_bs['masks']=masks
        data_dict_bs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
        data_dict_bs['image_grid_thw'] = inputs['image_grid_thw']
        data_dict_bs['img_metas'] = dict()
        data_dict_bs['img_metas']['image'] = torch.tensor(resize).unsqueeze(0).repeat(sample_num,1)
        data_dict_bs['img_metas']['sample_num'] = sample_num
        # data_dict_bs['img_metas']['SEG_num'] = 1
        data_dict_bs['img_metas']['SEG_num'] = sample_num
        # return data_dict
        return data_dict_bs
