#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License is distributed on an "AS IS" BASIS,
#    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#    See the License for the specific language governing permissions and
#    limitations under the License.
import numpy as np
import os
from .segment_anything import build_sam_vit_h,build_sam_vit_b
from transformers.models.qwen2_vl.modeling_qwen2_vl import (
    PatchEmbed,
    PatchMerger,
    Qwen2RMSNorm,
    Qwen2VLCausalLMOutputWithPast,
    Qwen2VLForConditionalGeneration,
    Qwen2VLModel,
    Qwen2VLPreTrainedModel,
    VisionAttention,
    VisionRotaryEmbedding,
    VisionSdpaAttention,
)
import cv2
from typing import List, Optional, Tuple, Union
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.nn import CrossEntropyLoss

from transformers import AutoConfig, AutoModelForCausalLM, \
                         LlamaConfig, LlamaModel, LlamaForCausalLM

from transformers import AutoConfig, AutoModelForCausalLM
from qwen2_5_vl import  Qwen2_5_VLConfig, Qwen2_5_VLModel,Qwen2_5_VLForConditionalGeneration
from qwen2_5_vl.modeling_qwen2_5_vl  import Qwen2_5_VisionTransformerPretrainedModel
from transformers.modeling_outputs import CausalLMOutputWithPast

from ..osprey_arch import OspreyMetaModel, OspreyMetaForCausalLM

from ..layer import Fine_Grained_Visual_Prompt_Extractor

###use sam2
# from sam2.build_sam import build_sam2

def latent_kl_between_features(z_pred, z_gt):
    """
    计算 z_pred 和 z_gt 两个特征分布之间的 KL 散度，
    假设两者都服从对角高斯分布，分布参数通过 batch 内统计得到。

    假设：
        q(z) = N(mu_q, sigma_q^2), mu_q = z_pred.mean(dim=0), sigma_q^2 = z_pred.var(dim=0)
        p(z) = N(mu_p, sigma_p^2), mu_p = z_gt.mean(dim=0), sigma_p^2 = z_gt.var(dim=0)

    参数:
        z_pred: Tensor, shape [batch_size, feature_dims]
        z_gt: Tensor, shape [batch_size, feature_dims]

    返回:
        kl_loss: Tensor, 标量，表示两个分布之间的 KL 散度
    """
    # 计算均值和方差
    mu_q = z_pred.mean(dim=0)      # [dims]
    var_q = z_pred.var(dim=0, unbiased=False)  # [dims]

    mu_p = z_gt.mean(dim=0)        # [dims]
    var_p = z_gt.var(dim=0, unbiased=False)  # [dims]

    # 避免除以零或 log(0)
    var_p = torch.clamp(var_p, min=1e-8)
    var_q = torch.clamp(var_q, min=1e-8)

    # 计算每个维度的 KL 散度（对角协方差）
    kl_per_dim = 0.5 * (
        torch.log(var_p / var_q) +
        (var_q + (mu_q - mu_p) ** 2) / var_p -
        1
    )

    # 对所有维度求和，得到每个样本的 KL，这里直接对所有维度求和
    kl_loss = kl_per_dim.mean()  # 或者 .mean() 如果你希望 per-dim 平均

    return kl_loss
def cosine_similarity_loss(z_pred, z_gt):
    """
    计算预测特征 z_pred 和目标特征 z_gt 之间的 Cosine 相似度 Loss

    参数:
        z_pred: Tensor, shape [batch_size, feature_dim]
        z_gt: Tensor, shape [batch_size, feature_dim]

    返回:
        loss: Tensor, 标量，值为 1 - cos(z_pred, z_gt)，越小表示越相似
    """
    # 计算每个样本的余弦相似度 [bs]
    cos_sim = torch.cosine_similarity(z_pred, z_gt, dim=1)  # shape: [bs]

    # loss = 1 - cos_sim，希望它接近 0（即 cos_sim 接近 1）
    loss = 1.0 - cos_sim

    # 对 batch 取平均
    return loss.mean()
def vis_pred_and_gt(pred_mask,gt,sent):
    num=pred_mask.shape[0]
    for i in range(num):
        save_gt_img=f'./output/ref_seg_output/{i}_{sent}.jpg'
        pred=pred_mask[i].cpu().detach().numpy()
        pred=np.where(pred > 0, 255, 0).astype(np.uint8)
        cur_gt=gt[i].to(dtype=torch.int32).cpu().detach().numpy().astype(np.uint8)
        cur_gt=cur_gt*255
        save_img = np.concatenate((pred, cur_gt), axis=1)
        cv2.imwrite(save_gt_img,save_img)
        print("{} has been saved.".format(save_gt_img))
def count_parameters(model, module_name=""):
    total_params = 0
    trainable_params = 0
    
    for name, param in model.named_parameters():
        # import pdb; pdb.set_trace()
        num_params = param.numel()
        total_params += num_params
        if param.requires_grad:
            trainable_params += num_params
            
    return {
        "module_name": module_name,
        "total_params": total_params,
        "trainable_params": trainable_params,
        "frozen_params": total_params - trainable_params
    }

def print_model_params(model, module_name=""):
    stats = count_parameters(model, module_name)
    
    def format_number(num):
        if num >= 1e9:
            return f"{num/1e9:.2f}B"
        elif num >= 1e6:
            return f"{num/1e6:.2f}M"
        elif num >= 1e3:
            return f"{num/1e3:.2f}K"
        return str(num)
    
    # print("\n" + "="*50)
    # print(f"\n{stats['module_name']} module parameters:")
    # print(f"Total parameters: {format_number(stats['total_params'])}")
    # print(f"Trainable parameters: {format_number(stats['trainable_params'])}")
    # print(f"Frozen parameters: {format_number(stats['frozen_params'])}")
    # print("\n" + "="*50)
    # 打印黄色高亮的模块参数信息
    print("\033[33m" + "\n" + "="*50 + "\033[0m")
    print("\033[33m" + f"\n{stats['module_name']} module parameters:" + "\033[0m")
    print("\033[33m" + f"Total parameters: {format_number(stats['total_params'])}" + "\033[0m")
    print("\033[33m" + f"Trainable parameters: {format_number(stats['trainable_params'])}" + "\033[0m")
    print("\033[33m" + f"Frozen parameters: {format_number(stats['frozen_params'])}" + "\033[0m")
    print("\033[33m" + "\n" + "="*50 + "\033[0m")
def tensor_to_mask(tensor, output_path='mask.png', cmap='grey'):
    """
    将PyTorch张量归一化并保存为掩码图像
    
    参数:
    tensor (torch.Tensor): 输入的张量，形状为[C, H, W]
    output_path (str): 输出图像的路径
    cmap (str): 颜色映射名称，默认为'viridis'
    """
    import numpy as np
    import matplotlib.pyplot as plt
    # 检查输入是否为PyTorch张量
    if not isinstance(tensor, torch.Tensor):
        raise TypeError("输入必须是PyTorch张量")
    
    # 检查张量维度
    if len(tensor.shape) != 3:
        raise ValueError("输入张量的形状必须为[C, H, W]")
    
    # 对每个通道进行归一化处理
    normalized_tensor = torch.zeros_like(tensor, dtype=torch.float32)
    for channel in range(tensor.shape[0]):
        channel_data = tensor[channel]
        if torch.max(channel_data) > torch.min(channel_data):
            normalized_tensor[channel] = (channel_data - torch.min(channel_data)) / (
                torch.max(channel_data) - torch.min(channel_data)
            )
    
    # 将张量转换为NumPy数组
    mask_np = normalized_tensor.cpu().detach().numpy()
    
    # 如果是单通道，使用灰度图；否则使用多通道图像
    if mask_np.shape[0] == 1:
        mask_to_save = mask_np[0]  # 单通道，形状为[H, W]
    else:
        # 对于多通道，将通道维度放在最后
        mask_to_save = np.transpose(mask_np, (1, 2, 0))
    mask_to_save[mask_to_save<0.5]=0
    # 创建图像并保存
    plt.figure(figsize=(10, 8))
    
    if mask_np.shape[0] == 1:
        # 单通道使用灰度图
        plt.imshow(mask_to_save, cmap=cmap)
    else:
        # 多通道直接显示
        plt.imshow(mask_to_save)
    
    plt.axis('off')
    plt.tight_layout()
    plt.savefig(output_path, bbox_inches='tight', pad_inches=0)
    plt.close()
    
    print(f"掩码图像已保存至: {output_path}")

def dice_loss(
    inputs: torch.Tensor,
    targets: torch.Tensor,
    num_masks: float,
    scale=1000,  # 100000.0,
    eps=1e-6,
):
    """
    Compute the DICE loss, similar to generalized IOU for masks
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    """
    inputs = inputs.sigmoid()
    inputs = inputs.flatten(1, 2)
    targets = targets.flatten(1, 2)
    numerator = 2 * (inputs / scale * targets).sum(-1)
    denominator = (inputs / scale).sum(-1) + (targets / scale).sum(-1)
    loss = 1 - (numerator + eps) / (denominator + eps)
    loss = loss.sum() / (num_masks + 1e-8)
    return loss


def sigmoid_ce_loss(
    inputs: torch.Tensor,
    targets: torch.Tensor,
    num_masks: float,
):
    """
    Args:
        inputs: A float tensor of arbitrary shape.
                The predictions for each example.
        targets: A float tensor with the same shape as inputs. Stores the binary
                 classification label for each element in inputs
                (0 for the negative class and 1 for the positive class).
    Returns:
        Loss tensor
    """
    loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction="none")
    loss = loss.flatten(1, 2).mean(1).sum() / (num_masks + 1e-8)
    return loss


@dataclass
class Qwen2_5_VLCausalLMOutputWithPast(Qwen2VLCausalLMOutputWithPast):
    pass


class OspreyConfig(Qwen2_5_VLConfig):
    model_type = "osprey"


class OspreyQwenModel(OspreyMetaModel, Qwen2_5_VLModel):
    config_class = OspreyConfig

    def __init__(self, config: Qwen2_5_VLConfig):
        super(OspreyQwenModel, self).__init__(config)


class OspreyQwenForCausalLM(Qwen2_5_VLForConditionalGeneration, OspreyMetaForCausalLM):
    config_class = OspreyConfig
    
    def __init__(self, config):
        # import ipdb;ipdb.set_trace()
        super(Qwen2_5_VLForConditionalGeneration, self).__init__(config)
        self.model = OspreyQwenModel(config)
        self.visual = Qwen2_5_VisionTransformerPretrainedModel._from_config(config.vision_config)
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
        # self.mask_extractor = MaskExtractor_Qwen()
        self.mask_extractor = Fine_Grained_Visual_Prompt_Extractor()
        self.rope_deltas = None  # cache rope_deltas here
        # Init SAM for grounding modules
        self.vision_pretrained='./checkpoints/sam_vit_h_4b8939.pth'
        self.train_mask_decoder=False
        config.out_dim = 256 #SAM
        self.config=config
        
        # Initialize weights and apply final processing
        self.post_init()
        self.bce_loss_weight = 2.0
        self.dice_loss_weight = 1.0
    def set_trainning_param(self,encoder=False,lm_head=True,mask_extractor=True):
        # Initialize modules
        if encoder==True:
            print('✅ Train Visual Encoder Parameters')
            for param in self.visual.parameters():
                param.requires_grad = True
        else:
            print('🔒 Froze Visual Encoder Parameters')
            for param in self.visual.parameters():
                param.requires_grad = False
        if lm_head==True:
            print('✅ Train lm_head')
            for param in self.lm_head.parameters():
                param.requires_grad = True
        else:
            print('🔒 Froze lm_head')
            for param in self.lm_head.parameters():
                param.requires_grad = False
        # Multi-GPU training support
        if mask_extractor==True:
            print('✅ Train mask_extractor')
            for param in self.mask_extractor.parameters():
                param.requires_grad = True
        else:
            print('🔒 Froze mask_extractor')
            for param in self.mask_extractor.parameters():
                param.requires_grad = False

        print('Params : visual')
        print_model_params(self.visual)
        print('Params : lm_head')
        print_model_params(self.lm_head)
        print('Params : LLM model')
        print_model_params(self.model)
        print('Params : MaskExtractor')
        print_model_params(self.mask_extractor)
    def get_visual_embs(self, pixel_values: torch.FloatTensor):
        with torch.no_grad():
            image_embeddings_list = []
            for i in range(pixel_values.shape[0]):
                torch.cuda.empty_cache()
                image_embeddings = self.visual_model.image_encoder(
                    pixel_values[i].unsqueeze(0)
                ) #SAM1
                # image_embeddings = self.visual_model.image_encoder(
                #     pixel_values[i].unsqueeze(0)
                # )['vision_features'] #SAM2
                image_embeddings_list.append(image_embeddings)
            torch.cuda.empty_cache()
            image_embeddings = torch.cat(image_embeddings_list, 0)
        return image_embeddings
    def initialize_grounding_modules(self,sam_pth=None,projector_pth=None):
        if sam_pth is not None:
            self.vision_pretrained=sam_pth
        print('freeze visual parameters')
        for param in self.visual.parameters():
            param.requires_grad = False
        self.visual_model = build_sam_vit_h(self.vision_pretrained).to(dtype=self.visual.dtype,device=self.visual.device)
        for param in self.visual_model.parameters():
            param.requires_grad = False
        if self.train_mask_decoder:
            self.visual_model.mask_decoder.train()
            for param in self.visual_model.mask_decoder.parameters():
                param.requires_grad = True
        # Projection layer
        in_dim = self.config.hidden_size
        out_dim = self.config.out_dim
        text_fc = [
            nn.Linear(in_dim, in_dim),
            nn.ReLU(inplace=True),
            nn.Linear(in_dim, out_dim),
            nn.Dropout(0.0),
        ]
        self.text_hidden_fcs = nn.ModuleList([nn.Sequential(*text_fc)])
        if projector_pth is not None and os.path.exists(projector_pth):
            self.text_hidden_fcs.load_state_dict(torch.load(projector_pth))
        self.text_hidden_fcs.train()
        for param in self.text_hidden_fcs.parameters():
            param.requires_grad = True
        print(f'SAM initialized from {self.vision_pretrained}')
        print("Params : visual_model")
        print_model_params(self.visual_model)
        print(f"FG_Visual_Prompt_Extractor load from {self.vision_pretrained} image encoder")
        # self.mask_extractor.set_local_encoder(self.visual_model.image_encoder)
    def get_model(self):
        return self.model
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        img_metas = None,
        masks = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        images: Optional[torch.FloatTensor] = None,
        return_dict: Optional[bool] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        rope_deltas: Optional[torch.LongTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        second_per_grid_ts: Optional[torch.Tensor] = None,
    ) -> Union[Tuple, Qwen2_5_VLCausalLMOutputWithPast]:
        ####
        mask_bce_loss=None
        mask_dice_loss=None
        lm_loss=None
        self_supervised_seg_loss=None
        latent_gt_ls=[]
        # self_supervised_lm_loss=None
        ####
        seg_token=self.tokenizer.encode('[SEG]')[0]
        mask_token=self.tokenizer.encode('<mask>')[0]
        mask_token_num=len(input_ids[input_ids==mask_token])
        seg_task=seg_token in input_ids
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        if inputs_embeds is None:
            inputs_embeds = self.model.embed_tokens(input_ids)
            if pixel_values is not None:
                N_sample=inputs_embeds.shape[0]
                bs=len(pixel_values)
                new_pixel_values=torch.concat(pixel_values,dim=0).type(self.visual.dtype)
                new_image_grid_thw=torch.stack(image_grid_thw)
                concat_image_embeds=self.visual(new_pixel_values, grid_thw=new_image_grid_thw)
                concat_image_embeds=concat_image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                each_pixel_ls=np.array([(each[2].item()//2)*(each[1].item()//2) for each in image_grid_thw])
                sum_pixel=each_pixel_ls.sum()
                image_embeds_bs=torch.split(concat_image_embeds,each_pixel_ls.tolist())
                assert sum_pixel==concat_image_embeds.shape[0] , f"sum_pixel:{sum_pixel}, concat_image_embeds:{concat_image_embeds.shape[0]} not equal."
                n_image_tokens = (input_ids == self.config.image_token_id).sum().item() #"<|image_pad|>"
                n_image_features = concat_image_embeds.shape[0]
                if n_image_tokens != n_image_features:
                    raise ValueError(
                        f"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}"
                    )
                for idx in range(bs):
                    mask = input_ids[idx].unsqueeze(0) == self.config.image_token_id
                    mask_unsqueezed = mask.unsqueeze(-1)
                    mask_expanded = mask_unsqueezed.expand_as(inputs_embeds[idx].unsqueeze(0)) #[1, 288, 2048]
                    image_mask = mask_expanded.to(inputs_embeds.device) #[1, 288, 2048]
                    cur_image_embeds=image_embeds_bs[idx]
                    inputs_embeds[idx] = inputs_embeds[idx].unsqueeze(0).masked_scatter(image_mask, cur_image_embeds).squeeze(0)
                ###################################################################
                
                bs_i,mask_idx=torch.where(input_ids==self.tokenizer('<mask>')['input_ids'][0])
                bs_i,pos_idx=torch.where(input_ids==self.tokenizer('<pos>')['input_ids'][0])
                
                self_supervised_seg_loss={'self_bce':0,'self_dice':0}
                if masks is not None and len(pos_idx)>0:
                    for idx in range(bs):
                        if idx%2==1: #SEG
                            # per_sample=img_metas['sample_num']
                            # per_sample=np.insert(per_sample,0,0)
                            # per_sample=per_sample.cumsum(-1) #array([0, 3, 5])
                            seg_token_mask = input_ids[:,1:]==self.tokenizer.encode('[SEG]')[0] # [batch_size, seq_length-1]
                            seg_token_mask=torch.cat(
                                [
                                    seg_token_mask,
                                    torch.zeros(seg_token_mask.shape[0],1).bool().to(seg_token_mask.device),
                                ],
                                dim=1,
                                )#[batch_size, seq_length
                        else:
                            mask_feats,pos_feats,aux_loss = self.mask_extractor(masks[idx],image_embeds_bs[idx],image_grid_thw[idx].unsqueeze(0),inputs_embeds[idx,mask_idx[bs_i==idx]].unsqueeze(0),images[idx],self.visual_model,self.text_hidden_fcs,self.lm_head)
                            self_supervised_seg_loss['self_bce']+=aux_loss[0]
                            self_supervised_seg_loss['self_dice']+=aux_loss[1]
                            assert len(pos_feats)==len(pos_feats) 
                            for i,j in enumerate(pos_idx[bs_i==idx]):
                                inputs_embeds[idx,j]=pos_feats[i]
                            for i,j in enumerate(mask_idx[bs_i==idx]):
                                inputs_embeds[idx,j]=mask_feats[i]
                            assert torch.equal(pos_idx-mask_idx, torch.ones_like(pos_idx)) ,\
                                f"pos_idx and mask_idx should be consecutive"
                            latent_gt_ls.append(pos_feats)
            if pixel_values_videos is not None:
                pixel_values_videos = pixel_values_videos.type(self.visual.dtype)
                video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw)
                n_video_tokens = (input_ids == self.config.video_token_id).sum().item()
                n_video_features = video_embeds.shape[0]
                if n_video_tokens != n_video_features:
                    raise ValueError(
                        f"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}"
                    )

                mask = input_ids == self.config.video_token_id
                mask_unsqueezed = mask.unsqueeze(-1)
                mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)
                video_mask = mask_expanded.to(inputs_embeds.device)

                video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)
                inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)

            if attention_mask is not None:
                attention_mask = attention_mask.to(inputs_embeds.device)

        if inputs_embeds is not None:
            inputs_embeds = inputs_embeds.bfloat16()
        
        self.model = self.model.bfloat16()

        # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme
        if position_ids is None and (attention_mask is None or attention_mask.ndim == 2):
            # calculate RoPE index once per generation in the pre-fill stage only
            if (
                (cache_position is not None and cache_position[0] == 0)
                or self.rope_deltas is None
                or (past_key_values is None or past_key_values.get_seq_length() == 0)
            ):
                position_ids, rope_deltas = self.get_rope_index(
                    input_ids,
                    image_grid_thw,
                    video_grid_thw,
                    second_per_grid_ts,
                    attention_mask,
                )
                self.rope_deltas = rope_deltas
            # then use the prev pre-calculated rope-deltas to get the correct position ids
            else:
                batch_size, seq_length, _ = inputs_embeds.shape
                delta = (
                    (cache_position[0] + self.rope_deltas).to(inputs_embeds.device)
                    if cache_position is not None
                    else 0
                )
                position_ids = torch.arange(seq_length, device=inputs_embeds.device)
                position_ids = position_ids.view(1, -1).expand(batch_size, -1)
                if cache_position is not None:  # otherwise `deltas` is an int `0`
                    delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0)
                position_ids = position_ids.add(delta)
                position_ids = position_ids.unsqueeze(0).expand(3, -1, -1)
        if inputs_embeds is not None:
            input_ids = None
        outputs = self.model(
            input_ids=input_ids, #[1, 276]
            attention_mask=attention_mask,  #bool [1,276]
            past_key_values=past_key_values, #None
            inputs_embeds=inputs_embeds, #None
            use_cache=use_cache, #None
            output_attentions=output_attentions, #False
            output_hidden_states=output_hidden_states, #False
            return_dict=return_dict #True
        )
        
        hidden_states = outputs[0]
        self.lm_head = self.lm_head.to(hidden_states.dtype)
        logits = self.lm_head(hidden_states)
        
        loss = None
        lantent_loss=0
        if labels is not None:
            # Shift so that tokens < n predict n
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            # Flatten the tokens
            loss_fct = CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.config.vocab_size)
            shift_labels = shift_labels.view(-1)
            # Enable model/pipeline parallelism
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)
            lm_loss = loss
            if len(seg_token_mask[seg_token_mask==True])>0:  
                select_seg_embed = hidden_states[seg_token_mask]
                latent_gt=torch.cat(latent_gt_ls)
                
                lantent_loss=cosine_similarity_loss(select_seg_embed,latent_gt)
                # lantent_loss = F.mse_loss(select_seg_embed, latent_gt)


                # pred_log=F.log_softmax(select_seg_embed,dim=1)
                # gt_los = F.softmax(latent_gt,dim=1)
                # kl = nn.KLDivLoss(reduction='batchmean')
                # lantent_loss = kl(pred_log, gt_los)

                SAM_adaptor_states = self.text_hidden_fcs[0](select_seg_embed)
                mask_bce_loss = 0
                mask_dice_loss = 0
                num_masks = 0
                pred_mask_ls=[]
                for bs_i in range(1,bs,2):
                    sam_image_embeddings=self.get_visual_embs(images[bs_i][0].unsqueeze(0).contiguous()) #[1, 3, 1024, 1024] - > [1, 256, 64, 64]                    
                    (
                        sparse_embeddings, #[3, 1, 256]
                        dense_embeddings, #[3, 256, 64, 64]
                    ) = self.visual_model.prompt_encoder(
                        points=None,
                        boxes=None,
                        masks=None,
                        text_embeds=SAM_adaptor_states.unsqueeze(1),
                    )
                    sparse_embeddings = sparse_embeddings.to(SAM_adaptor_states.dtype)
                    low_res_masks, iou_predictions = self.visual_model.mask_decoder(
                        image_embeddings=sam_image_embeddings, #[2, 256, 64, 64]
                        image_pe=self.visual_model.prompt_encoder.get_dense_pe(),
                        sparse_prompt_embeddings=sparse_embeddings, #[3, 256, 64, 64]
                        dense_prompt_embeddings=dense_embeddings, #[3, 1, 256]
                        multimask_output=False,
                    ) #[3, 1, 256, 256]
                    pred_mask = self.visual_model.postprocess_masks(
                        low_res_masks, #[3, 1, 256, 256]
                        input_size=img_metas['image'][bs_i//2], #[686, 1024]
                        original_size=masks[bs_i][0].shape #[500, 358]
                    ) # [3, 1, 500, 358]
                    pred_mask_ls.append(pred_mask)
                    gt_mask=masks[bs_i].to(dtype=hidden_states.dtype,device=hidden_states.device)
                    assert (
                        gt_mask.shape[-2:] == pred_mask.shape[-2:]
                    ), "gt_mask.shape: {}, pred_mask.shape: {}".format(
                        gt_mask.shape, pred_mask.shape
                    )
                    mask_bce_loss +=  (
                        sigmoid_ce_loss(pred_mask.squeeze(1), gt_mask, num_masks=gt_mask.shape[0])
                        * gt_mask.shape[0]
                    )
                    mask_dice_loss += (
                        dice_loss(pred_mask.squeeze(1), gt_mask, num_masks=gt_mask.shape[0])
                        * gt_mask.shape[0]
                    )
                    ###viz###
                    # vis_pred_and_gt(pred_mask.squeeze(1),gt_mask,f'bce_{mask_bce_loss:.2f}_dice_{mask_dice_loss:.2f}')
                    #########
                    num_masks += gt_mask.shape[0]
                    # import ipdb;ipdb.set_trace()
                mask_bce_loss = self.bce_loss_weight * mask_bce_loss / (num_masks + 1e-8)
                mask_dice_loss = self.dice_loss_weight * mask_dice_loss / (num_masks + 1e-8)
                mask_loss = mask_bce_loss + mask_dice_loss
                pred=shift_logits.argmax(-1)
                seg_num=len(pred[pred==self.tokenizer('[SEG]')['input_ids'][0]])
                # logger.info(f'LLM_loss:{loss}\tDice_loss:{mask_dice_loss.item()}\tBCE_loss:{mask_bce_loss.item()}\tGT_seg_num:{len(seg_token_mask[seg_token_mask==True])}\tPred_seg_num:{seg_num}')
                loss+=mask_loss
                
            loss+=lantent_loss

            # self_supervised_seg_loss_bce = self.bce_loss_weight * self_supervised_seg_loss['self_bce'] / (mask_token_num + 1e-8)
            # self_supervised_seg_loss_dice = self.dice_loss_weight * self_supervised_seg_loss['self_dice'] / (mask_token_num + 1e-8)
            # self_supervised_seg_loss = self_supervised_seg_loss_bce + self_supervised_seg_loss_dice
            # loss += self_supervised_seg_loss
        if not return_dict:
            output = (logits,) + outputs[1:]

            return (loss,) + output if loss is not None else output

        
        return Qwen2_5_VLCausalLMOutputWithPast(
            loss=loss,
            logits=logits,
            mask_bce_loss=mask_bce_loss,
            mask_dice_loss=mask_dice_loss,
            self_supervised_seg_loss=0,
            # self_supervised_seg_loss=self_supervised_seg_loss,
            lantent_loss=lantent_loss,
            # self_supervised_lm_loss=self_supervised_lm_loss,
            lm_loss=lm_loss,
            past_key_values=outputs.past_key_values,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )

    
    def prepare_inputs_for_generation(
        self,
        input_ids,
        past_key_values=None,
        attention_mask=None,
        inputs_embeds=None,
        cache_position=None,
        position_ids=None,
        use_cache=True,
        pixel_values=None,
        pixel_values_videos=None,
        image_grid_thw=None,
        video_grid_thw=None,
        second_per_grid_ts=None,
        masks=None,
        **kwargs,
    ):
        # Overwritten -- in specific circumstances we don't want to forward image inputs to the model
        # import ipdb;ipdb.set_trace()
        model_inputs = super().prepare_inputs_for_generation(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            inputs_embeds=inputs_embeds,
            cache_position=cache_position,
            position_ids=position_ids,
            pixel_values=pixel_values,
            pixel_values_videos=pixel_values_videos,
            image_grid_thw=image_grid_thw,
            video_grid_thw=video_grid_thw,
            second_per_grid_ts=second_per_grid_ts,
            use_cache=use_cache,
            masks=masks,
            **kwargs,
        )

        # Qwen2-5-VL position_ids are prepareed with rope_deltas in forward
        model_inputs["position_ids"] = None

        if cache_position[0] != 0:
            model_inputs["pixel_values"] = None
            model_inputs["pixel_values_videos"] = None

        return model_inputs

    @torch.no_grad()
    def evaluate(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict_in_generate: Optional[bool] = None,
        pixel_values: Optional[torch.Tensor] = None,
        pixel_values_videos: Optional[torch.FloatTensor] = None,
        image_grid_thw: Optional[torch.LongTensor] = None,
        video_grid_thw: Optional[torch.LongTensor] = None,
        rope_deltas: Optional[torch.LongTensor] = None,
        cache_position: Optional[torch.LongTensor] = None,
        second_per_grid_ts: Optional[torch.Tensor] = None,
        max_new_tokens: int = 250,
        ### from lisa_plus
        offset: Optional[torch.LongTensor] = None,
        masks_list: List[torch.FloatTensor] = None,
        label_list: List[torch.Tensor] = None,
        output_masks_flag: bool = False,
        output_logits_flag: bool = False,
       **kwargs
    
    ):  
        original_mask_shape = kwargs.get('original_mask_shape')
        seg_token_idx = kwargs.get('seg_token_idx')
        sam_images = kwargs.get('sam_images')
        sam_resizes = kwargs.get('sam_resizes')
        # for inference with general next-token prediction tasks with word ids as output
        if not output_masks_flag:
            output = self.generate(
                input_ids=input_ids,
                attention_mask=attention_mask,
                pixel_values=pixel_values,
                image_grid_thw=image_grid_thw,
                max_new_tokens=max_new_tokens,
                return_dict_in_generate=return_dict_in_generate,
                inference=True,
            )
            return output
        # import ipdb;ipdb.set_trace()
        # for inference with SEG task and sam image input
        if seg_token_idx is not None and sam_images is not None and output_masks_flag:
            outputs = self.generate(
                    input_ids=input_ids, #torch.Size([1, 388])
                    do_sample=True,
                    output_hidden_states=output_hidden_states, #True
                    return_dict_in_generate=True, 
                    output_scores=True,
                    output_logits=True,
                    temperature=0.2,
                    max_new_tokens=512,
                    use_cache=True,
                    num_beams=1,
                    pixel_values=[pixel_values], #torch.Size([1380, 1176])
                    attention_mask=attention_mask, #torch.Size([1, 388])
                    image_grid_thw=[image_grid_thw.squeeze(0)] #torch.Size([1, 3])
                )
            # outputs.hidden_states: (num_ar_steps * (llm_layer_num * [bs, num_current_step_tokens, dim]))
            # outputs.sequences.shape [bs, prefill_token_num + new_token_num + 1]
            
            
            # last_hidden_state: [bs, prefill_token_num + new_token_num - 1, dim]
            last_hidden_state = torch.concatenate([hidden_states[-1] for hidden_states in outputs.hidden_states], dim=1) 
            # [bs, prefill_token_num + new_token_num, dim]
            last_hidden_state = torch.concatenate([last_hidden_state.new_zeros(1, 1, last_hidden_state.shape[-1]), last_hidden_state], dim=1)
            # [bs, prefill_token_num + new_token_num, dim]
            output_ids = outputs.sequences
           
            image_embeddings = self.get_visual_embs(sam_images.unsqueeze(0).to(torch.bfloat16).to(output_ids.device))
            batch_size = image_embeddings.shape[0]
            # import ipdb;ipdb.set_trace()
            seg_token_mask = output_ids == seg_token_idx
            pred_embeddings = last_hidden_state[seg_token_mask] # 相当于取出所有的实例，比如第1个样本是1个，第2个样本是2个
            #####
            # self.tokenizer.batch_decode(torch.argmax(self.lm_head(last_hidden_state),dim=2)[seg_token_mask])
            #####
            # get the logit of the predicted word
            if output_logits_flag:
                # [bs, vocab_size]
                logits_tuple = outputs.logits
                prompt_length = input_ids.shape[-1]
                output_confidences = [0.] * prompt_length # 输入prompt的score预设为0
                # logits_tuple 的第一个元素对应生成的第一步，对应生成 token在 generated_seq 中的位置 prompt_length + 0
                for t, logits in enumerate(logits_tuple):
                    # TODO softmax make the value to be near 1
                    prob = F.softmax(logits, dim=-1)
                    # 取得当前生成的 token id（推理时 batch_size=1）
                    token_id = output_ids[0, prompt_length + t]
                    token_prob = prob[0, token_id].item()
                    output_confidences.append(token_prob)
                assert len(output_confidences) == output_ids.shape[-1]

            if len(pred_embeddings) == 0:
                if output_logits_flag:
                    return output_ids, [], [output_confidences]
                else:
                    return output_ids, []

            seg_token_counts = seg_token_mask.int().sum(-1)  # [bs, ]

            seg_token_offset = seg_token_counts.cumsum(-1)
            seg_token_offset = torch.cat(
                [torch.zeros(1).long().to(pred_embeddings.device), seg_token_offset], dim=0
            )

            # seg_token_offset = seg_token_offset[offset]

            pred_embeddings_ = []
            for i in range(len(seg_token_offset) - 1):
                start_i, end_i = seg_token_offset[i], seg_token_offset[i + 1]
                pred_embeddings_.append(pred_embeddings[start_i:end_i])
            pred_embeddings = pred_embeddings_

            multimask_output = False
            pred_masks = []
            self.text_hidden_fcs[0].to(pred_embeddings[i].dtype)
            # 一个batch中一共出现的tokens
            for i in range(len(pred_embeddings)):
                SAM_adaptor_states = self.text_hidden_fcs[0](pred_embeddings[i])
                (
                        sparse_embeddings, #[3, 1, 256]
                        dense_embeddings, #[3, 256, 64, 64]
                ) = self.visual_model.prompt_encoder(
                    points=None,
                    boxes=None,
                    masks=None,
                    text_embeds=SAM_adaptor_states.unsqueeze(1),
                )
                sparse_embeddings = sparse_embeddings.to(SAM_adaptor_states.dtype)
                low_res_masks, iou_predictions = self.visual_model.mask_decoder(
                    image_embeddings=image_embeddings, #[2, 256, 64, 64]
                    image_pe=self.visual_model.prompt_encoder.get_dense_pe(),
                    sparse_prompt_embeddings=sparse_embeddings, #[3, 256, 64, 64]
                    dense_prompt_embeddings=dense_embeddings, #[3, 1, 256]
                    multimask_output=False,
                ) #[3, 1, 256, 256]
                pred_mask = self.visual_model.postprocess_masks(
                    low_res_masks, #[3, 1, 256, 256]
                    input_size=sam_resizes, #[686, 1024]
                    original_size=original_mask_shape #[500, 358]
                ) # [3, 1, 500, 358]
                pred_masks.append(pred_mask)
            # assume batchsize = 1
            # accumulate prob with iou prediction in SAM
            output_confidences_np = torch.tensor(output_confidences).to(seg_token_mask.device)
            output_confidences_np[seg_token_mask[0]] =  output_confidences_np[seg_token_mask[0]] * iou_predictions.squeeze()
            output_confidences = output_confidences_np.cpu().tolist()
            if output_logits_flag:
                return output_ids, pred_masks, [output_confidences]
            else:
                return output_ids, pred_masks
AutoConfig.register("osprey", OspreyConfig)
AutoModelForCausalLM.register(OspreyConfig, OspreyQwenForCausalLM)
