"""
Sentence span extraction utilities for gradient attribution.

Provides mapping from text sentences to token position spans, enabling
aggregation of per-token gradients to per-sentence importance scores.
"""

import bisect
from typing import Any

import torch
from torch import Tensor


def _normalize_with_mapping(text: str) -> tuple[str, list[int]]:
    """Normalize whitespace and track mapping to original character indices."""
    normalized_chars: list[str] = []
    normalized_to_original: list[int] = []
    prev_was_space = False

    for idx, ch in enumerate(text):
        if ch.isspace():
            # Collapse any whitespace run (including CRLF) into a single space.
            if normalized_chars and not prev_was_space:
                normalized_chars.append(" ")
                normalized_to_original.append(idx)
            prev_was_space = True
            continue

        normalized_chars.append(ch)
        normalized_to_original.append(idx)
        prev_was_space = False

    # Trim trailing normalized spaces so sentence-level matching is stable.
    if normalized_chars and normalized_chars[-1] == " ":
        normalized_chars.pop()
        normalized_to_original.pop()

    return "".join(normalized_chars), normalized_to_original


def get_sentence_token_spans(
    text: str,
    sentences: list[str],
    tokenizer: Any,
    add_special_tokens: bool = False,
) -> list[tuple[int, int]]:
    """
    Map text sentences to their corresponding token position spans.

    Given a full text and its constituent sentences, determines the start and
    end token indices for each sentence in the tokenized representation.

    Args:
        text: The full document text.
        sentences: List of sentences that compose the text. Should be in order
            and their concatenation should approximately match the text.
        tokenizer: A HuggingFace tokenizer with offset mapping support.
        add_special_tokens: Whether to include special tokens (BOS/EOS) in
            tokenization. Default False for compatibility with CHG.

    Returns:
        List of (start_idx, end_idx) tuples, one per sentence. Indices are
        inclusive of start, exclusive of end (Python slice convention).
        If a sentence couldn't be mapped, returns (0, 0) for that entry.

    Example:
        >>> text = "Hello world. How are you?"
        >>> sentences = ["Hello world.", "How are you?"]
        >>> spans = get_sentence_token_spans(text, sentences, tokenizer)
        >>> # spans might be [(0, 3), (3, 7)] depending on tokenizer
    """
    # Tokenize full text with offset mapping
    encoding = tokenizer(
        text,
        add_special_tokens=add_special_tokens,
        return_offsets_mapping=True,
    )

    offset_mapping = encoding["offset_mapping"]  # List of (char_start, char_end) per token

    # Build sorted arrays for O(log n) char→token lookups via binary search.
    # Assumes offset_mapping is in non-decreasing character order (standard HF behaviour).
    tok_starts = [s for s, _ in offset_mapping]
    tok_ends = [e for _, e in offset_mapping]

    def _first_token(char_start: int, char_end: int) -> int | None:
        # First token whose char range overlaps [char_start, char_end).
        idx = bisect.bisect_left(tok_ends, char_start + 1)
        if idx < len(tok_ends) and tok_starts[idx] < char_end:
            return idx
        return None

    def _last_token(char_start: int, char_end: int) -> int | None:
        # Last token whose char range overlaps [char_start, char_end).
        idx = bisect.bisect_left(tok_starts, char_end) - 1
        if idx >= 0 and tok_ends[idx] > char_start:
            return idx
        return None

    # Normalize whitespace in both text and sentence strings while preserving
    # an index mapping to original text offsets for token conversion.
    normalized_text, norm_to_orig = _normalize_with_mapping(text)

    # Find each sentence's character span in the full text
    spans = []
    search_start_norm = 0

    for sentence in sentences:
        sentence_clean = sentence.strip()
        if not sentence_clean:
            spans.append((0, 0))
            continue

        normalized_sentence, _ = _normalize_with_mapping(sentence_clean)
        if not normalized_sentence:
            spans.append((0, 0))
            continue

        # Find sentence in normalized text starting from the last position.
        norm_start = normalized_text.find(normalized_sentence, search_start_norm)
        if norm_start == -1:
            spans.append((0, 0))
            continue

        norm_end = norm_start + len(normalized_sentence)
        search_start_norm = norm_end

        # Convert normalized offsets back to original text offsets.
        char_start = norm_to_orig[norm_start]
        char_end = norm_to_orig[norm_end - 1] + 1  # Exclusive

        # Map character positions to token positions via binary search
        token_start = _first_token(char_start, char_end)
        last = _last_token(char_start, char_end)
        token_end = last + 1 if last is not None else None  # exclusive end

        if token_start is not None and token_end is not None:
            spans.append((token_start, token_end))
        else:
            spans.append((0, 0))

    return spans


def get_sentence_token_spans_from_pretokenized(
    sentences: list[str],
    tokenizer: Any,
    add_special_tokens: bool = False,
) -> tuple[Tensor, list[tuple[int, int]]]:
    """
    Tokenize sentences individually and return combined tokens with spans.

    Alternative to get_sentence_token_spans when sentences may not exactly
    reconstruct the original text. Tokenizes each sentence separately and
    concatenates, tracking span boundaries.

    Args:
        sentences: List of sentences to tokenize.
        tokenizer: A HuggingFace tokenizer.
        add_special_tokens: Whether to add special tokens per sentence.
            Typically False to avoid multiple BOS/EOS tokens.

    Returns:
        Tuple of:
        - input_ids: Tensor of shape [seq_len] with all tokens concatenated.
        - spans: List of (start_idx, end_idx) tuples for each sentence.

    Example:
        >>> sentences = ["Hello world.", "How are you?"]
        >>> input_ids, spans = get_sentence_token_spans_from_pretokenized(
        ...     sentences, tokenizer
        ... )
    """
    all_input_ids = []
    spans = []
    current_pos = 0

    for sentence in sentences:
        if not sentence.strip():
            spans.append((current_pos, current_pos))
            continue

        encoding = tokenizer(
            sentence,
            add_special_tokens=add_special_tokens,
            return_tensors="pt",
        )

        sent_ids = encoding["input_ids"].squeeze(0)
        sent_len = len(sent_ids)

        all_input_ids.append(sent_ids)
        spans.append((current_pos, current_pos + sent_len))
        current_pos += sent_len

    if all_input_ids:
        input_ids = torch.cat(all_input_ids, dim=0)
    else:
        input_ids = torch.tensor([], dtype=torch.long)

    return input_ids, spans


def aggregate_to_sentences(
    token_values: Tensor,
    spans: list[tuple[int, int]],
    aggregation: str = "mean",
) -> Tensor:
    """
    Aggregate per-token values to per-sentence values using span information.

    Args:
        token_values: Tensor of shape [seq_len] or [seq_len, ...] with
            per-token values (e.g., gradient norms, attention weights).
        spans: List of (start_idx, end_idx) tuples from get_sentence_token_spans.
        aggregation: How to aggregate tokens within a sentence:
            - "mean": Average of token values (default)
            - "sum": Sum of token values
            - "max": Maximum token value

    Returns:
        Tensor of shape [num_sentences] with aggregated values per sentence.

    Example:
        >>> grad_norms = torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5])
        >>> spans = [(0, 2), (2, 5)]
        >>> sentence_scores = aggregate_to_sentences(grad_norms, spans, "mean")
        >>> # sentence_scores = [0.15, 0.4]
    """
    sentence_values = []

    for start, end in spans:
        if start >= end:
            # Empty span
            sentence_values.append(token_values.new_tensor(0.0))
            continue

        # Clamp indices to valid range
        start = max(0, start)
        end = min(len(token_values), end)

        span_values = token_values[start:end]

        if len(span_values) == 0:
            sentence_values.append(token_values.new_tensor(0.0))
            continue

        if aggregation == "mean":
            agg_value = span_values.mean()
        elif aggregation == "sum":
            agg_value = span_values.sum()
        elif aggregation == "max":
            agg_value = span_values.max()
        else:
            raise ValueError(f"Unknown aggregation method: {aggregation}")

        sentence_values.append(agg_value)

    if not sentence_values:
        return token_values.new_empty((0,))

    return torch.stack(sentence_values)
