import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
import numpy as np
from typing import Optional
from torch.autograd import Function
from torch.autograd.function import once_differentiable
from torch.nn.init import constant_, xavier_uniform_
import math
from torch.nn import CrossEntropyLoss
# helpers
def _is_power_of_2(n):
    if (not isinstance(n, int)) or (n < 0):
        raise ValueError("invalid input for _is_power_of_2: {} (type: {})".format(n, type(n)))
    return (n & (n - 1) == 0) and n != 0
def multi_scale_deformable_attn_pytorch(
    value: torch.Tensor,
    value_spatial_shapes: torch.Tensor,
    sampling_locations: torch.Tensor,
    attention_weights: torch.Tensor,
) -> torch.Tensor:

    bs, _, num_heads, embed_dims = value.shape
    _, num_queries, num_heads, num_levels, num_points, _ = sampling_locations.shape
    value_list = value.split([H_ * W_ for H_, W_ in value_spatial_shapes], dim=1)
    sampling_grids = 2 * sampling_locations - 1
    sampling_value_list = []
    for level, (H_, W_) in enumerate(value_spatial_shapes):
        # bs, H_*W_, num_heads, embed_dims ->
        # bs, H_*W_, num_heads*embed_dims ->
        # bs, num_heads*embed_dims, H_*W_ ->
        # bs*num_heads, embed_dims, H_, W_
        value_l_ = (
            value_list[level].flatten(2).transpose(1, 2).reshape(bs * num_heads, embed_dims, H_, W_)
        )
        # bs, num_queries, num_heads, num_points, 2 ->
        # bs, num_heads, num_queries, num_points, 2 ->
        # bs*num_heads, num_queries, num_points, 2
        sampling_grid_l_ = sampling_grids[:, :, :, level].transpose(1, 2).flatten(0, 1)
        # bs*num_heads, embed_dims, num_queries, num_points
        sampling_value_l_ = F.grid_sample(
            value_l_, sampling_grid_l_, mode="bilinear", padding_mode="zeros", align_corners=False
        )
        sampling_value_list.append(sampling_value_l_)
    # (bs, num_queries, num_heads, num_levels, num_points) ->
    # (bs, num_heads, num_queries, num_levels, num_points) ->
    # (bs, num_heads, 1, num_queries, num_levels*num_points)
    attention_weights = attention_weights.transpose(1, 2).reshape(
        bs * num_heads, 1, num_queries, num_levels * num_points
    )
    output = (
        (torch.stack(sampling_value_list, dim=-2).flatten(-2) * attention_weights)
        .sum(-1)
        .view(bs, num_heads * embed_dims, num_queries)
    )
    return output.transpose(1, 2).contiguous()


class MultiScaleDeformableAttention(nn.Module):
    """Multi-Scale Deformable Attention Module used in Deformable-DETR

    `Deformable DETR: Deformable Transformers for End-to-End Object Detection.
    <https://arxiv.org/pdf/2010.04159.pdf>`_.

    Args:
        embed_dim (int): The embedding dimension of Attention. Default: 256.
        num_heads (int): The number of attention heads. Default: 8.
        num_levels (int): The number of feature map used in Attention. Default: 4.
        num_points (int): The number of sampling points for each query
            in each head. Default: 4.
        img2col_steps (int): The step used in image_to_column. Defualt: 64.
            dropout (float): Dropout layer used in output. Default: 0.1.
        batch_first (bool): if ``True``, then the input and output tensor will be
            provided as `(bs, n, embed_dim)`. Default: False. `(n, bs, embed_dim)`
    """

    def __init__(
        self,
        embed_dim: int = 256,
        num_heads: int = 8,
        num_levels: int = 4,
        num_points: int = 4,
        img2col_step: int = 64,
        batch_first: bool = False,
    ):
        super().__init__()
        if embed_dim % num_heads != 0:
            raise ValueError(
                "embed_dim must be divisible by num_heads, but got {} and {}".format(
                    embed_dim, num_heads
                )
            )
        head_dim = embed_dim // num_heads

        self.batch_first = batch_first

        self.im2col_step = img2col_step
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.num_levels = num_levels
        self.num_points = num_points
        self.sampling_offsets = nn.Linear(embed_dim, num_heads * num_levels * num_points * 2)
        self.attention_weights = nn.Linear(embed_dim, num_heads * num_levels * num_points)
        self.value_proj = nn.Linear(embed_dim, embed_dim)
        self.output_proj = nn.Linear(embed_dim, embed_dim)

        self.init_weights()

    def _reset_parameters(self):
        return self.init_weights()

    def init_weights(self):
        """
        Default initialization for Parameters of Module.
        """
        constant_(self.sampling_offsets.weight.data, 0.0)
        thetas = torch.arange(self.num_heads, dtype=torch.float32) * (
            2.0 * math.pi / self.num_heads
        )
        grid_init = torch.stack([thetas.cos(), thetas.sin()], -1)
        grid_init = (
            (grid_init / grid_init.abs().max(-1, keepdim=True)[0])
            .view(self.num_heads, 1, 1, 2)
            .repeat(1, self.num_levels, self.num_points, 1)
        )
        for i in range(self.num_points):
            grid_init[:, :, i, :] *= i + 1
        with torch.no_grad():
            self.sampling_offsets.bias = nn.Parameter(grid_init.view(-1))
        constant_(self.attention_weights.weight.data, 0.0)
        constant_(self.attention_weights.bias.data, 0.0)
        xavier_uniform_(self.value_proj.weight.data)
        constant_(self.value_proj.bias.data, 0.0)
        xavier_uniform_(self.output_proj.weight.data)
        constant_(self.output_proj.bias.data, 0.0)

    def freeze_sampling_offsets(self):
        self.sampling_offsets.weight.requires_grad = False
        self.sampling_offsets.bias.requires_grad = False

    def freeze_attention_weights(self):
        self.attention_weights.weight.requires_grad = False
        self.attention_weights.bias.requires_grad = False

    def forward(
        self,
        query: torch.Tensor,
        key: Optional[torch.Tensor] = None,
        value: Optional[torch.Tensor] = None,
        query_pos: Optional[torch.Tensor] = None,
        key_padding_mask: Optional[torch.Tensor] = None,
        reference_points: Optional[torch.Tensor] = None,
        spatial_shapes: Optional[torch.Tensor] = None,
        level_start_index: Optional[torch.Tensor] = None,
        **kwargs
    ) -> torch.Tensor:

        """Forward Function of MultiScaleDeformableAttention

        Args:
            query (torch.Tensor): Query embeddings with shape
                `(num_query, bs, embed_dim)`
            key (torch.Tensor): Key embeddings with shape
                `(num_key, bs, embed_dim)`
            value (torch.Tensor): Value embeddings with shape
                `(num_key, bs, embed_dim)`
            query_pos (torch.Tensor): The position embedding for `query`. Default: None.
            key_padding_mask (torch.Tensor): ByteTensor for `query`, with shape `(bs, num_key)`,
                indicating which elements within `key` to be ignored in attention.
            reference_points (torch.Tensor): The normalized reference points
                with shape `(bs, num_query, num_levels, 2)`,
                all elements is range in [0, 1], top-left (0, 0),
                bottom-right (1, 1), including padding are.
                or `(N, Length_{query}, num_levels, 4)`, add additional
                two dimensions `(h, w)` to form reference boxes.
            spatial_shapes (torch.Tensor): Spatial shape of features in different levels.
                With shape `(num_levels, 2)`, last dimension represents `(h, w)`.
            level_start_index (torch.Tensor): The start index of each level. A tensor with
                shape `(num_levels, )` which can be represented as
                `[0, h_0 * w_0, h_0 * w_0 + h_1 * w_1, ...]`.

        Returns:
            torch.Tensor: forward results with shape `(num_query, bs, embed_dim)`
        """
        # During inference, dtype may be inconsistent
        if self.sampling_offsets.weight.dtype!=query.dtype or self.sampling_offsets.bias.dtype!=query.dtype:
            self.sampling_offsets.to(dtype=query.dtype)
            self.attention_weights.to(dtype=query.dtype)
            self.value_proj.to(dtype=query.dtype)
            self.output_proj.to(dtype=query.dtype)
        if value is None:
            value = query

        if query_pos is not None:
            query = query + query_pos

        if not self.batch_first:
            # change to (bs, num_query ,embed_dims)
            query = query.permute(1, 0, 2)
            value = value.permute(1, 0, 2)

        bs, num_query, _ = query.shape
        bs, num_value, _ = value.shape

        assert (spatial_shapes[:, 0] * spatial_shapes[:, 1]).sum() == num_value

        value = self.value_proj(value)
        if key_padding_mask is not None:
            value = value.masked_fill(key_padding_mask[..., None], float(0))
        value = value.view(bs, num_value, self.num_heads, -1)
        sampling_offsets = self.sampling_offsets(query).view(
            bs, num_query, self.num_heads, self.num_levels, self.num_points, 2
        )
        attention_weights = self.attention_weights(query).view(
            bs, num_query, self.num_heads, self.num_levels * self.num_points
        )
        attention_weights = attention_weights.softmax(-1)
        attention_weights = attention_weights.view(
            bs,
            num_query,
            self.num_heads,
            self.num_levels,
            self.num_points,
        )
        reference_points=reference_points.to(device=sampling_offsets.device,dtype=sampling_offsets.dtype)
        # bs, num_query, num_heads, num_levels, num_points, 2
        if reference_points.shape[-1] == 2:
            offset_normalizer = torch.stack([spatial_shapes[..., 1], spatial_shapes[..., 0]], -1).to(device=sampling_offsets.device,dtype=sampling_offsets.dtype)
            sampling_locations = (
                reference_points[:, :, None, :, None, :]
                + sampling_offsets / offset_normalizer[None, None, None, :, None, :]
            )
        elif reference_points.shape[-1] == 4:
            sampling_locations = (
                reference_points[:, :, None, :, None, :2]
                + sampling_offsets
                / self.num_points
                * reference_points[:, :, None, :, None, 2:]
                * 0.5
            )
        else:
            raise ValueError(
                "Last dim of reference_points must be 2 or 4, but get {} instead.".format(
                    reference_points.shape[-1]
                )
            )

        output = multi_scale_deformable_attn_pytorch(
            value, spatial_shapes, sampling_locations, attention_weights
        )

        output = self.output_proj(output)

        if not self.batch_first:
            output = output.permute(1, 0, 2)

        return output

def visualize_normalized_tensor(tensor, title="Normalized Tensor Visualization", save_path=None, mean=None, std=None, reference_pts=None):
    """
    Visualize a normalized [3, H, W] PyTorch Tensor with optional denormalization.

    Args:
        tensor (torch.Tensor): Input tensor with shape [3, H, W], typically a normalized RGB image
        title (str): Image title
        save_path (str, optional): If provided, save the image to this path (e.g., 'output.png')
        mean (list/tuple, optional): Mean values used for original normalization, e.g., [0.485, 0.456, 0.406]
        std (list/tuple, optional): Std values used for original normalization, e.g., [0.229, 0.224, 0.225]
                                    If mean and std are provided, denormalization will be performed
    """
    import matplotlib.pyplot as plt
    tensor=tensor.float()
    # Clone tensor to avoid modifying original data
    tensor_vis = tensor.clone().detach()

    # Perform denormalization if mean and std are provided
    if mean is not None and std is not None:
        mean = torch.tensor(mean).view(-1, 1, 1).to(tensor.device)
        std = torch.tensor(std).view(-1, 1, 1).to(tensor.device)
        tensor_vis = tensor_vis * std + mean
    # Clamp to [0, 1] range to avoid outliers
    tensor_vis = torch.clamp(tensor_vis, 0, 1)
    # Convert to [H, W, 3] and numpy
    img_np = tensor_vis.permute(1, 2, 0).cpu().numpy()
    # Display using matplotlib
    plt.figure(figsize=(10, 10))
    plt.imshow(img_np)
    plt.title(title)
    plt.axis('off')

    # Draw reference points
    if reference_pts is not None:
        # Convert to numpy, ensure shape is [N, 2] with values [x, y] in range [0, 1]
        if isinstance(reference_pts, torch.Tensor):
            reference_pts = reference_pts.detach().cpu().numpy()
        elif isinstance(reference_pts, list):
            reference_pts = np.array(reference_pts)
        else:
            raise TypeError(f"reference_pts should be list or Tensor, but got {type(reference_pts)}")

        H, W, _ = img_np.shape

        # Iterate over each reference point [x, y] in [0, 1]
        for pt in reference_pts:
            x_norm, y_norm = pt[0], pt[1]

            # Map normalized coordinates [0,1] to pixel coordinates [0, W-1], [0, H-1]
            x_pixel = int(x_norm * (W - 1))
            y_pixel = int(y_norm * (H - 1))

            # Draw a red dot using matplotlib
            plt.scatter(x_pixel, y_pixel, color='red', s=50, edgecolors='white', linewidths=1, zorder=5)

    plt.tight_layout()
    # Save image (optional)
    if save_path:
        plt.savefig(save_path, bbox_inches='tight', pad_inches=0, dpi=100)
        print(f"Image saved to: {save_path}")

def resize_longest_side(image, target_size):
    """
    Resize image while maintaining aspect ratio, scaling the longest side to target_size.

    Args:
        image: PIL.Image object
        target_size: Target size for the longest side (integer)

    Returns:
        Resized PIL.Image object
    """
    # Get original image dimensions
    width, height = image.size

    # Calculate scale ratio
    max_dim = max(width, height)
    scale = target_size / max_dim

    # Calculate new dimensions (keep integer pixels)
    new_width = int(round(width * scale))
    new_height = int(round(height * scale))

    # Resize image using high-quality Lanczos filter
    resized_image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)

    return resized_image

def resize_mask_longest_side(mask_tensor, target_size):
    """
    Resize 0-1 mask tensor while maintaining aspect ratio using nearest neighbor interpolation (pure PyTorch).
    """
    # Handle dimensions: ensure input is (B, C, H, W) format (interpolate requires batch dimension)
    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))

    # Use interpolate for scaling with nearest neighbor mode
    resized_mask = F.interpolate(
        mask_tensor,
        size=(new_height, new_width),
        mode='nearest-exact'  # Exact nearest neighbor to avoid intermediate values
    )

    # Restore original dimensions (remove batch dimension)
    if resized_mask.shape[0] == 1:
        resized_mask = resized_mask.squeeze(0)

    return resized_mask

def count_parameters(model, module_name=""):
    total_params = 0
    trainable_params = 0

    for name, param in model.named_parameters():
        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("\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 normalize_tensor(tensor):
    # Calculate global min and max values
    min_val = tensor.min()
    max_val = tensor.max()

    # Handle case where max equals min (avoid division by zero)
    if max_val == min_val:
        return torch.ones_like(tensor) * 0.5

    # Linear normalization to [0, 1]
    normalized = (tensor - min_val) / (max_val - min_val)
    return normalized

def dice_loss(
    inputs: torch.Tensor,
    targets: torch.Tensor,
    num_masks=1,
    scale=1000,
    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=1,
):
    """
    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

class MLP(nn.Module):

    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int,
                 num_layers: int) -> None:
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(
            nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))

    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x

class MaskPooling(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, mask):
        if not x.shape[-2:] == mask.shape[-2:]:
            # reshape mask to x
            mask = F.interpolate(mask.to(x.dtype), size=x.shape[-2:], mode='bilinear', align_corners=False)

        B, C, H ,W = x.shape
        B, Q, H, W = mask.shape
        mask = (mask > 0).to(mask.dtype)
        denorm = mask.sum(dim=(-1, -2), keepdim=True) + 1e-8

        mask_pooled_x = torch.einsum(
            "bchw,bqhw->bqc",
            x,
            (mask / denorm).to(x.dtype),
        )
        return mask_pooled_x

def manual_multi_head_attention(Q, K, V, num_heads, mask=None):
    """
    Manual multi-head attention computation.
    Args:
        Q: Query tensor, shape [batch_size, seq_len_q, d_model]
        K: Key tensor, shape [batch_size, seq_len_kv, d_model]
        V: Value tensor, shape [batch_size, seq_len_kv, d_model]
        num_heads: Number of attention heads
        mask: Optional mask, shape [batch_size, seq_len_q, seq_len_kv] (True means masked)
    Returns:
        output: Attention output, shape [batch_size, seq_len_q, d_model]
        attn_weights: Attention weights, shape [batch_size, num_heads, seq_len_q, seq_len_kv]
    """
    batch_size, seq_len_q, d_model = Q.shape
    seq_len_kv = K.shape[1]
    head_dim = d_model // num_heads

    # 1. Linear transform and split heads [batch_size, seq_len, num_heads, head_dim]
    Q = Q.view(batch_size, seq_len_q, num_heads, head_dim).transpose(1, 2)  # [batch, num_heads, seq_len_q, head_dim]
    K = K.view(batch_size, seq_len_kv, num_heads, head_dim).transpose(1, 2)
    V = V.view(batch_size, seq_len_kv, num_heads, head_dim).transpose(1, 2)

    # 2. Calculate scaled dot-product attention scores
    scores = torch.matmul(Q, K.transpose(-2, -1)) / (head_dim ** 0.5)  # [batch, num_heads, seq_len_q, seq_len_kv]

    # 3. Apply mask (e.g., causal mask or padding mask)
    if mask is not None:
        scores = scores.masked_fill(mask.unsqueeze(1), -1e9)

    # 4. Stable Softmax computation
    attn_weights = F.softmax(scores, dim=-1)  # [batch, num_heads, seq_len_q, seq_len_kv]

    # 5. Weighted sum with attention
    output = torch.matmul(attn_weights, V)  # [batch, num_heads, seq_len_q, head_dim]

    # 6. Merge heads and linear transform (simplified, may need additional linear layer in practice)
    output = output.transpose(1, 2).contiguous()  # [batch, seq_len_q, num_heads, head_dim]
    output = output.view(batch_size, seq_len_q, d_model)  # [batch, seq_len_q, d_model]

    return output, attn_weights

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model, num_heads, dropout=0.1):
        super().__init__()
        assert d_model % num_heads == 0, "d_model must be divisible by num_heads"

        self.d_model = d_model
        self.num_heads = num_heads
        self.head_dim = d_model // num_heads

        # Linear transformation layers
        self.W_q = nn.Linear(d_model, d_model)
        self.W_k = nn.Linear(d_model, d_model)
        self.W_v = nn.Linear(d_model, d_model)
        self.W_o = nn.Linear(d_model, d_model)

        # Dropout and initialization
        self.dropout = nn.Dropout(dropout)
        self._init_weights()

    def _init_weights(self):
        nn.init.xavier_uniform_(self.W_q.weight)
        nn.init.xavier_uniform_(self.W_k.weight)
        nn.init.xavier_uniform_(self.W_v.weight)
        nn.init.xavier_uniform_(self.W_o.weight)

    def forward(self, Q, K, V, key_padding_mask=None, attn_mask=None):
        """
        Args:
            Q: Query, shape [batch_size, seq_len_q, d_model]
            K: Key, shape [batch_size, seq_len_kv, d_model]
            V: Value, shape [batch_size, seq_len_kv, d_model]
            key_padding_mask: Padding mask, shape [batch_size, seq_len_kv] (True means padding position)
            attn_mask: Attention mask, shape [seq_len_q, seq_len_kv] (e.g., causal mask)
        Returns:
            output: Attention output, shape [batch_size, seq_len_q, d_model]
            attn_weights: Attention weights, shape [batch_size, num_heads, seq_len_q, seq_len_kv]
        """
        batch_size, seq_len_q = Q.size(0), Q.size(1)
        seq_len_kv = K.size(1)

        # 1. Linear transform and split heads
        Q = self.W_q(Q).view(batch_size, seq_len_q, self.num_heads, self.head_dim).transpose(1, 2)
        K = self.W_k(K).view(batch_size, seq_len_kv, self.num_heads, self.head_dim).transpose(1, 2)
        V = self.W_v(V).view(batch_size, seq_len_kv, self.num_heads, self.head_dim).transpose(1, 2)
        # 2. Calculate scaled dot-product attention scores
        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.head_dim ** 0.5)

        # 3. Apply key_padding_mask (set padding positions to -inf)
        if key_padding_mask is not None:
            # Expand mask dimensions to match attention scores [batch_size, num_heads, seq_len_q, seq_len_kv]
            scores = scores.masked_fill(
                key_padding_mask.unsqueeze(1).unsqueeze(2),  # [batch, 1, 1, seq_len_kv]
                -1e6
            )
        # 4. Apply attn_mask (e.g., causal mask)
        if attn_mask is not None:
            scores = scores.masked_fill(attn_mask.unsqueeze(0).unsqueeze(1), float('-inf'))

        # 5. Compute Softmax and Dropout
        attn_weights = F.softmax(scores, dim=-1)
        attn_weights = self.dropout(attn_weights)

        # 6. Weighted sum
        output = torch.matmul(attn_weights, V)

        # 7. Merge heads and output transform
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, self.d_model)
        output = self.W_o(output)

        return output, attn_weights

class Fine_Grained_Visual_Prompt_Extractor(nn.Module):
    def __init__(self, mask_shape=336, embed_dim=2048, out_dim=2048):
        super(Fine_Grained_Visual_Prompt_Extractor, self).__init__()
        self.local_size=1024
        self.mask_shape = mask_shape
        self.mask_pooling = MaskPooling()
        self.mask_ca=MultiHeadAttention(out_dim,num_heads=8)
        self.mask_proj = nn.Linear(out_dim, out_dim)
        self.mask_proj_out = nn.Sequential(
                    nn.Linear(out_dim, out_dim),
                    nn.GELU(),
                    nn.Linear(out_dim, out_dim),
                    nn.GELU(),
                    nn.Linear(out_dim, out_dim)
                )
        self.mask_linear = MLP(mask_shape*mask_shape, embed_dim, out_dim, 3)
        self.coords_proj = MLP(4,embed_dim,out_dim,2)
        text_fc = [
            nn.Linear(256, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, 2048),
            nn.Dropout(0.0),
        ]
        self.re_encoder_projector = nn.ModuleList([nn.Sequential(*text_fc)])
        # Local 2 Global Deformable-CA
        self.SAM2Qwen_CA=MultiScaleDeformableAttention(embed_dim = out_dim,
                                        num_heads = 8,
                                        num_levels = 1,
                                        num_points = 8,
                                        img2col_step = 64,
                                        batch_first = True)
    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)
                )
                image_embeddings_list.append(image_embeddings)
            torch.cuda.empty_cache()
            image_embeddings = torch.cat(image_embeddings_list, 0)
        return image_embeddings

    def preprocess(self, x,size=None) -> torch.Tensor:
        """Normalize pixel values and pad to a square input."""
        if size is None:
            size=self.local_size
        h, w = x.shape[-2:]
        padh = size - h
        padw = size - w
        x = F.pad(x, (0, padw, 0, padh))
        return x

    def mask_to_box(self, mask):
        mask_bool = (mask > 0)
        coords = torch.nonzero(mask_bool, as_tuple=False)

        if coords.numel() == 0:
            B = mask.size(0)
            return torch.zeros((B, 4), dtype=torch.int32, device=mask.device)
        batch_indices = coords[:, 0]
        y_coords = coords[:, 1]
        x_coords = coords[:, 2]

        B = mask.size(0)
        boxes = torch.zeros((B, 4), dtype=torch.int32, device=mask.device)

        for b in range(B):
            mask_b = (batch_indices == b)
            if not mask_b.any():
                boxes[b] = torch.tensor([0, 0, 0, 0], dtype=torch.int32, device=mask.device)
                continue

            y_b = y_coords[mask_b]
            x_b = x_coords[mask_b]

            y0, y1 = y_b.min(), y_b.max()
            x0, x1 = x_b.min(), x_b.max()

            w = x1 - x0 + 1
            h = y1 - y0 + 1

            boxes[b] = torch.tensor([x0, y0, w, h], dtype=torch.int32, device=mask.device)

        return boxes

    def focal_crop(self,images,mask):
        min_box_w = 48
        min_box_h = 48
        cropped_imgs=[]
        cropped_masks=[]
        img_h, img_w = mask.shape[-2:]

        for idx,each in enumerate(self.mask_to_box(mask)):
            x0, y0, w, h = each
            xc, yc = x0 + w/2, y0 + h/2
            w, h = max(w, min_box_w), max(h, min_box_h)
            x0, y0 = int(xc - w / 2), int(yc - h / 2)
            cropped_masks.append(mask[idx][max(y0-h, 0):min(y0+2*h, img_h), max(x0-w, 0):min(x0+2*w, img_w)])
            cropped_imgs.append(images[:,max(y0-h, 0):min(y0+2*h, img_h), max(x0-w, 0):min(x0+2*w, img_w)])
        return cropped_imgs,cropped_masks

    def generate_2d_sine_positional_encoding(self,H, W, dim, device=None):
        pos_enc = torch.zeros((H, W, dim), device=device)
        d_model = dim
        div_term = torch.exp(torch.arange(0, d_model, 2, device=device) * (-math.log(10000.0) / d_model))

        y_coords, x_coords = torch.meshgrid(
            torch.arange(H, device=device), torch.arange(W, device=device), indexing='ij'
        )

        y_pos = y_coords.unsqueeze(-1)
        x_pos = x_coords.unsqueeze(-1)

        pos_enc[..., 0::2] = torch.sin(x_pos * div_term) + torch.sin(y_pos * div_term)
        pos_enc[..., 1::2] = torch.cos(x_pos * div_term) + torch.cos(y_pos * div_term)

        return pos_enc

    def getpoint_reference_from_mask(self,mask: torch.Tensor):
        H, W = mask.shape
        true_points = torch.nonzero(mask, as_tuple=False)
        if len(mask[mask==True])==0:
            reference_point = torch.tensor([0, 0], dtype=torch.float32)
            return reference_point
        idx = torch.randint(0, len(true_points), (1,)).item()
        y, x = true_points[idx]

        x_norm = x.float() / (W - 1) if W > 1 else torch.tensor(0.5)
        y_norm = y.float() / (H - 1) if H > 1 else torch.tensor(0.5)

        reference_point = torch.tensor([x_norm, y_norm], dtype=torch.float32)
        return reference_point

    def get_aux_lm_loss(self,mask_token,lm_head):
        self.ce_loss=CrossEntropyLoss()
        N = mask_token.shape[1]
        pred=lm_head(mask_token).flatten(0,1)
        label = [151669] * N
        label = torch.tensor(label).to(device=mask_token.device)
        lm_loss=self.ce_loss(pred, label)
        return lm_loss

    def get_self_supervised_seg_loss(self,mask_token,sam_image_embeddings,masks,visual_model,text_hidden_fcs,img_longest_side_size,lm_head):
        (
            sparse_embeddings,
            dense_embeddings,
        ) = visual_model.prompt_encoder(
            points=None,
            boxes=None,
            masks=None,
            text_embeds=text_hidden_fcs[0](mask_token.permute(1,0,2)),
        )
        sparse_embeddings = sparse_embeddings.to(mask_token.dtype)
        sam_image_embeddings = sam_image_embeddings.to(mask_token.dtype)
        low_res_masks, iou_predictions = visual_model.mask_decoder(
            image_embeddings=sam_image_embeddings,
            image_pe=visual_model.prompt_encoder.get_dense_pe(),
            sparse_prompt_embeddings=sparse_embeddings,
            dense_prompt_embeddings=dense_embeddings,
            multimask_output=False,
        )
        pred_mask = visual_model.postprocess_masks(
            low_res_masks,
            input_size=img_longest_side_size,
            original_size=masks.shape[1:]
        ).to(mask_token.dtype)
        mask_bce_loss =  (
                sigmoid_ce_loss(pred_mask.squeeze(1), masks.to(mask_token.dtype), num_masks=masks.shape[0])
                * masks.shape[0]
            )
        mask_dice_loss = (
                dice_loss(pred_mask.squeeze(1), masks, num_masks=masks.shape[0])
                * masks.shape[0]
            )

        return mask_bce_loss , mask_dice_loss

    def forward(self, masks,image_embeds,image_grid_thw,mask_token=None,images=None,visual_model=None,text_hidden_fcs=None,lm_head=None):
        pos_feats = []
        if masks is None:
            return pos_feats, pos_feats
        num_imgs = len(masks)
        # SAM encoder Focal Crop image
        cropped_imgs,cropped_masks = self.focal_crop(images,masks)
        for idx,each in enumerate(cropped_imgs):
            cropped_imgs[idx]=self.preprocess(resize_mask_longest_side(cropped_imgs[idx],self.local_size))
        cropped_imgs = torch.stack(cropped_imgs)
        for idx,each in enumerate(cropped_masks):
            cropped_masks[idx]=self.preprocess(resize_mask_longest_side(cropped_masks[idx],64),64)
        cropped_masks = torch.stack(cropped_masks)
        crop_feature=visual_model.image_encoder(cropped_imgs)

        crop_feature=self.re_encoder_projector[0](self.mask_pooling(crop_feature,cropped_masks))

        # Deformable Attention + SAM encoder full image CA
        pe=self.generate_2d_sine_positional_encoding(images.shape[1],images.shape[2],256,image_embeds.device).permute(2,0,1)
        pe=self.preprocess(resize_mask_longest_side(pe,64),size=64).unsqueeze(0)
        img_longest_side_size=resize_mask_longest_side(images,self.local_size).shape[1:]
        images=self.preprocess(resize_mask_longest_side(images,self.local_size))
        mask_local=self.preprocess(resize_mask_longest_side(masks,64),64)
        sam_feature_org = visual_model.image_encoder(images.unsqueeze(0))
        sam_feature = sam_feature_org + pe
        mask_feat=[]
        for each_mask in mask_local.bool():
            mask_feat.append(sam_feature.squeeze(0)[:,each_mask].permute(1,0))
        max_len = max(sublist.shape[0] for sublist in mask_feat)
        attention_mask = torch.zeros((len(mask_local), max_len), dtype=torch.long)
        for i,each in enumerate(mask_feat):
            attention_mask[i,mask_feat[i].shape[0]:]=1
            padding_len=max_len-mask_feat[i].shape[0]
            mask_feat[i]=torch.cat((mask_feat[i],torch.zeros(padding_len,256).to(mask_feat[i].device)),dim=0)
        mask_feat=torch.stack(mask_feat).to(mask_token.dtype)
        mask_feat=self.re_encoder_projector[0](mask_feat)

        query=crop_feature
        key=value=mask_feat
        key_padding_mask=attention_mask.to(torch.bool).to(mask_feat.device)

        mask_token=self.mask_ca(query,key,value,key_padding_mask=key_padding_mask)[0]
        mask_token=mask_token.permute(1,0,2)
        mask_token=self.mask_proj(mask_token)

        reference_points=[]
        for idx in range(masks.shape[0]):
            reference_points.append(self.getpoint_reference_from_mask(masks[idx]))
        reference_points=torch.stack(reference_points).unsqueeze(0).unsqueeze(2)
        spatial_shapes=torch.tensor([image_grid_thw[0][1]//2,image_grid_thw[0][2]//2]).unsqueeze(0)
        mask_token=self.SAM2Qwen_CA(mask_token,image_embeds.unsqueeze(0),image_embeds.unsqueeze(0),reference_points=reference_points,spatial_shapes=spatial_shapes,level_start_index=[0])
        mask_token=self.mask_proj_out(mask_token)

        masks = F.interpolate(masks.unsqueeze(0).float(), size=self.mask_shape, mode='bilinear', align_corners=False).squeeze(0)
        # Coords embed
        xywh=self.mask_to_box(masks)/self.mask_shape
        coords_embed=self.coords_proj(xywh.to(dtype=mask_feat.dtype, device=mask_feat.device))
        # Spatial token
        spatial_feat = self.mask_linear(masks.reshape(masks.shape[0], -1).to(dtype=mask_feat.dtype))
        pos_feats = spatial_feat+coords_embed

        # Self-supervised segmentation loss
        self_supervised_seg_loss=self.get_self_supervised_seg_loss(pos_feats.unsqueeze(0),sam_feature_org,masks,visual_model,text_hidden_fcs,img_longest_side_size,lm_head)

        return  mask_token.squeeze(0),pos_feats,self_supervised_seg_loss
