# coding=utf-8
# Copyright 2026 The ConCor-1 authors.
#
# 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.
"""Processor for ConCor-1.

Turns an image-text pair into the flat multimodal sequence ConCor-1 expects
(image tokens, text tokens, bridge tokens) and turns the model's per-bridge
logits back into a set of image-text correspondences: a character-level text
span set and a binary image mask per correspondence.
"""

from __future__ import annotations

import string
from typing import Dict, List, Optional, Sequence, Tuple, Union

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

from transformers.feature_extraction_utils import BatchFeature
from transformers.processing_utils import ProcessorMixin

ImageInput = Union["PIL.Image.Image", Sequence["PIL.Image.Image"]]  # noqa: F821

_TRIM_CHARACTERS = frozenset(string.punctuation)


def compute_span_iou(spans_a: List[List[int]], spans_b: List[List[int]]) -> float:
    """IoU of two character-span sets, treated as 1-D binary masks over the text."""
    characters_a = {index for start, end in spans_a for index in range(start, end)}
    characters_b = {index for start, end in spans_b for index in range(start, end)}
    union = len(characters_a | characters_b)
    if union == 0:
        return 0.0
    return len(characters_a & characters_b) / union


def compute_mask_iou(mask_a: np.ndarray, mask_b: np.ndarray) -> float:
    """IoU of two binary masks."""
    flat_a = mask_a.astype(bool).ravel()
    flat_b = mask_b.astype(bool).ravel()
    union = np.count_nonzero(flat_a | flat_b)
    if union == 0:
        return 0.0
    return np.count_nonzero(flat_a & flat_b) / union


class ConCor1Processor(ProcessorMixin):
    r"""Constructs a ConCor-1 processor from a Qwen3.5 image processor and tokenizer.

    Args:
        image_processor: the Qwen3.5 (`Qwen2VLImageProcessorFast`) image processor.
        tokenizer: the Qwen3.5 tokenizer.
        num_bridge_tokens (`int`, *optional*, defaults to 385):
            Number of bridge tokens appended to every sequence.
        bridge_token_id_start (`int`, *optional*, defaults to 248077):
            Token id of the first bridge token.
        patch_size (`int`, *optional*, defaults to 16): vision patch size.
        merge_size (`int`, *optional*, defaults to 2): patch-merger factor.
        min_pixels (`int`, *optional*, defaults to 1003520):
            `min_pixels` forwarded to the image processor.
        max_pixels (`int`, *optional*, defaults to 1003520):
            `max_pixels` forwarded to the image processor. The released
            checkpoint was trained with `min_pixels == max_pixels`, i.e. a fixed
            ~1.0 M pixel budget.
        presence_threshold (`float`, *optional*, defaults to 0.1)
        text_threshold (`float`, *optional*, defaults to 0.45)
        image_threshold (`float`, *optional*, defaults to 0.45)
        nms_iou_threshold (`float`, *optional*, defaults to 0.5)

    Example:

    ```python
    >>> text = "A close-up of a brown bear sitting in lush green grass."
    >>> inputs = processor(images=image, text=text, return_tensors="pt")
    >>> outputs = model(**inputs)
    >>> correspondences = processor.post_process_correspondences(
    ...     outputs, text=text, target_sizes=[(image.height, image.width)]
    ... )
    ```
    """

    attributes = ["image_processor", "tokenizer"]
    # Qwen3.5's image processor, named concretely: transformers 5.x deprecates
    # `"AutoImageProcessor"` here, and its newer releases reject the (slow, fast)
    # tuple form. Requires torchvision, as all fast image processors do.
    image_processor_class = "Qwen2VLImageProcessorFast"
    tokenizer_class = "AutoTokenizer"

    def __init__(
        self,
        image_processor=None,
        tokenizer=None,
        num_bridge_tokens: int = 385,
        bridge_token_id_start: int = 248077,
        patch_size: int = 16,
        merge_size: int = 2,
        num_mask_upsample_blocks: int = 2,
        min_pixels: int = 1003520,
        max_pixels: int = 1003520,
        presence_threshold: float = 0.1,
        text_threshold: float = 0.45,
        image_threshold: float = 0.45,
        nms_iou_threshold: float = 0.5,
        **kwargs,
    ):
        self.num_bridge_tokens = num_bridge_tokens
        self.bridge_token_id_start = bridge_token_id_start
        self.patch_size = patch_size
        self.merge_size = merge_size
        self.num_mask_upsample_blocks = num_mask_upsample_blocks
        self.min_pixels = min_pixels
        self.max_pixels = max_pixels
        self.presence_threshold = presence_threshold
        self.text_threshold = text_threshold
        self.image_threshold = image_threshold
        self.nms_iou_threshold = nms_iou_threshold
        super().__init__(image_processor, tokenizer, **kwargs)

    @property
    def bridge_token_ids(self) -> List[int]:
        start = self.bridge_token_id_start
        return list(range(start, start + self.num_bridge_tokens))

    # ── Pre-processing ───────────────────────────────────────────────────────

    def __call__(
        self,
        images: Optional[ImageInput] = None,
        text: Optional[Union[str, List[str]]] = None,
        return_tensors: str = "pt",
        **kwargs,
    ) -> BatchFeature:
        """Build ConCor-1's flat multimodal sequences for one or more image-text pairs.

        The text is *not* wrapped in a chat template: ConCor-1 consumes the raw
        text (a caption, a list of category names, a referring expression, ...)
        directly, and predicts its text masks over exactly these tokens.

        Args:
            images: one PIL image, or a list with one image per text.
            text: the paired text, or a list of texts.
            return_tensors: only `"pt"` is supported.

        Returns:
            [`BatchFeature`] with `input_ids`, `attention_mask`,
            `visual_token_mask`, `text_token_mask`, `bridge_token_mask` and —
            when images are given — `pixel_values` and `image_grid_thw`.
        """
        if return_tensors != "pt":
            raise ValueError(f"ConCor1Processor only supports return_tensors='pt', got {return_tensors!r}")
        if text is None:
            raise ValueError(
                "ConCor-1 always grounds text: pass the image's paired `text` "
                "(a caption, category list or referring expression)."
            )

        texts = [text] if isinstance(text, str) else list(text)
        if images is None:
            image_list = []
        elif isinstance(images, (list, tuple)):
            image_list = list(images)
        else:
            image_list = [images]

        if image_list and len(image_list) != len(texts):
            raise ValueError(
                f"Got {len(image_list)} image(s) and {len(texts)} text(s); pass one image per text."
            )

        pixel_values, image_grid_thw = None, None
        num_visual_tokens = [0] * len(texts)
        if image_list:
            image_inputs = self.image_processor(
                images=image_list,
                return_tensors="pt",
                min_pixels=kwargs.pop("min_pixels", self.min_pixels),
                max_pixels=kwargs.pop("max_pixels", self.max_pixels),
            )
            pixel_values = image_inputs["pixel_values"]
            image_grid_thw = image_inputs["image_grid_thw"]
            num_visual_tokens = [
                int(grid.prod().item()) // (self.merge_size ** 2) for grid in image_grid_thw
            ]

        vision_start_id = self.tokenizer.convert_tokens_to_ids("<|vision_start|>")
        vision_end_id = self.tokenizer.convert_tokens_to_ids("<|vision_end|>")
        image_pad_id = self.tokenizer.convert_tokens_to_ids("<|image_pad|>")
        bridge_ids = self.bridge_token_ids

        sequences, visual_masks, text_masks, bridge_masks = [], [], [], []
        for index, sample_text in enumerate(texts):
            ids: List[int] = []
            visual_mask: List[bool] = []
            text_mask: List[bool] = []

            if num_visual_tokens[index] > 0:
                ids.append(vision_start_id)
                ids.extend([image_pad_id] * num_visual_tokens[index])
                ids.append(vision_end_id)
                visual_mask.extend([False] + [True] * num_visual_tokens[index] + [False])
                text_mask.extend([False] * (num_visual_tokens[index] + 2))

            text_ids = self.tokenizer.encode(sample_text, add_special_tokens=False)
            if not text_ids:
                raise ValueError(f"Text {index} is empty after tokenization: {sample_text!r}")
            ids.extend(text_ids)
            visual_mask.extend([False] * len(text_ids))
            text_mask.extend([True] * len(text_ids))

            ids.extend(bridge_ids)
            visual_mask.extend([False] * len(bridge_ids))
            text_mask.extend([False] * len(bridge_ids))

            bridge_mask = [False] * (len(ids) - len(bridge_ids)) + [True] * len(bridge_ids)

            sequences.append(ids)
            visual_masks.append(visual_mask)
            text_masks.append(text_mask)
            bridge_masks.append(bridge_mask)

        # Right-pad to the longest sequence, as in training.
        max_length = max(len(ids) for ids in sequences)
        pad_token_id = self.tokenizer.pad_token_id or 0
        batch_size = len(sequences)

        input_ids = torch.full((batch_size, max_length), pad_token_id, dtype=torch.long)
        attention_mask = torch.zeros((batch_size, max_length), dtype=torch.long)
        visual_token_mask = torch.zeros((batch_size, max_length), dtype=torch.bool)
        text_token_mask = torch.zeros((batch_size, max_length), dtype=torch.bool)
        bridge_token_mask = torch.zeros((batch_size, max_length), dtype=torch.bool)

        for index, ids in enumerate(sequences):
            length = len(ids)
            input_ids[index, :length] = torch.tensor(ids, dtype=torch.long)
            attention_mask[index, :length] = 1
            visual_token_mask[index, :length] = torch.tensor(visual_masks[index], dtype=torch.bool)
            text_token_mask[index, :length] = torch.tensor(text_masks[index], dtype=torch.bool)
            bridge_token_mask[index, :length] = torch.tensor(bridge_masks[index], dtype=torch.bool)

        data = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "visual_token_mask": visual_token_mask,
            "text_token_mask": text_token_mask,
            "bridge_token_mask": bridge_token_mask,
        }
        if pixel_values is not None:
            data["pixel_values"] = pixel_values
            data["image_grid_thw"] = image_grid_thw
        return BatchFeature(data=data)

    # ── Post-processing ──────────────────────────────────────────────────────

    def text_logits_to_spans(
        self,
        text_logits: np.ndarray,
        offset_mapping: Sequence[Tuple[int, int]],
        text: str,
        text_threshold: float,
    ) -> List[List[int]]:
        """Convert one bridge token's text-mask logits into character spans.

        Active text tokens are mapped to their character ranges, merged into
        maximal contiguous spans and trimmed of leading/trailing whitespace and
        punctuation. A text mask may consist of several disjoint spans, which is
        how discontinuous and co-referring mentions are represented.
        """
        probabilities = 1.0 / (1.0 + np.exp(-text_logits.astype(np.float64)))
        active = probabilities >= text_threshold

        characters = set()
        for token_index in np.where(active)[0]:
            if token_index < len(offset_mapping):
                start, end = offset_mapping[token_index]
                characters.update(range(int(start), int(end)))
        if not characters:
            return []

        ordered = sorted(characters)
        spans: List[List[int]] = []
        start = ordered[0]
        end = start + 1
        for character in ordered[1:]:
            if character == end:
                end += 1
            else:
                spans.append([start, end])
                start, end = character, character + 1
        spans.append([start, end])

        trimmed: List[List[int]] = []
        for span_start, span_end in spans:
            span_start = max(0, span_start)
            span_end = min(len(text), span_end)
            while span_start < span_end and (
                text[span_start].isspace() or text[span_start] in _TRIM_CHARACTERS
            ):
                span_start += 1
            while span_end > span_start and (
                text[span_end - 1].isspace() or text[span_end - 1] in _TRIM_CHARACTERS
            ):
                span_end -= 1
            if span_start < span_end:
                trimmed.append([span_start, span_end])
        return trimmed

    @staticmethod
    def suppress_duplicate_correspondences(
        masks: List[Optional[np.ndarray]],
        spans: List[List[List[int]]],
        presence_scores: np.ndarray,
        nms_iou_threshold: float,
    ) -> List[int]:
        """Greedy NMS over correspondences, ranked by presence score.

        A candidate is suppressed only when it duplicates a higher-scoring one on
        *both* sides of the correspondence: image-mask IoU and text-span IoU both
        exceed `nms_iou_threshold`.
        """
        count = len(spans)
        if nms_iou_threshold <= 0.0 or count <= 1:
            return list(range(count))

        kept: List[int] = []
        for candidate in np.argsort(-presence_scores):
            candidate = int(candidate)
            suppressed = False
            for reference in kept:
                if masks[candidate] is not None and masks[reference] is not None:
                    if compute_mask_iou(masks[candidate], masks[reference]) < nms_iou_threshold:
                        continue
                if compute_span_iou(spans[candidate], spans[reference]) >= nms_iou_threshold:
                    suppressed = True
                    break
            if not suppressed:
                kept.append(candidate)
        return sorted(kept)

    def post_process_correspondences(
        self,
        outputs,
        text: Union[str, List[str]],
        target_sizes: Optional[Sequence[Tuple[int, int]]] = None,
        presence_threshold: Optional[float] = None,
        text_threshold: Optional[float] = None,
        image_threshold: Optional[float] = None,
        nms_iou_threshold: Optional[float] = None,
        return_masks: bool = True,
    ) -> List[List[Dict]]:
        """Turn per-bridge logits into a set of image-text correspondences.

        Bridge tokens whose presence probability passes `presence_threshold` are
        kept; each keeps a character-span text mask and a binary image mask
        (logits bilinearly upsampled to `target_sizes`, then thresholded), and
        duplicates are removed with NMS.

        Args:
            outputs: the [`ConCor1Output`] returned by the model.
            text: the same text that was passed to `__call__`.
            target_sizes: `(height, width)` per sample the image masks are
                resized to — normally the original image size. Defaults to the
                processed image resolution.
            presence_threshold / text_threshold / image_threshold /
            nms_iou_threshold: override the defaults (0.1 / 0.45 / 0.45 / 0.5).
            return_masks: set to `False` to skip mask upsampling and return text
                masks only.

        Returns:
            One list of correspondences per sample, sorted by decreasing
            presence score. Each correspondence is a dict with
            `bridge_index`, `presence_score`, `text_spans`, `text_phrases` and
            (unless disabled) `mask`.
        """
        presence_threshold = (
            self.presence_threshold if presence_threshold is None else presence_threshold
        )
        text_threshold = self.text_threshold if text_threshold is None else text_threshold
        image_threshold = self.image_threshold if image_threshold is None else image_threshold
        nms_iou_threshold = (
            self.nms_iou_threshold if nms_iou_threshold is None else nms_iou_threshold
        )

        texts = [text] if isinstance(text, str) else list(text)
        presence_probabilities = torch.sigmoid(outputs.presence_logits.float()).cpu().numpy()
        batch_size = presence_probabilities.shape[0]
        if len(texts) != batch_size:
            raise ValueError(f"Got {len(texts)} text(s) for a batch of {batch_size}.")

        results: List[List[Dict]] = []
        for sample in range(batch_size):
            sample_text = texts[sample]
            offset_mapping = self.tokenizer(
                sample_text, return_offsets_mapping=True, add_special_tokens=False
            )["offset_mapping"]

            active = np.where(presence_probabilities[sample] >= presence_threshold)[0]

            spans_per_bridge: List[List[List[int]]] = []
            masks_per_bridge: List[Optional[np.ndarray]] = []
            for bridge_index in active:
                text_logits = outputs.text_mask_logits[sample, bridge_index].float().cpu().numpy()
                spans_per_bridge.append(
                    self.text_logits_to_spans(
                        text_logits, offset_mapping, sample_text, text_threshold
                    )
                )
                masks_per_bridge.append(
                    self._decode_image_mask(
                        outputs,
                        sample=sample,
                        bridge_index=int(bridge_index),
                        target_size=None if target_sizes is None else tuple(target_sizes[sample]),
                        image_threshold=image_threshold,
                    )
                    if return_masks
                    else None
                )

            kept = self.suppress_duplicate_correspondences(
                masks_per_bridge,
                spans_per_bridge,
                presence_probabilities[sample][active],
                nms_iou_threshold,
            )

            correspondences = []
            for position in kept:
                bridge_index = int(active[position])
                spans = spans_per_bridge[position]
                correspondence = {
                    "bridge_index": bridge_index,
                    "presence_score": float(presence_probabilities[sample, bridge_index]),
                    "text_spans": spans,
                    "text_phrases": [sample_text[start:end] for start, end in spans],
                }
                if return_masks:
                    correspondence["mask"] = masks_per_bridge[position]
                correspondences.append(correspondence)

            correspondences.sort(key=lambda item: -item["presence_score"])
            results.append(correspondences)
        return results

    def _decode_image_mask(
        self,
        outputs,
        sample: int,
        bridge_index: int,
        target_size: Optional[Tuple[int, int]],
        image_threshold: float,
    ) -> Optional[np.ndarray]:
        """Reshape, upsample and threshold one bridge token's image-mask logits."""
        if outputs.image_mask_logits is None:
            return None

        height = int(outputs.image_mask_grid_hw[sample, 0].item())
        width = int(outputs.image_mask_grid_hw[sample, 1].item())
        if height == 0 or width == 0:
            return None

        logits = outputs.image_mask_logits[sample, bridge_index, : height * width]
        logits = logits.float().reshape(1, 1, height, width)

        if target_size is None:
            target_size = (height * self.mask_cell_size, width * self.mask_cell_size)

        # Upsample in logit space, then sigmoid — the mask resolution is set by the
        # decoder's cell grid, not by the interpolation.
        upsampled = F.interpolate(
            logits, size=tuple(int(value) for value in target_size), mode="bilinear", align_corners=False
        )
        probabilities = torch.sigmoid(upsampled)[0, 0].cpu().numpy()
        return probabilities >= image_threshold

    @property
    def mask_cell_size(self) -> int:
        """Pixels per side of one image-mask cell in the processed image (4 px)."""
        return self.patch_size // (2 ** self.num_mask_upsample_blocks)


__all__ = ["ConCor1Processor"]