"""Name-based views into an atom-axis property."""

from __future__ import annotations

from operator import attrgetter
from typing import Any

import numpy as np

from .esmfold2_protein_structure import index_by_atom_name


class AtomIndexer:
    """Select named atoms from one property of a structure-like object.

    The wrapper intentionally remains small because ``ProteinChain.atom37`` and
    related public properties expose it directly.
    """

    __slots__ = ("_get_property", "dim", "property", "structure")

    def __init__(self, structure: Any, property: str, dim: int):
        self.structure = structure
        self.property = property
        self.dim = dim
        self._get_property = attrgetter(property)

    def __getitem__(self, atom_names: str | list[str]) -> np.ndarray:
        values = self._get_property(self.structure)
        return index_by_atom_name(values, atom_names, dim=self.dim)