"""
Gradient-based attribution for sentence importance scoring.

Implements the core method for computing per-sentence importance scores
by measuring gradient flow through CHG-identified facilitating heads.
This provides a causally-grounded sentence importance measure.
"""

import logging
from collections.abc import Callable
from typing import Any

import torch
import torch.nn.functional as F
from torch import Tensor, nn

from ..utils.sentence_spans import aggregate_to_sentences
from .chg import CHG

logger = logging.getLogger(__name__)


def _is_oom_error(exc: RuntimeError) -> bool:
    msg = str(exc).lower()
    return "out of memory" in msg or "cuda error: out of memory" in msg


class CausalCircuitAttribution:
    """
    Computes sentence importance via gradient attribution through facilitating heads.

    The key insight is that CHG identifies which attention heads are causally
    necessary for the task/dataset. By computing gradients of the summary
    likelihood with respect to input embeddings, and doing so with only
    facilitating heads active, we measure how much each input token contributes
    to summary generation through the causal circuit.

    This provides a principled, causally-grounded sentence importance score.

    Attributes:
        model: The transformer model.
        tokenizer: The tokenizer for the model.
        chg: CHG wrapper for applying head masks.
        device: Computation device.

    Example:
        >>> attr = CausalCircuitAttribution(model, tokenizer)
        >>> mask = results.get_facilitating_mask()
        >>> scores = attr.compute_sentence_importance(
        ...     document="The quick brown fox...",
        ...     sentences=["The quick brown fox.", "..."],
        ...     summary="Fox jumps over dog.",
        ...     facilitating_mask=mask,
        ... )
    """

    def __init__(
        self,
        model: nn.Module,
        tokenizer: Any,
        enable_gradient_checkpointing: bool = True,
    ):
        """
        Initialize CausalCircuitAttribution.

        Args:
            model: A HuggingFace transformer model (e.g., from AutoModelForCausalLM).
            tokenizer: The corresponding tokenizer.
            enable_gradient_checkpointing: If True, enables gradient checkpointing
                to reduce memory usage at the cost of compute. Recommended for
                longer documents on limited GPU memory.
        """
        self.model = model
        self.tokenizer = tokenizer
        self.chg = CHG(model)
        self.device = next(model.parameters()).device

        # We only need gradients wrt input embeddings, not model weights.
        for param in self.model.parameters():
            param.requires_grad = False

        if enable_gradient_checkpointing:
            if hasattr(model, "gradient_checkpointing_enable"):
                model.gradient_checkpointing_enable()

    def _run_gradient_pass(
        self,
        embeds: Tensor,
        input_len: int,
        mask: Tensor,
        loss_fn: Callable[[Tensor], Tensor],
    ) -> Tensor:
        """Run a masked forward/backward pass and return per-token gradient norms."""
        with self.chg.register_mask_hooks(mask):
            outputs = self.model(inputs_embeds=embeds)
            loss = loss_fn(outputs.logits)
        self.model.zero_grad(set_to_none=True)
        loss.backward()
        grads = embeds.grad[0, :input_len, :]
        return grads.norm(dim=-1).detach()

    def compute_token_gradients(
        self,
        input_ids: Tensor,
        target_ids: Tensor,
        facilitating_mask: Tensor,
    ) -> Tensor:
        """
        Compute gradient norms for each input token position.

        Runs a forward pass with the facilitating head mask applied, computes
        the log-likelihood loss for predicting the target tokens, and returns
        the L2 norm of gradients with respect to input embeddings.

        Args:
            input_ids: Token IDs for the input document. Shape [1, seq_len].
            target_ids: Token IDs for the target summary. Shape [1, target_len].
            facilitating_mask: Binary mask from CHGResults.get_facilitating_mask().
                Shape [num_layers, num_heads].

        Returns:
            Tensor of gradient norms per input position. Shape [seq_len].
        """
        self.model.eval()
        embed_layer = self.model.get_input_embeddings()
        embeds = embed_layer(input_ids).clone().detach().requires_grad_(True)
        mask = facilitating_mask.unsqueeze(0).to(self.device)

        def loss_fn(logits: Tensor) -> Tensor:
            log_probs = F.log_softmax(logits[0, -1, :], dim=-1)
            return -log_probs[target_ids[0, 0]]

        return self._run_gradient_pass(embeds, input_ids.shape[1], mask, loss_fn)

    def compute_token_gradients_full_target(
        self,
        input_ids: Tensor,
        target_ids: Tensor,
        facilitating_mask: Tensor,
    ) -> Tensor:
        """
        Compute gradient norms using full target sequence likelihood.

        More thorough version that uses teacher-forced generation to compute
        gradients for the full target sequence. More expensive but potentially
        more accurate for longer summaries.

        Args:
            input_ids: Token IDs for input document. Shape [1, seq_len].
            target_ids: Token IDs for target summary. Shape [1, target_len].
            facilitating_mask: Binary mask. Shape [num_layers, num_heads].

        Returns:
            Tensor of gradient norms per input position. Shape [seq_len].
        """
        self.model.eval()
        full_ids = torch.cat([input_ids, target_ids], dim=1)
        embed_layer = self.model.get_input_embeddings()
        embeds = embed_layer(full_ids).clone().detach().requires_grad_(True)
        mask = facilitating_mask.unsqueeze(0).to(self.device)
        input_len = input_ids.shape[1]
        target_len = target_ids.shape[1]

        def loss_fn(logits: Tensor) -> Tensor:
            target_logits = logits[0, input_len - 1:input_len + target_len - 1, :]
            return F.cross_entropy(target_logits, target_ids[0], reduction="mean")

        return self._run_gradient_pass(embeds, input_len, mask, loss_fn)

    def compute_sentence_importance(
        self,
        document: str,
        sentences: list[str],
        summary: str,
        facilitating_mask: Tensor,
        spans: list[tuple[int, int]] | None = None,
        use_full_target: bool = True,
        aggregation: str = "mean",
        max_input_tokens: int | None = None,
        max_target_tokens: int | None = None,
        oom_retry_fast: bool = True,
    ) -> Tensor:
        """
        Compute importance scores for each sentence in a document.

        Main entry point for sentence-level gradient attribution. Tokenizes
        the document, computes per-token gradients through facilitating heads,
        and aggregates to sentence-level scores.

        Args:
            document: The full document text.
            sentences: List of sentences in the document.
            summary: The reference summary.
            facilitating_mask: Binary mask from CHGResults.get_facilitating_mask().
            spans: Pre-computed sentence token spans. If None, computed automatically.
            use_full_target: If True, uses full target likelihood (more accurate).
                If False, uses single-token approximation (faster).
            aggregation: How to aggregate token scores to sentences:
                "mean", "sum", or "max".
            max_input_tokens: Optional max tokens for document truncation.
            max_target_tokens: Optional max tokens for summary truncation.
            oom_retry_fast: If True, retry with single-token gradients on OOM.

        Returns:
            Tensor of importance scores. Shape [num_sentences].
        """
        from ..utils.sentence_spans import get_sentence_token_spans

        # Tokenize document
        doc_encoding = self.tokenizer(
            document,
            add_special_tokens=True,
            truncation=max_input_tokens is not None,
            max_length=max_input_tokens,
            return_tensors="pt",
        )
        input_ids = doc_encoding["input_ids"].to(self.device)

        # Tokenize summary
        summary_encoding = self.tokenizer(
            summary,
            add_special_tokens=False,
            truncation=max_target_tokens is not None,
            max_length=max_target_tokens,
            return_tensors="pt",
        )
        target_ids = summary_encoding["input_ids"].to(self.device)

        input_len = int(input_ids.shape[1])
        target_len = int(target_ids.shape[1])

        if target_len == 0:
            logger.warning(
                "Empty target after tokenization/truncation; returning NaN scores "
                "(input_tokens=%d, target_tokens=%d).",
                input_len,
                target_len,
            )
            return torch.full((len(sentences),), float("nan"), dtype=torch.float32)

        # Get sentence spans if not provided
        if spans is None:
            spans = get_sentence_token_spans(
                document, sentences, self.tokenizer, add_special_tokens=True
            )

        empty_spans = sum(1 for start, end in spans if start >= end)
        if spans:
            empty_ratio = empty_spans / len(spans)
            if empty_ratio > 0.3:
                logger.warning(
                    "High empty-span ratio in sentence mapping: %.2f (%d/%d).",
                    empty_ratio,
                    empty_spans,
                    len(spans),
                )

        # Compute token gradients with OOM fallback.
        try:
            if use_full_target:
                grad_norms = self.compute_token_gradients_full_target(
                    input_ids, target_ids, facilitating_mask
                )
            else:
                grad_norms = self.compute_token_gradients(
                    input_ids, target_ids, facilitating_mask
                )
        except RuntimeError as exc:
            if not _is_oom_error(exc):
                raise

            if torch.cuda.is_available():
                torch.cuda.empty_cache()

            if not (oom_retry_fast and use_full_target):
                logger.warning(
                    "OOM during sentence attribution; returning NaN scores "
                    "(input_tokens=%d, target_tokens=%d, use_full_target=%s).",
                    input_len,
                    target_len,
                    use_full_target,
                )
                return torch.full((len(sentences),), float("nan"), dtype=torch.float32)

            logger.warning(
                "OOM with full-target attribution; retrying single-token mode "
                "(input_tokens=%d, target_tokens=%d).",
                input_len,
                target_len,
            )
            try:
                grad_norms = self.compute_token_gradients(
                    input_ids, target_ids, facilitating_mask
                )
            except RuntimeError as retry_exc:
                if not _is_oom_error(retry_exc):
                    raise
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                logger.warning(
                    "OOM in fallback attribution; returning NaN scores "
                    "(input_tokens=%d, target_tokens=%d).",
                    input_len,
                    target_len,
                )
                return torch.full((len(sentences),), float("nan"), dtype=torch.float32)

        # Aggregate to sentences
        sentence_scores = aggregate_to_sentences(grad_norms, spans, aggregation)

        # Normalize to [0, 1]
        finite = torch.isfinite(sentence_scores)
        if finite.any():
            max_score = sentence_scores[finite].max()
            if max_score > 0:
                sentence_scores = sentence_scores / max_score

        return sentence_scores

    def compute_sentence_importance_batch(
        self,
        documents: list[str],
        sentences_list: list[list[str]],
        summaries: list[str],
        facilitating_mask: Tensor,
        use_full_target: bool = True,
        aggregation: str = "mean",
        max_input_tokens: int | None = None,
        max_target_tokens: int | None = None,
        oom_retry_fast: bool = True,
    ) -> list[Tensor]:
        """
        Compute importance scores for multiple documents.

        Processes documents one at a time (batch size 1) to manage memory,
        but provides a convenient interface for scoring entire datasets.

        Args:
            documents: List of document texts.
            sentences_list: List of sentence lists, one per document.
            summaries: List of reference summaries.
            facilitating_mask: Shared facilitating mask for all documents.
            use_full_target: Whether to use full target likelihood.
            aggregation: Token-to-sentence aggregation method.
            max_input_tokens: Optional max tokens for document truncation.
            max_target_tokens: Optional max tokens for summary truncation.
            oom_retry_fast: Whether to retry single-token attribution on OOM.

        Returns:
            List of importance score tensors, one per document.
        """
        all_scores = []

        for doc, sents, summary in zip(documents, sentences_list, summaries, strict=False):
            scores = self.compute_sentence_importance(
                document=doc,
                sentences=sents,
                summary=summary,
                facilitating_mask=facilitating_mask,
                use_full_target=use_full_target,
                aggregation=aggregation,
                max_input_tokens=max_input_tokens,
                max_target_tokens=max_target_tokens,
                oom_retry_fast=oom_retry_fast,
            )
            all_scores.append(scores)

        return all_scores
