"""SAM 3's text-prompted segmentation, split into the three graphs the browser runs.

The split follows `Sam3Model.forward` (Hugging Face Transformers, `models/sam3/modeling_sam3.py`)
for a text prompt with no box prompts:

- **vision**: the Perception Encoder ViT and its neck, run once per image. Its outputs are the
  first three feature levels (288², 144² and 72², 256 channels); `Sam3Model` drops the fourth.
- **text**: the CLIP text encoder and `text_projection`, run once per phrase.
- **detector**: the DETR encoder and decoder, the box head, the dot-product scoring and the mask
  decoder, run once per phrase over the vision graph's features.

The 72² level's sine position encoding depends only on the input size, which is fixed at
1008 × 1008, so the detector carries it as a constant instead of taking it as an input.
"""

import torch
from torch import nn
from transformers.models.sam3.modeling_sam3 import box_cxcywh_to_xyxy, inverse_sigmoid

IMAGE_SIZE = 1008
TEXT_TOKENS = 32


class Sam3Vision(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.vision_encoder = model.vision_encoder

    def forward(self, pixel_values: torch.Tensor):
        levels = self.vision_encoder(pixel_values).fpn_hidden_states
        return levels[0], levels[1], levels[2]


class Sam3Text(nn.Module):
    def __init__(self, model):
        super().__init__()
        self.text_encoder = model.text_encoder
        self.text_projection = model.text_projection

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor):
        hidden = self.text_encoder(
            input_ids=input_ids, attention_mask=attention_mask, return_dict=True
        ).last_hidden_state
        return self.text_projection(hidden)


class Sam3Detector(nn.Module):
    def __init__(self, model, position_encoding: torch.Tensor):
        super().__init__()
        self.detr_encoder = model.detr_encoder
        self.detr_decoder = model.detr_decoder
        self.dot_product_scoring = model.dot_product_scoring
        self.mask_decoder = model.mask_decoder
        self.register_buffer("position_encoding", position_encoding)

    def forward(
        self,
        features_288: torch.Tensor,
        features_144: torch.Tensor,
        features_72: torch.Tensor,
        text_features: torch.Tensor,
        text_mask: torch.Tensor,
    ):
        mask = text_mask.bool()
        encoder = self.detr_encoder(
            vision_features=[features_72],
            text_features=text_features,
            vision_pos_embeds=[self.position_encoding],
            text_mask=mask,
        )
        decoder = self.detr_decoder(
            vision_features=encoder.last_hidden_state,
            text_features=encoder.text_features,
            vision_pos_encoding=encoder.pos_embeds_flattened,
            text_mask=mask,
            spatial_shapes=encoder.spatial_shapes,
        )
        hidden = decoder.intermediate_hidden_states[-1:]
        offsets = self.detr_decoder.box_head(hidden)
        boxes = (inverse_sigmoid(decoder.reference_boxes[-1:]) + offsets).sigmoid()
        logits = self.dot_product_scoring(
            decoder_hidden_states=hidden,
            text_features=encoder.text_features,
            text_mask=mask,
        ).squeeze(-1)
        masks = self.mask_decoder(
            decoder_queries=hidden[0],
            backbone_features=[features_288, features_144, features_72],
            encoder_hidden_states=encoder.last_hidden_state,
            prompt_features=text_features,
            prompt_mask=mask,
        )
        return (
            logits[0],
            box_cxcywh_to_xyxy(boxes)[0],
            decoder.presence_logits[-1],
            masks.pred_masks,
            masks.semantic_seg,
        )


def position_encoding_72(model) -> torch.Tensor:
    """The 72² level's position encoding for a 1008² image, as the vision model computes it."""
    with torch.no_grad():
        dummy = torch.zeros(1, 3, IMAGE_SIZE, IMAGE_SIZE)
        return model.vision_encoder(dummy).fpn_position_encoding[2].clone()