Vision¶

class VisionModel:
    """Unified interface for extracting DNN activations from images.

    Composes a :class:`~vneurotk.vision.model.backend.base.BaseBackend`
    and a :class:`~vneurotk.vision.model.selector.ModuleSelector`.
    Activations are returned as-is; any further processing (pooling,
    embedding, etc.) is left to the user.

    Parameters
    ----------
    model_id : str
        Model identifier passed directly to the backend, e.g.
        ``"facebook/dinov2-base"`` (transformers) or ``"resnet50"`` (timm).
    backend : str
        Backend to use: ``"transformers"`` (default), ``"timm"``, or
        ``"thingsvision"``.
    selector : ModuleSelector or None
        Layer selection strategy.  Defaults to
        :class:`~vneurotk.vision.model.selector.BlockLevelSelector`.
    device : str
        Inference device (default ``"cpu"``).
    pretrained : bool
        Load pretrained weights (default ``True``).
    """

    def __init__(
        self,
        model_id: str,
        backend: str = "transformers",
        selector: ModuleSelector | None = None,
        device: str = "cpu",
        pretrained: bool = True,
    ) -> None:
        self._selector = selector or BlockLevelSelector()

        self._backend = self._build_backend(backend, device)
        self._backend.load(model_id, pretrained=pretrained)

        self._bind_selector()

        logger.info(
            "VisionModel ready | model={} | backend={} | modules={}",
            model_id,
            backend,
            len(self._module_names),
        )

    # ------------------------------------------------------------------
    # Class-method constructor
    # ------------------------------------------------------------------

    @classmethod
    def from_model(
        cls,
        model: Any,
        backend: BaseBackend,
        selector: ModuleSelector | None = None,
    ) -> VisionModel:
        """Build a VisionModel from an already-loaded model.

        Parameters
        ----------
        model : nn.Module
            Pre-loaded PyTorch model.
        backend : BaseBackend
            Backend instance with *model* already assigned.
        selector : ModuleSelector or None
            Defaults to :class:`BlockLevelSelector`.

        Returns
        -------
        VisionModel
        """
        inst = object.__new__(cls)
        inst._selector = selector or BlockLevelSelector()
        inst._backend = backend
        inst._backend.model = model

        inst._bind_selector()
        return inst

    # ------------------------------------------------------------------
    # Internal setup
    # ------------------------------------------------------------------

    def _bind_selector(self) -> None:
        """Select modules, register hooks, and build the module-type map.

        Called by ``__init__``, ``from_model``, and ``set_selector`` so the
        wiring logic lives in exactly one place.
        """
        module_names = self._selector.select(self._backend.enumerate_modules())
        self._backend.register_hooks(module_names)
        self._module_names = module_names
        self._module_type_map: dict[str, str] = {m.name: m.module_type for m in self._backend.enumerate_modules()}

    # ------------------------------------------------------------------
    # Extraction
    # ------------------------------------------------------------------

    @staticmethod
    def _normalize_images(image: Any) -> dict:
        """Normalise any supported input type to a ``{stim_id: image}`` dict.

        Parameters
        ----------
        image : dict or single image
            ``dict`` → returned as-is.
            Single image (PIL, ndarray, str, Path) → ``{0: image}``.

        Returns
        -------
        dict

        Raises
        ------
        TypeError
            If *image* is a ``list``.  Use ``{0: img0, 1: img1, …}`` or any
            mapping with meaningful stim IDs instead — a list would silently
            assign integer stim IDs that won't align with
            ``BaseData.trial_stim_ids``.
        """
        if isinstance(image, list):
            raise TypeError(
                "batch input requires a dict mapping stim_id → image; "
                "got a list.  Use {0: img0, 1: img1, …} for integer stim IDs, "
                "or {name: img, …} for named stimuli."
            )
        if isinstance(image, dict):
            return image
        return {0: image}

    def extract(
        self,
        image: Any,
        batch_size: int = DEFAULT_BATCH_SIZE,
        show_progress: bool = True,
    ) -> VisualRepresentations:
        """Extract DNN activations for one image or a collection of stimuli.

        Parameters
        ----------
        image : PIL.Image.Image or np.ndarray or str or Path or dict
            Single image → ``n_sample=1``, ``stim_id=0``.
            ``dict`` mapping stim_ids to images → ``n_sample=len(dict)``.
            String / ``pathlib.Path`` values are opened automatically.
            ``list`` is **not** accepted — use a ``dict`` with explicit
            stim IDs to preserve alignment with ``BaseData.trial_stim_ids``.
        batch_size : int
            Number of images per GPU forward pass.  Default 32.
            Ignored for single-image input.
        show_progress : bool
            Display a tqdm progress bar over batches.  Automatically
            suppressed for single-image input.  Default ``True``.

        Returns
        -------
        VisualRepresentations
        """
        if batch_size <= 0:
            raise ValueError(f"batch_size must be positive, got {batch_size}")

        images = self._normalize_images(image)
        is_single = not isinstance(image, dict)
        return self.extract_for_modules(
            images,
            self._module_names,
            batch_size=1 if is_single else batch_size,
            show_progress=False if is_single else show_progress,
        )

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def model_id(self) -> str:
        """Model identifier (e.g. ``'facebook/dinov2-base'``, ``'resnet50'``)."""
        return self._backend.get_model_meta().model_id

    @property
    def module_names(self) -> list[str]:
        """Names of all currently hooked modules."""
        return list(self._module_names)

    @property
    def module_list(self) -> list:
        """All modules available in the loaded model.

        Returns
        -------
        list[ModuleInfo]
            One entry per named module, ordered as ``model.named_modules()``.
            Each entry exposes ``.name``, ``.module_type``, ``.depth``,
            and ``.n_params``.
        """
        return self._backend.enumerate_modules()

    # ------------------------------------------------------------------
    # Inspection & reconfiguration
    # ------------------------------------------------------------------

    def print_modules(self, max_depth: int | None = None, console: Any = None) -> None:
        """Print a tree-style summary of all model modules.

        Parameters
        ----------
        max_depth : int or None
            Maximum nesting depth to display.  ``None`` shows all levels.
        console : rich.console.Console or None
            Rich console to use for output.  Pass ``Console(record=True)`` to
            capture the output for SVG / HTML export.
        """
        _print_modules(self.module_list, max_depth=max_depth, console=console)

    @staticmethod
    def _filter_modules(
        module_list: list,
        types: set[str],
        names: set[str],
    ) -> list:
        """Filter *module_list* by module type and/or exact name.

        Returns all modules whose ``module_type`` is in *types* OR whose
        ``name`` is in *names*.  Both sets may be empty; an entry matches if
        it satisfies at least one non-empty criterion.

        Parameters
        ----------
        module_list : list[ModuleInfo]
            Full module list from :attr:`module_list`.
        types : set[str]
            ``module_type`` values to include.
        names : set[str]
            Exact ``name`` values to include.

        Returns
        -------
        list[ModuleInfo]
            Ordered subset of *module_list* (preserves original order).
        """
        return [m for m in module_list if (types and m.module_type in types) or (names and m.name in names)]

    def set_selector(
        self,
        selector: ModuleSelector | list | None = None,
        *,
        module_type: str | list[str] | None = None,
        module_name: str | list[str] | None = None,
    ) -> None:
        """Replace the layer selector and re-register hooks.

        Accepts the following forms (combinable):

        - ``set_selector(BlockLevelSelector())`` — explicit selector object
        - ``set_selector(["layer.0", "layer.1"])`` — list of module names / ModuleInfo
        - ``set_selector(module_type="Dinov2Layer")`` — all modules of that type
        - ``set_selector(module_type=["Dinov2Layer", "LayerNorm"])`` — multiple types
        - ``set_selector(module_name="encoder.layer.3")`` — single module by name
        - ``set_selector(module_name=["enc.0", "enc.6"])`` — multiple names
        - ``set_selector(module_type="Dinov2Layer", module_name="layernorm")``
          — union of both filters

        *selector* and (*module_type* / *module_name*) are mutually exclusive.

        Parameters
        ----------
        selector : ModuleSelector, list, or None
            Explicit selector object or list of module names / ModuleInfo objects.
        module_type : str, list[str], or None
            Hook all modules whose ``module_type`` is in this set.
        module_name : str, list[str], or None
            Hook modules whose ``name`` is in this set (exact match).

        Raises
        ------
        ValueError
            If no arguments are supplied, *selector* is combined with filters,
            or the resulting module list is empty.
        """
        using_filters = module_type is not None or module_name is not None
        if selector is None and not using_filters:
            raise ValueError("Provide selector, module_type, or module_name.")
        if selector is not None and using_filters:
            raise ValueError("selector is mutually exclusive with module_type/module_name.")

        if using_filters:
            types = {module_type} if isinstance(module_type, str) else set(module_type or [])
            names = {module_name} if isinstance(module_name, str) else set(module_name or [])
            matched = self._filter_modules(self.module_list, types, names)
            if not matched:
                raise ValueError(
                    f"No modules matched module_type={module_type!r}, "
                    f"module_name={module_name!r}. "
                    f"Available types: {sorted({m.module_type for m in self.module_list})}"
                )
            selector = CustomSelector(matched)
        elif isinstance(selector, list):
            selector = CustomSelector(selector)

        assert selector is not None
        self._selector = selector
        self._bind_selector()
        logger.info("Selector updated | modules={}", len(self._module_names))

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def extract_for_modules(
        self,
        images: dict,
        module_names: list[str],
        batch_size: int,
        show_progress: bool = True,
    ) -> VisualRepresentations:
        """Extract activations for a subset of modules without altering state.

        Temporarily re-registers hooks for *module_names*, runs extraction,
        then restores the original hook configuration.

        Parameters
        ----------
        images : dict
            ``{stim_id: image}`` mapping.
        module_names : list[str]
            Subset of module names to extract.
        batch_size : int
            Images per forward pass.
        show_progress : bool
            Show tqdm progress bar.

        Returns
        -------
        VisualRepresentations
        """
        with self._hooked_for(module_names):
            return self._extract_batch(images, batch_size=batch_size, show_progress=show_progress)

    @contextlib.contextmanager
    def _hooked_for(self, module_names: list[str]):
        """Temporarily re-register hooks for *module_names*, then restore."""
        original = self._module_names
        self._backend.register_hooks(module_names)
        self._module_names = module_names
        try:
            yield
        finally:
            self._backend.register_hooks(original)
            self._module_names = original

    def _build_vr_list(
        self,
        stim_ids: list,
        features: dict[str, np.ndarray],
    ) -> list[VisualRepresentation]:
        """Assemble VisualRepresentation objects from batched features.

        Pairs each layer's activation array with *stim_ids*, using the
        current model metadata and module-type map.  The caller guarantees
        that rows in every array are aligned to *stim_ids* in the same order.

        Parameters
        ----------
        stim_ids : list
            Ordered stimulus IDs produced by :meth:`_prepare_images`.
        features : dict[str, np.ndarray]
            Layer name → activation array of shape ``(n_stim, ...)``.

        Returns
        -------
        list[VisualRepresentation]
        """
        model_meta = self._backend.get_model_meta()
        return [
            VisualRepresentation(
                model=model_meta.model_id,
                module_name=layer,
                module_type=self._module_type_map.get(layer, ""),
                stim_ids=stim_ids,
                array=arr,
            )
            for layer, arr in features.items()
        ]

    def _extract_batch(self, images: dict, batch_size: int, show_progress: bool = False) -> VisualRepresentations:
        stim_ids, loaded = self._prepare_images(images)
        features = self._run_batches(loaded, batch_size, show_progress)
        vr_list = self._build_vr_list(stim_ids, features)
        logger.info(
            "Extracted | n={} | batch_size={} | modules={}",
            len(stim_ids),
            batch_size,
            len(vr_list),
        )
        return VisualRepresentations(vr_list)

    def _run_batches(
        self,
        loaded: list,
        batch_size: int,
        show_progress: bool = False,
    ) -> dict[str, np.ndarray]:
        """Run batched forward passes and concatenate activations per layer.

        Parameters
        ----------
        loaded : list
            Pre-loaded images (PIL / ndarray / Tensor).
        batch_size : int
            Images per forward pass.
        show_progress : bool
            Show tqdm progress bar.

        Returns
        -------
        dict[str, np.ndarray]
            Layer name → concatenated activation array of shape ``(n_images, ...)``.
        """
        all_features: dict[str, list[np.ndarray]] = {}
        n_chunks = (len(loaded) + batch_size - 1) // batch_size
        iterator = tqdm(
            range(0, len(loaded), batch_size),
            total=n_chunks,
            desc="VisionModel",
            unit="batch",
            disable=not show_progress or n_chunks <= 1,
        )
        for start in iterator:
            chunk = loaded[start : start + batch_size]
            chunk_feats = self._forward_chunk(chunk)
            for layer, arr in chunk_feats.items():
                all_features.setdefault(layer, []).append(arr)
        return {layer: np.concatenate(arrs, axis=0) for layer, arrs in all_features.items()}

    @staticmethod
    def _prepare_images(images: dict) -> tuple[list[Any], list[Any]]:
        """Resolve image sources and return (stim_ids, loaded_images).

        Handles string and Path entries by opening them with PIL.
        Already-loaded images (PIL, ndarray, Tensor) pass through unchanged.

        Parameters
        ----------
        images : dict
            ``{stim_id: image_or_path}`` mapping.

        Returns
        -------
        tuple[list, list]
            ``(stim_ids, loaded_images)`` in the same order as *images*.
        """
        from pathlib import Path

        from PIL import Image as PILImage

        stim_ids = list(images.keys())
        loaded: list[Any] = []
        for sid in stim_ids:
            img = images[sid]
            if isinstance(img, (str, Path)):
                img = PILImage.open(img).convert("RGB")
            loaded.append(img)
        return stim_ids, loaded

    def _forward_chunk(self, images: list[Any]) -> dict[str, np.ndarray]:
        """Run one forward pass and return batched activations per layer.

        Parameters
        ----------
        images : list
            Batch of images (PIL / ndarray / str / Path).

        Returns
        -------
        dict[str, np.ndarray]
            Layer name → array of shape ``(B, ...)``.
        """
        inputs = self._backend.preprocess(images)
        with self._backend.collecting() as collect:
            self._backend.forward(inputs)
            activations = collect()
        return {name: act.numpy() for name, act in activations.items()}

    @classmethod
    def register_backend(cls, name: str, backend_cls: type[BaseBackend]) -> None:
        """Register a custom backend class under *name*.

        Parameters
        ----------
        name : str
            Key used in the ``backend=`` argument of :class:`VisionModel`.
        backend_cls : type[BaseBackend]
            Backend class to register.  Must be a concrete subclass of
            :class:`~vneurotk.vision.model.backend.base.BaseBackend`.
        """
        _BACKEND_REGISTRY[name] = backend_cls

    @staticmethod
    def _build_backend(backend: str, device: str) -> BaseBackend:
        if backend not in _BACKEND_REGISTRY:
            raise ValueError(f"Unknown backend {backend!r}. Available: {sorted(_BACKEND_REGISTRY)}")
        entry = _BACKEND_REGISTRY[backend]
        if isinstance(entry, str):
            module_path, class_name = entry.rsplit(":", 1)
            BackendClass: type[BaseBackend] = getattr(importlib.import_module(module_path), class_name)
        else:
            BackendClass = entry
        return BackendClass(device=device)

class ModuleSelector(ABC):
    """Abstract base class for layer selection strategies.

    Subclasses implement :meth:`select`, which receives a list of
    :class:`~vneurotk.vision.meta.ModuleInfo` objects and returns an
    ordered list of module name strings to hook.
    """

    @abstractmethod
    def select(self, modules: list[ModuleInfo]) -> list[str]:
        """Return layer names to hook.

        Parameters
        ----------
        modules : list[ModuleInfo]
            All named modules enumerated by the backend, as returned by
            :meth:`~vneurotk.vision.model.backend.base.BaseBackend.enumerate_modules`.

        Returns
        -------
        list[str]
            Ordered module names to register hooks on.
        """

class BlockLevelSelector(ModuleSelector):
    """Select major block-level modules appropriate for the architecture.

    Uses regex patterns matched against module names.  Architecture
    patterns are tried in order; the first match wins.  Falls back to
    top-level children (depth == 1) if no pattern matches.

    Parameters
    ----------
    max_depth : int
        Maximum nesting depth to include (default 2).  Controls how
        deeply nested sub-blocks are included.
    include_patterns : list[str] or None
        Additional regex patterns to include alongside defaults.
    """

    _ARCH_PATTERNS: list[tuple[str, int]] = [
        (r"^blocks\.\d+$", 2),  # timm ViT
        (r"^encoder\.layers\.\d+$", 3),  # HF ViT (plural)
        (r"^encoder\.layer\.\d+$", 3),  # HF DINOv2 (singular)
        (r"^model\.layer\.\d+$", 3),  # HF DINOv3
        (r"^layer\d+\.\d+$", 3),  # ResNet
        (r"^features\.\d+$", 2),  # VGG / EfficientNet
        (r"^stages\.\d+$", 2),  # ConvNeXt
        (r"^layers\.\d+$", 2),  # Swin / generic
        (r"^vision_model\.encoder\.layers\.\d+$", 4),  # SigLIP / SigLIP2
    ]

    def __init__(
        self,
        max_depth: int = 2,
        include_patterns: list[str] | None = None,
        arch_patterns: list[tuple[str, int]] | None = None,
    ) -> None:
        self.max_depth = max_depth
        self._extra = [re.compile(p) for p in (include_patterns or [])]
        raw = arch_patterns if arch_patterns is not None else self._ARCH_PATTERNS
        self._compiled_patterns: list[tuple[re.Pattern, int]] = [(re.compile(p), d) for p, d in raw]

    @classmethod
    def default_patterns(cls) -> list[tuple[str, int]]:
        """Return a copy of the built-in architecture patterns.

        Returns
        -------
        list[tuple[str, int]]
            Each element is ``(regex_pattern, max_depth)``.
            Mutating the returned list does not affect the class default.
        """
        return list(cls._ARCH_PATTERNS)

    @staticmethod
    def _module_depth(name: str) -> int:
        """Return the nesting depth of a module given its dotted name.

        Parameters
        ----------
        name : str
            Module name as produced by ``model.named_modules()``,
            e.g. ``"encoder.layer.3"`` → depth 3.

        Returns
        -------
        int
            ``name.count(".") + 1``.  Empty string returns ``0``.
        """
        return name.count(".") + 1 if name else 0

    def select(self, modules: list[ModuleInfo]) -> list[str]:
        """Select block-level layers from *modules*.

        Parameters
        ----------
        modules : list[ModuleInfo]

        Returns
        -------
        list[str]
        """
        selected: list[str] = []

        for m in modules:
            matched = any(pat.match(m.name) and m.depth <= max_d for pat, max_d in self._compiled_patterns)
            if not matched:
                matched = any(pat.search(m.name) for pat in self._extra)

            if matched:
                selected.append(m.name)

        if not selected:
            selected = [m.name for m in modules if m.depth == 1]

        return selected

class AllLeafSelector(ModuleSelector):
    """Select all leaf modules (modules with no children).

    Parameters
    ----------
    exclude_types : tuple[type, ...] or None
        Module types to skip.  Defaults to activation and regularization
        layers that carry no representational content.
    """

    _DEFAULT_EXCLUDE = (
        nn.Dropout,
        nn.Identity,
        nn.ReLU,
        nn.GELU,
        nn.SiLU,
        nn.Sigmoid,
        nn.Softmax,
    )

    def __init__(self, exclude_types: tuple | None = None) -> None:
        raw = exclude_types if exclude_types is not None else self._DEFAULT_EXCLUDE
        self._exclude_names: frozenset[str] = frozenset(t.__name__ for t in raw)

    def select(self, modules: list[ModuleInfo]) -> list[str]:
        """Return names of all non-excluded leaf modules.

        Parameters
        ----------
        modules : list[ModuleInfo]

        Returns
        -------
        list[str]
        """
        return [m.name for m in modules if m.is_leaf and m.module_type not in self._exclude_names]

class CustomSelector(ModuleSelector):
    """Use an explicit user-supplied list of layer names.

    Parameters
    ----------
    layer_names : list[str] or list[ModuleInfo]
        Exact module names as they appear in ``model.named_modules()``,
        or :class:`~vneurotk.vision.meta.ModuleInfo`
        objects as returned by :attr:`VisionModel.module_list`.

    Raises
    ------
    ValueError
        During :meth:`select` if any name is not found in the module list.
    """

    def __init__(self, layer_names: list) -> None:
        self.layer_names = [item.name if hasattr(item, "name") else item for item in layer_names]

    def select(self, modules: list[ModuleInfo]) -> list[str]:
        """Validate and return the configured layer names.

        Parameters
        ----------
        modules : list[ModuleInfo]

        Returns
        -------
        list[str]

        Raises
        ------
        ValueError
            If any layer name is absent from *modules*.
        """
        available = {m.name for m in modules}
        missing = [n for n in self.layer_names if n not in available]
        if missing:
            raise ValueError(
                f"Layer(s) not found in model: {missing}. Inspect available names with model.named_modules()."
            )
        return list(self.layer_names)