"""OpenWorldSAM: self-contained HuggingFace model (no detectron2).

Original paper: "Extending SAM2 for Universal Image Segmentation with Language Prompts"
(Xiao et al., NeurIPS 2025 Spotlight). Original code: GinnyXiao/OpenWorldSAM (Apache-2.0).

Architecture:
  evf_sam2        — EvfSam2Model (SAM2 Hiera-Large + BEiT-3 multimodal encoder)
  text_hidden_fcs — 3-layer projection MLP: BEiT-3 hidden_dim → query_dim
  positional_tokens — learnable positional embeddings [num_tokens, query_dim]
  cross_attention_transformer — 3-layer cross-attention stack
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision

from transformers import PreTrainedModel, AutoTokenizer

# Absolute imports: the repo root must be on sys.path (handled by the FiftyOne
# loader or by HF trust_remote_code loading from a local snapshot directory).
from configuration_openworld_sam import OpenWorldSAMConfig  # noqa: E402

# Trigger Hydra config registration before any SAM2 imports
from model import evf_sam2 as _evf_module  # noqa: F401, E402
from model.evf_sam2 import EvfSam2Model  # noqa: E402


# ---------------------------------------------------------------------------
# Cross-attention transformer (inlined from model/open_world_sam2.py)
# ---------------------------------------------------------------------------

class _CrossAttentionLayer(nn.Module):
    def __init__(self, embedding_dim, num_heads, mlp_dim, dropout=0.1):
        super().__init__()
        self.self_attn_norm = nn.LayerNorm(embedding_dim)
        self.self_attn = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True)
        self.self_attn_dropout = nn.Dropout(dropout)

        self.cross_attn_norm = nn.LayerNorm(embedding_dim)
        self.cross_attn = nn.MultiheadAttention(embedding_dim, num_heads, dropout=dropout, batch_first=True)
        self.cross_attn_dropout = nn.Dropout(dropout)

        self.mlp_norm = nn.LayerNorm(embedding_dim)
        self.mlp = nn.Sequential(
            nn.Linear(embedding_dim, mlp_dim),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(mlp_dim, embedding_dim),
            nn.Dropout(dropout),
        )

    def forward(self, vlm_features, image_embeddings):
        # Self-attention
        r = vlm_features
        x = self.self_attn_norm(vlm_features)
        x, _ = self.self_attn(x, x, x)
        x = r + self.self_attn_dropout(x)

        # Cross-attention
        r = x
        x = self.cross_attn_norm(x)
        x, _ = self.cross_attn(query=x, key=image_embeddings, value=image_embeddings)
        x = r + self.cross_attn_dropout(x)

        # MLP
        r = x
        x = self.mlp_norm(x)
        x = self.mlp(x)
        return r + x


class _CrossAttentionTransformer(nn.Module):
    def __init__(self, embedding_dim, num_heads, mlp_dim, num_layers=3, dropout=0.1):
        super().__init__()
        self.layers = nn.ModuleList([
            _CrossAttentionLayer(embedding_dim, num_heads, mlp_dim, dropout)
            for _ in range(num_layers)
        ])

    def forward(self, vlm_features, image_embeddings):
        x = vlm_features
        for layer in self.layers:
            x = layer(x, image_embeddings)
        return x


# ---------------------------------------------------------------------------
# Pure-torch helpers (replacing detectron2 structures)
# ---------------------------------------------------------------------------

def _masks_to_boxes(masks):
    """bool tensor [N, H, W] → float tensor [N, 4] xyxy bounding boxes."""
    n = masks.shape[0]
    boxes = torch.zeros((n, 4), dtype=torch.float32, device=masks.device)
    for i in range(n):
        m = masks[i]
        rows = m.any(dim=1).nonzero(as_tuple=False)
        cols = m.any(dim=0).nonzero(as_tuple=False)
        if rows.numel() > 0 and cols.numel() > 0:
            boxes[i] = torch.tensor(
                [cols[0].item(), rows[0].item(), cols[-1].item() + 1, rows[-1].item() + 1],
                dtype=torch.float32,
                device=masks.device,
            )
    return boxes


# ---------------------------------------------------------------------------
# Main PreTrainedModel
# ---------------------------------------------------------------------------

class OpenWorldSAMModel(PreTrainedModel):
    """OpenWorldSAM zero-shot segmentation model (HuggingFace trust_remote_code).

    Usage::

        from transformers import AutoModel
        model = AutoModel.from_pretrained(
            "neerajaabhyankar/openworld-sam", trust_remote_code=True
        )
        # batched_inputs: list of dicts with keys:
        #   "image"            — float32 SAM-normalised tensor [3, 1024, 1024]
        #   "evf_image"        — float32 BEiT-3 tensor [3, 224, 224]
        #   "height", "width"  — original image dimensions (int)
        #   "prompt"           — list[str] of text prompts
        #   "unique_categories" — list[int] of category ids (one per prompt)
        outputs = model(batched_inputs)
        # outputs: list of dicts, one per image; key "instances" holds masks/scores/class_ids
    """

    config_class = OpenWorldSAMConfig

    def __init__(self, config: OpenWorldSAMConfig):
        super().__init__(config)

        # EVF-SAM2 backbone (SAM2 Hiera-L + BEiT-3)
        from model.configuration_evf import EvfConfig
        evf_cfg = EvfConfig(
            hidden_size=1024,
            sam_scale=config.sam_scale,
            mm_extractor_scale=config.mm_extractor_scale,
        )
        self.evf_sam2 = EvfSam2Model(evf_cfg)

        # Projection MLP: BEiT-3 hidden (1024) → query_dim (256)
        in_dim = 1024  # BEiT-3 large hidden size
        qd = config.query_dim
        self.text_hidden_fcs = nn.ModuleList([
            nn.Sequential(nn.Linear(in_dim, in_dim), nn.ReLU(), nn.Linear(in_dim, qd))
        ])

        # Learnable positional tokens [num_tokens, query_dim]
        self.positional_tokens = nn.Parameter(torch.randn(config.num_tokens, qd))

        # Cross-attention transformer
        self.cross_attention_transformer = _CrossAttentionTransformer(
            embedding_dim=qd,
            num_heads=8,
            mlp_dim=qd * 4,
            num_layers=config.cross_attention_layers,
        )

        # SAM2 feature size schedule (matches Hiera-L)
        self._bb_feat_sizes = [(256, 256), (128, 128), (64, 64)]

        # Preprocessing buffers
        self.register_buffer(
            "pixel_mean",
            torch.tensor(config.pixel_mean).view(-1, 1, 1),
            persistent=False,
        )
        self.register_buffer(
            "pixel_std",
            torch.tensor(config.pixel_std).view(-1, 1, 1),
            persistent=False,
        )

        # Tokenizer loaded lazily on first forward
        self._tokenizer = None

        # Required by transformers>=5's from_pretrained (sets
        # self.all_tied_weights_keys and other bookkeeping consumed by
        # _finalize_model_loading); harmless no-op pre-checkpoint-load init.
        self.post_init()

    # ------------------------------------------------------------------
    # Tokenizer
    # ------------------------------------------------------------------

    @property
    def tokenizer(self):
        if self._tokenizer is None:
            self._tokenizer = AutoTokenizer.from_pretrained(
                self.config.tokenizer_name_or_path,
                padding_side="right",
                use_fast=False,
            )
        return self._tokenizer

    def _tokenize_prompts(self, prompts):
        tok = self.tokenizer
        ids = [tok(p, return_tensors="pt").input_ids[0] for p in prompts]
        ids = torch.nn.utils.rnn.pad_sequence(ids, batch_first=True, padding_value=tok.pad_token_id)
        masks = ids.ne(tok.pad_token_id)
        trunc = tok.model_max_length
        return ids[:, :trunc].to(self.device), masks[:, :trunc].to(self.device)

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------

    def forward(self, batched_inputs):
        """
        Args:
            batched_inputs: list of dicts, one per image:
                "image"             float32 [3, 1024, 1024], SAM normalised
                "evf_image"         float32 [3, 224, 224],  BEiT-3 normalised
                "height", "width"   int, original image size
                "prompt"            list[str]
                "unique_categories" list[int]

        Returns:
            list of dicts, one per image. Key "instances" is a dict:
                "masks"     bool tensor [N, H, W]
                "scores"    float tensor [N]
                "class_ids" long tensor [N]
        """
        dtype = torch.float32
        images = torch.stack([x["image"].to(dtype=dtype, device=self.device) for x in batched_inputs])
        images_evf = torch.stack([x["evf_image"].to(dtype=dtype, device=self.device) for x in batched_inputs])
        original_size_list = [(x["height"], x["width"]) for x in batched_inputs]

        # Build flattened prompt list with per-image offsets
        offset = [0]
        all_prompts = []
        for x in batched_inputs:
            all_prompts.extend(x["prompt"])
            offset.append(offset[-1] + len(x["prompt"]))

        input_ids, attention_masks = self._tokenize_prompts(all_prompts)
        batch_size = len(batched_inputs)

        # SAM2 visual encoder
        with torch.no_grad():
            backbone_out = self.evf_sam2.visual_model.forward_image(images)
            _, image_embeddings, _, _ = self.evf_sam2.visual_model._prepare_backbone_features(backbone_out)

        image_embeddings = [e.to(dtype) for e in image_embeddings]
        if self.evf_sam2.visual_model.directly_add_no_mem_embed:
            image_embeddings[-1] = image_embeddings[-1] + self.evf_sam2.visual_model.no_mem_embed

        # Expand images_evf per prompt count if using visual tokens
        if self.config.use_visual_tokens:
            imgs_list = []
            for i in range(batch_size):
                n = offset[i + 1] - offset[i]
                imgs_list.append(images_evf[i].unsqueeze(0).expand(n, -1, -1, -1).contiguous())
            images_evf_expanded = torch.cat(imgs_list, dim=0)
        else:
            images_evf_expanded = None

        # BEiT-3 multimodal encoding
        with torch.no_grad():
            if images_evf_expanded is not None:
                out = self.evf_sam2.mm_extractor.beit3(
                    visual_tokens=images_evf_expanded,
                    textual_tokens=input_ids,
                    text_padding_position=~attention_masks,
                )
            else:
                out = self.evf_sam2.mm_extractor.beit3(
                    visual_tokens=None,
                    textual_tokens=input_ids,
                    text_padding_position=~attention_masks,
                )
        feat = out["encoder_out"][:, :1, ...]  # [total_prompts, 1, hidden]
        feat = self.text_hidden_fcs[0](feat)   # [total_prompts, 1, query_dim]

        # Split back per image
        feat = torch.split(feat, [offset[i + 1] - offset[i] for i in range(batch_size)])

        # Multi-scale image feature tensor
        feats = [
            e.permute(1, 2, 0).view(batch_size, -1, *sz)
            for e, sz in zip(image_embeddings[::-1], self._bb_feat_sizes[::-1])
        ][::-1]
        _features = {"image_embed": feats[-1], "high_res_feats": feats[:-1]}

        processed_results = []
        for img_idx in range(batch_size):
            img_feat = feat[img_idx]  # [num_prompts, 1, query_dim]

            # Build batch_feat_with_tokens
            tokens_list = []
            for pf in img_feat:
                repeated = pf.expand(self.config.num_tokens, -1, -1)  # [num_tokens, 1, query_dim]
                tokens_list.append(repeated + self.positional_tokens.unsqueeze(1))
            batch_feat_with_tokens = torch.cat(tokens_list, dim=0)  # [total_tokens, 1, query_dim]

            # Cross-attention with skip connection
            if self.config.use_cross_attention:
                img_embed = _features["image_embed"][img_idx].flatten(1).transpose(0, 1).unsqueeze(0)
                # batch_feat_with_tokens: [total_tokens, 1, qd] → squeeze middle → [1, total_tokens, qd]
                bft = batch_feat_with_tokens.squeeze(1).unsqueeze(0)
                enhanced = self.cross_attention_transformer(bft, img_embed)  # [1, total_tokens, qd]
                # Reshape back to [total_tokens, 1, qd] and add skip connection
                enhanced = enhanced.squeeze(0).unsqueeze(1)
                batch_feat_with_tokens = batch_feat_with_tokens + enhanced

            # SAM2 prompt encoder + mask decoder
            sparse_embeddings, dense_embeddings = self.evf_sam2.visual_model.sam_prompt_encoder(
                points=None, boxes=None, masks=None, text_embeds=batch_feat_with_tokens,
            )
            sparse_embeddings = sparse_embeddings.to(batch_feat_with_tokens.dtype)

            high_res_features = [
                f[img_idx].unsqueeze(0) for f in _features["high_res_feats"]
            ]

            # The decoder's repeat_image path duplicates image_embeddings and
            # image_pe once per candidate query, so scoring all num_classes *
            # num_tokens candidates in one call can reach double-digit GiB
            # for large vocabularies. `chunk_size` bounds peak memory
            # rather than total candidate count.
            image_embed_img = _features["image_embed"][img_idx].unsqueeze(0)
            image_pe = self.evf_sam2.visual_model.sam_prompt_encoder.get_dense_pe()
            chunk_size = self.config.mask_decoder_chunk_size
            num_total_tokens = sparse_embeddings.shape[0]

            low_res_masks_chunks = []
            iou_pred_chunks = []

            with torch.no_grad():
                for start in range(0, num_total_tokens, chunk_size):
                    end = min(start + chunk_size, num_total_tokens)
                    chunk_low_res_masks, chunk_iou_pred, _, _ = self.evf_sam2.visual_model.sam_mask_decoder(
                        image_embeddings=image_embed_img,
                        image_pe=image_pe,
                        sparse_prompt_embeddings=sparse_embeddings[start:end],
                        dense_prompt_embeddings=dense_embeddings[start:end],
                        multimask_output=False,
                        repeat_image=True,
                        high_res_features=high_res_features,
                    )
                    low_res_masks_chunks.append(chunk_low_res_masks)
                    iou_pred_chunks.append(chunk_iou_pred)
            low_res_masks = torch.cat(low_res_masks_chunks, dim=0)
            iou_pred = torch.cat(iou_pred_chunks, dim=0)

            pred_masks = low_res_masks.squeeze(1)  # [total_tokens, H_low, W_low]
            pred_logits = iou_pred.squeeze(1)       # [total_tokens]

            # Assign class labels: each prompt gets num_tokens predictions
            unique_categories = batched_inputs[img_idx]["unique_categories"]
            num_total = pred_masks.shape[0]
            class_indices = torch.div(
                torch.arange(num_total, device=self.device),
                self.config.num_tokens, rounding_mode="floor"
            )
            class_labels = torch.tensor(
                [unique_categories[int(i)] for i in class_indices],
                dtype=torch.long, device=self.device,
            )

            # Filter on low-res masks first; only the survivors get upsampled
            # to the original image size (upsampling the full query set before
            # filtering allocates one [num_queries, H, W] float32 tensor that
            # can reach double-digit GiB for large vocabularies).
            instances = self._instance_inference(
                pred_masks, pred_logits, class_labels, original_size_list[img_idx]
            )

            processed_results.append({"instances": instances})

        return processed_results

    def _postprocess_masks(self, masks, orig_hw):
        return F.interpolate(
            masks.float().unsqueeze(0), orig_hw, mode="bilinear", align_corners=False
        ).squeeze(0)

    def _instance_inference(self, pred_masks, iou_scores, class_labels, orig_hw):
        """Returns dict with keys: masks (bool), scores (float), class_ids (long)."""
        pred_masks = pred_masks.squeeze(1) if pred_masks.ndim == 4 else pred_masks

        # Top-K filter
        if self.config.top_k_on:
            k = min(self.config.detections_per_image, pred_masks.shape[0])
            idx = torch.argsort(iou_scores, descending=True)[:k]
            pred_masks, iou_scores, class_labels = pred_masks[idx], iou_scores[idx], class_labels[idx]

        # IoU threshold filter
        keep = iou_scores >= self.config.iou_thresh
        pred_masks, iou_scores, class_labels = pred_masks[keep], iou_scores[keep], class_labels[keep]

        if pred_masks.shape[0] == 0:
            empty = torch.empty(0, device=self.device)
            return {
                "masks": torch.empty((0, *orig_hw), dtype=torch.bool, device=self.device),
                "scores": empty,
                "class_ids": empty.long(),
            }


        # NMS on low-res masks — box IoU is scale-invariant, so this doesn't
        # need the full-res masks either.
        low_res_binary_masks = pred_masks > 0
        if self.config.nms_on:
            boxes = _masks_to_boxes(low_res_binary_masks)
            nms_keep = torchvision.ops.nms(boxes, iou_scores, self.config.nms_thresh)
            pred_masks = pred_masks[nms_keep]
            iou_scores = iou_scores[nms_keep]
            class_labels = class_labels[nms_keep]

        # Upsample only the surviving masks to the original image size
        pred_masks = self._postprocess_masks(pred_masks, orig_hw)
        binary_masks = pred_masks > 0

        return {
            "masks": binary_masks,
            "scores": iou_scores,
            "class_ids": class_labels,
        }

    # ------------------------------------------------------------------
    # Convenience preprocessing (mirrors demo/inference_utils.py)
    # ------------------------------------------------------------------

    def preprocess_image(self, image):
        """Normalise a uint8 HWC numpy array (RGB) → float32 [3, 1024, 1024] tensor on model device."""
        if not isinstance(image, np.ndarray):
            image = np.array(image)
        tensor = torch.as_tensor(np.ascontiguousarray(image.transpose(2, 0, 1))).float()
        tensor = F.interpolate(
            tensor.unsqueeze(0), (1024, 1024), mode="bilinear", align_corners=False
        ).squeeze(0)
        return (tensor - self.pixel_mean.cpu()) / self.pixel_std.cpu()

    def preprocess_image_beit3(self, image):
        """Normalise a uint8 HWC numpy array (RGB) → float32 [3, 224, 224] tensor."""
        from torchvision import transforms
        from PIL import Image as PILImage

        if isinstance(image, np.ndarray):
            pil = PILImage.fromarray(image.astype(np.uint8))
        else:
            pil = image
        tf = transforms.Compose([
            transforms.ToTensor(),
            transforms.Resize((224, 224), interpolation=3, antialias=None),
            transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
        ])
        return tf(pil)