Source code for httk.atomistic.models.structure.plain

"""
Backend wrapping an spglib-like (lattice, positions, numbers) triple.
"""

from typing import Any, Self

from httk.atomistic.elements import symbol_of
from httk.atomistic.models._vector_guards import is_basis_3x3, is_coords_nx3, try_surdvector
from httk.atomistic.models.cell.cell import Cell
from httk.atomistic.models.sites.sites import Sites
from httk.atomistic.models.species.species import Species
from httk.atomistic.models.structure.backend import StructureBackend


def _is_number(value: Any) -> bool:
    return isinstance(value, (int, float)) and not isinstance(value, bool)


def _is_primitive_triple(obj: Any) -> bool:
    if not isinstance(obj, (list, tuple)) or len(obj) != 3:
        return False
    lattice, positions, numbers = obj
    if not is_basis_3x3(lattice) or not is_coords_nx3(positions):
        return False
    if not isinstance(numbers, (list, tuple)) or not all(_is_number(z) for z in numbers):
        return False
    positions_vector = try_surdvector(positions)
    nsites = positions_vector.dim[0] if positions_vector is not None and len(positions_vector.dim) == 2 else 0
    return len(numbers) == nsites




class PlainStructure(StructureBackend):
    """Represent a crystal structure backed by an spglib-like triple.

    The native representation is a length-3 ``(lattice, positions, numbers)`` list
    or tuple, where ``lattice`` is 3x3, ``positions`` is Nx3 reduced coordinates,
    and ``numbers`` is the length-N sequence of atomic numbers. The quartet is
    derived lazily and cached: ``cell`` is a ``Cell``, ``sites`` a ``Sites``,
    ``species`` one single-element ``Species`` per distinct atomic number, and
    ``unwrap`` returns the original triple.
    :param obj: The primitive structure triple to wrap.
    :param \\*\\*hints: Backend-selection hints.
    """

    _raw: Any
    _lattice: Any
    _positions: Any
    _numbers: tuple[int, ...]
    _cell_cache: Cell | None
    _sites_cache: Sites | None
    _species_cache: tuple[Species, ...] | None
    _species_at_sites_cache: tuple[str, ...] | None

    @classmethod
    def _backend_adopt(cls, obj: Any, **hints: Any) -> Self | None:
        r"""Adopt a primitive structure triple.

        :param obj: The source object to adopt.
        :param \**hints: Backend-selection hints.
        :return: An initialized backend, or ``None`` when this backend declines ``obj``.
        """
        if hints and hints.get("kind", "plain") != "plain":
            return None
        if not _is_primitive_triple(obj):
            return None
        return cls(obj, **hints)

    def __init__(self, obj: Any, **hints: Any) -> None:
        lattice, positions, numbers = obj
        self._raw = obj
        self._lattice = lattice
        self._positions = positions
        self._numbers = tuple(int(z) for z in numbers)
        self._cell_cache = None
        self._sites_cache = None
        self._species_cache = None
        self._species_at_sites_cache = None

    @property


    def cell(self) -> Cell:
        """Expose the cell derived from the lattice."""
        if self._cell_cache is None:
            self._cell_cache = Cell(self._lattice)
        return self._cell_cache


    @property


    def sites(self) -> Sites:
        """Expose the reduced coordinates derived from the positions."""
        if self._sites_cache is None:
            self._sites_cache = Sites(self._positions)
        return self._sites_cache


    @property


    def species(self) -> tuple[Species, ...]:
        """Expose one species definition for each distinct atomic number."""
        if self._species_cache is None:
            distinct: list[int] = []
            seen: set[int] = set()
            for z in self._numbers:
                if z not in seen:
                    seen.add(z)
                    distinct.append(z)
            self._species_cache = tuple(
                Species(name=symbol_of(z), chemical_symbols=(symbol_of(z),), concentration=(1.0,)) for z in distinct
            )
        return self._species_cache


    @property


    def species_at_sites(self) -> tuple[str, ...]:
        """Expose the species name occupying each site."""
        if self._species_at_sites_cache is None:
            self._species_at_sites_cache = tuple(symbol_of(z) for z in self._numbers)
        return self._species_at_sites_cache




    def unwrap(self) -> Any:
        """Return the original primitive structure triple.

        :return: The wrapped lattice, positions, and atomic numbers.
        """
        return self._raw