"""
Advanced PyTorch hook utilities for transformer analysis.

Beyond the ``register_forward_pre_hook`` / ``register_forward_hook`` pair used
internally by :class:`~inteprelens.core.tracer.TransformerTracer`, PyTorch
exposes several other hook points that are useful for mechanistic
interpretability.  This module wraps each one in a context manager so handles
are always cleaned up, and provides runnable examples for
``meta-llama/Llama-3.2-1B``.

Hook surface overview
---------------------
+------------------------------------------+-------------------------------+------------------+
| API                                      | Fires                         | PyTorch version  |
+==========================================+===============================+==================+
| ``register_full_backward_hook``          | backward, per module          | ≥1.8             |
+------------------------------------------+-------------------------------+------------------+
| ``register_forward_pre_hook(with_kwargs)``| forward, inspect/modify kwargs| ≥2.0             |
+------------------------------------------+-------------------------------+------------------+
| ``Tensor.register_hook``                 | backward, per tensor          | always           |
+------------------------------------------+-------------------------------+------------------+
| ``autograd.graph.Node.register_hook``    | backward, per graph node      | ≥2.0             |
+------------------------------------------+-------------------------------+------------------+
| ``register_module_full_backward_pre_hook``| backward, fires *before* grad | ≥2.1             |
+------------------------------------------+-------------------------------+------------------+
"""

from __future__ import annotations

import contextlib
from collections.abc import Callable, Generator
from typing import Literal

from torch import Tensor, nn

# ---------------------------------------------------------------------------
# 1. Full backward hook — capture grad_input / grad_output per module
# ---------------------------------------------------------------------------

@contextlib.contextmanager
def capture_gradients(
    module: nn.Module,
    site: Literal["input", "output"] = "output",
) -> Generator[list[tuple[Tensor | None, ...]], None, None]:
    """
    Capture gradient tensors flowing through *module* during ``.backward()``.

    Uses ``register_full_backward_hook``, which supersedes the deprecated
    ``register_backward_hook`` and correctly handles modules with multiple
    inputs/outputs.

    Args:
        module: The module to attach the hook to.
        site: ``"output"`` captures ``grad_output`` (gradient w.r.t. the module
            output, flowing *into* the module from downstream); ``"input"``
            captures ``grad_input`` (gradient w.r.t. the module input, flowing
            *out* to upstream).

    Yields:
        A list that accumulates ``(grad_input,)`` or ``(grad_output,)`` tuples
        for every backward call that traverses *module*.

    Example — gradient of attention output w.r.t. a loss::

        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        from inteprelens import LensAPI
        from inteprelens.models.adapters import get_adapter
        from inteprelens.core.hooks import capture_gradients

        model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16
        ).cuda()
        adapter = get_adapter(model)
        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")

        inputs = tokenizer("The Eiffel Tower is in", return_tensors="pt").to("cuda")
        target_id = tokenizer(" Paris", add_special_tokens=False)["input_ids"][0]

        attn_module = adapter.get_attention_module(layer=0)
        with capture_gradients(attn_module, site="output") as grads:
            out = model(**inputs)
            loss = -out.logits[0, -1, target_id]
            loss.backward()
        # grads[0][0] shape: [batch, seq, hidden]
        grad_tensor = grads[0][0]
        print("grad norm:", grad_tensor.float().norm().item())
    """
    captured: list[tuple[Tensor | None, ...]] = []

    def _hook(
        mod: nn.Module,
        grad_input: tuple[Tensor | None, ...],
        grad_output: tuple[Tensor | None, ...],
    ) -> None:
        captured.append(grad_output if site == "output" else grad_input)

    handle = module.register_full_backward_hook(_hook)
    try:
        yield captured
    finally:
        handle.remove()


# ---------------------------------------------------------------------------
# 2. Forward pre-hook with kwargs — inspect / modify attention_mask etc.
# ---------------------------------------------------------------------------

@contextlib.contextmanager
def capture_kwargs_hook(
    module: nn.Module,
) -> Generator[list[dict], None, None]:
    """
    Capture (and optionally modify) keyword arguments passed to *module*.

    Uses ``register_forward_pre_hook(with_kwargs=True)`` (PyTorch ≥ 2.0).
    This gives access to kwargs such as ``attention_mask``, ``position_ids``,
    and ``past_key_values`` that are invisible to the classic pre-hook.

    Args:
        module: The module to attach the hook to (e.g., a transformer block).

    Yields:
        A list that accumulates ``{**kwargs}`` dicts captured before each
        forward call on *module*.

    Example — log attention_mask shape for every layer-0 forward call::

        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        from inteprelens.models.adapters import get_adapter
        from inteprelens.core.hooks import capture_kwargs_hook

        model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16
        ).cuda()
        adapter = get_adapter(model)
        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")

        inputs = tokenizer(["Paris is the capital of", "Hello world"],
                           return_tensors="pt", padding=True).to("cuda")
        block0 = adapter.get_layers()[0]
        with capture_kwargs_hook(block0) as all_kwargs:
            with torch.no_grad():
                model(**inputs)
        mask = all_kwargs[0].get("attention_mask")
        if mask is not None:
            print("attention_mask shape:", mask.shape)  # [batch, seq]
    """
    captured: list[dict] = []

    def _hook(
        mod: nn.Module,
        args: tuple,
        kwargs: dict,
    ) -> None:
        captured.append(dict(kwargs))

    handle = module.register_forward_pre_hook(_hook, with_kwargs=True)
    try:
        yield captured
    finally:
        handle.remove()


# ---------------------------------------------------------------------------
# 3. Tensor gradient hook — probe gradient on a specific activation
# ---------------------------------------------------------------------------

@contextlib.contextmanager
def tensor_grad_hook(
    get_tensor_fn: Callable[[], Tensor],
    retain_grad: bool = True,
) -> Generator[list[Tensor], None, None]:
    """
    Register a gradient hook on an arbitrary activation tensor.

    ``Tensor.register_hook(fn)`` calls *fn* with the gradient flowing through
    that tensor during ``.backward()``.  The tensor must be part of the
    computation graph (``requires_grad=True`` or a leaf with ``retain_grad()``
    called before backward).

    Args:
        get_tensor_fn: Zero-argument callable that returns the target tensor
            *after* the forward pass and *before* ``backward()``.  Typically a
            closure over a captured activation (see example).
        retain_grad: If ``True``, call ``tensor.retain_grad()`` so that
            non-leaf tensors keep their ``.grad`` attribute after backward.

    Yields:
        A list that accumulates gradient tensors as they arrive.

    Example — gradient of final hidden state w.r.t. cross-entropy loss::

        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        from inteprelens.core.hooks import tensor_grad_hook

        model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16,
            output_hidden_states=True,
        ).cuda()
        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
        inputs = tokenizer("The capital of France is", return_tensors="pt").to("cuda")

        hidden_ref: list[Tensor] = []
        def _fwd_hook(mod, inp, out):
            hidden_ref.append(out[0])  # final hidden states

        handle = model.register_forward_hook(_fwd_hook)
        out = model(**inputs)
        handle.remove()

        hidden = hidden_ref[0]   # [batch, seq, hidden]
        grads: list[Tensor] = []
        def _get(): return hidden

        with tensor_grad_hook(_get) as grads:
            loss = -out.logits[0, -1].log_softmax(-1).max()
            loss.backward()
        print("hidden grad shape:", grads[0].shape)  # [batch, seq, hidden]
    """
    captured: list[Tensor] = []
    tensor = get_tensor_fn()
    if retain_grad and tensor.requires_grad:
        tensor.retain_grad()

    def _hook(grad: Tensor) -> None:
        captured.append(grad.detach())

    handle = tensor.register_hook(_hook)
    try:
        yield captured
    finally:
        handle.remove()


# ---------------------------------------------------------------------------
# 4. Autograd graph node hook — low-level backward traversal
# ---------------------------------------------------------------------------

@contextlib.contextmanager
def autograd_node_hook(
    tensor: Tensor,
) -> Generator[list[tuple[Tensor, ...]], None, None]:
    """
    Attach a hook to the autograd graph node that *produced* ``tensor``.

    ``torch.autograd.graph.Node.register_hook`` (PyTorch ≥ 2.0) fires once
    per backward traversal of the node, receiving the gradient inputs to that
    node (i.e., the gradients flowing *into* the operation that created the
    tensor from upstream).

    Args:
        tensor: A non-leaf tensor whose ``.grad_fn`` node will be hooked.

    Yields:
        A list accumulating gradient-input tuples per backward call.

    Example — hook on the MLP output node::

        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        from inteprelens.models.adapters import get_adapter
        from inteprelens.core.hooks import autograd_node_hook

        model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16
        ).cuda()
        adapter = get_adapter(model)
        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
        inputs = tokenizer("Hello world", return_tensors="pt").to("cuda")

        mlp_out: list[Tensor] = []
        fwd_handle = adapter.get_mlp(layer=0).register_forward_hook(
            lambda m, i, o: mlp_out.append(o)
        )
        out = model(**inputs)
        fwd_handle.remove()

        act = mlp_out[0]  # [batch, seq, hidden]; has .grad_fn
        with autograd_node_hook(act) as node_grads:
            loss = -out.logits[0, -1].log_softmax(-1).max()
            loss.backward()
        print("node grad inputs:", [g.shape for g in node_grads[0]])
    """
    captured: list[tuple[Tensor, ...]] = []

    if tensor.grad_fn is None:
        raise ValueError(
            "tensor has no grad_fn — it must be a non-leaf tensor in the computation graph."
        )

    def _hook(grad_inputs: tuple[Tensor, ...], grad_outputs: tuple[Tensor, ...]) -> None:
        captured.append(tuple(g.detach() for g in grad_inputs if g is not None))

    handle = tensor.grad_fn.register_hook(_hook)
    try:
        yield captured
    finally:
        handle.remove()


# ---------------------------------------------------------------------------
# 5. Module full-backward-pre-hook — fires before module backward (≥ 2.1)
# ---------------------------------------------------------------------------

@contextlib.contextmanager
def capture_backward_pre(
    module: nn.Module,
) -> Generator[list[tuple[Tensor | None, ...]], None, None]:
    """
    Capture ``grad_output`` **before** a module's backward pass executes.

    ``register_module_full_backward_pre_hook`` (PyTorch ≥ 2.1) fires
    immediately when backward reaches *module*, before any gradient is
    propagated to inputs.  This is useful for inspecting or modifying the
    gradient signal entering a layer.

    Args:
        module: The module to attach the hook to.

    Yields:
        A list accumulating ``grad_output`` tuples.

    Example — inspect gradient entering layer-0 attention before backprop::

        from transformers import AutoModelForCausalLM, AutoTokenizer
        import torch
        from inteprelens.models.adapters import get_adapter
        from inteprelens.core.hooks import capture_backward_pre

        model = AutoModelForCausalLM.from_pretrained(
            "meta-llama/Llama-3.2-1B", torch_dtype=torch.bfloat16
        ).cuda()
        adapter = get_adapter(model)
        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B")
        inputs = tokenizer("The Eiffel Tower is in Paris", return_tensors="pt").to("cuda")

        attn = adapter.get_attention_module(layer=0)
        with capture_backward_pre(attn) as pre_grads:
            out = model(**inputs)
            loss = -out.logits[0, -1].log_softmax(-1).max()
            loss.backward()
        # pre_grads[0][0]: grad_output entering attn-layer-0 backward
        print("pre-backward grad shape:", pre_grads[0][0].shape)
    """
    captured: list[tuple[Tensor | None, ...]] = []

    def _hook(
        mod: nn.Module,
        grad_output: tuple[Tensor | None, ...],
    ) -> None:
        captured.append(grad_output)

    handle = module.register_full_backward_pre_hook(_hook)
    try:
        yield captured
    finally:
        handle.remove()
