"""Optional PyTorch hook utilities for collecting per-layer activations."""

from __future__ import annotations

from typing import Any, Callable


class TorchActivationCollector:
    """Collect layer activations from a PyTorch model via forward hooks."""

    def __init__(
        self,
        layer_filter: Callable[[str, Any], bool] | None = None,
        reducer: str = "mean",
    ) -> None:
        if reducer not in {"mean", "full"}:
            raise ValueError("reducer must be 'mean' or 'full'")
        self.layer_filter = layer_filter
        self.reducer = reducer
        self._handles: list[Any] = []
        self._latest: dict[str, list[float]] = {}

    @staticmethod
    def _extract_tensor(output: Any) -> Any:
        try:
            import torch
        except ImportError as exc:  # pragma: no cover
            raise ImportError("PyTorch is required for TorchActivationCollector") from exc

        if isinstance(output, torch.Tensor):
            return output
        if isinstance(output, (list, tuple)):
            for item in output:
                if isinstance(item, torch.Tensor):
                    return item
        if isinstance(output, dict):
            for item in output.values():
                if isinstance(item, torch.Tensor):
                    return item
        raise TypeError("hook output does not contain a tensor")

    def _make_hook(self, layer_name: str) -> Callable[[Any, Any, Any], None]:
        def hook(_module: Any, _inputs: Any, output: Any) -> None:
            tensor = self._extract_tensor(output).detach().float().reshape(-1)
            if self.reducer == "mean":
                values = [float(tensor.mean().item())]
            else:
                values = [float(v) for v in tensor.tolist()]
            self._latest[layer_name] = values

        return hook

    def attach(self, model: Any) -> None:
        """Attach hooks to matching named modules."""
        self.close()
        for name, module in model.named_modules():
            if self.layer_filter is not None and not self.layer_filter(name, module):
                continue
            handle = module.register_forward_hook(self._make_hook(name))
            self._handles.append(handle)

    def pop_batch(self) -> dict[str, list[float]]:
        """Return and clear the latest captured batch activations."""
        batch = dict(self._latest)
        self._latest.clear()
        return batch

    def close(self) -> None:
        while self._handles:
            self._handles.pop().remove()
        self._latest.clear()