Source code for molfoundry.atomid

from typing import Dict, List, Optional

from .core import TSObject, coerce_type
from .exceptions import LogicError

_class_ref = TSObject.get_class_ref("AtomId")


# noinspection PyPep8Naming


class AtomId(TSObject):
    """Identifies a chemical element and exposes its data from the shared ``Atom``
    table.

    Construct one by symbol or by atomic number::

        carbon = AtomId.bySymbol("C")
        carbon = AtomId.byId(6)
        print(carbon.protons, carbon.valenceElectrons, carbon.orbitalNames)

    ``AtomId`` is also what ``Vertex.atomId`` returns. An AtomId that comes from a
    non-atom vertex (e.g. an electron), an unknown symbol, or an unknown atomic
    number is *invalid*: ``valid`` is ``False``, ``id`` is available, and every
    element property raises ``LogicError``.
    """

    def __init__(self) -> None:
        # A bare AtomId is invalid; use bySymbol / byId (or Vertex.atomId).
        super().__init__(_class_ref)



    @staticmethod
    def bySymbol(symbol: str) -> "AtomId":
        """Look up an element by symbol (e.g. ``"C"``). Invalid if unknown."""
        return AtomId.wrap(TSObject.call_static(_class_ref, "fromSymbol", symbol))


    # noinspection PyShadowingBuiltins


    @staticmethod
    def byId(id: int) -> "AtomId":
        """Look up an element by atomic number / id (e.g. ``6`` -> carbon).
        Invalid if no element has that atomic number."""
        return AtomId.wrap(TSObject.call_static(_class_ref, "fromId", id))


    def __int__(self) -> int:
        return int(self._ref.getId())

    def __str__(self) -> str:
        # Matches MØD's ``operator<<``: the atomic number, not the symbol.
        return str(int(self))

    def __eq__(self, other: object) -> bool:
        # MØD compares AtomId to AtomId only (by id); not to a bare int.
        if isinstance(other, AtomId):
            return int(self) == int(other)
        return NotImplemented

    def __hash__(self) -> int:
        # Defining ``__eq__`` otherwise makes the type unhashable; hash on the id
        # so equal AtomIds hash alike.
        return hash(int(self))

    def __repr__(self) -> str:
        return f"AtomId({self._ref.getSymbol()!r})" if self.valid else "AtomId(invalid)"

    def _require_valid(self) -> None:
        if not self.valid:
            raise LogicError("Invalid AtomId has no element data")

    @property
    def valid(self) -> bool:
        return bool(self._ref.isValid())

    @property
    def id(self) -> int:
        """The atomic number (proton count); 0 for an invalid AtomId."""
        return int(self._ref.getId())

    @property
    def symbol(self) -> str:
        self._require_valid()
        return self._ref.getSymbol()

    @property
    def protons(self) -> int:
        self._require_valid()
        return int(self._ref.getProtons())

    @property
    def electrons(self) -> int:
        self._require_valid()
        return int(self._ref.getElectrons())

    @property
    def standardAtomicWeight(self) -> Optional[float]:
        self._require_valid()
        value = coerce_type(self._ref.getStandardAtomicWeight())
        return float(value) if value is not None else None

    @property
    def mostAbundantIsotope(self) -> Optional[int]:
        self._require_valid()
        value = coerce_type(self._ref.getMostAbundantIsotope())
        return int(value) if value is not None else None

    @property
    def principalQuantumNumber(self) -> int:
        self._require_valid()
        return int(self._ref.getPrincipalQuantumNumber())

    @property
    def valenceElectrons(self) -> int:
        self._require_valid()
        return int(self._ref.getValenceElectrons())

    @property
    def orbitals(self) -> List[List[int]]:
        """Electron occupancy per orbital, in fill order: one inner list per
        occupied orbital, each entry a 1 (single electron) or 2 (pair)."""
        self._require_valid()
        return [[int(x) for x in orbital] for orbital in self._ref.getOrbitals()]

    @property
    def orbitalNames(self) -> List[str]:
        """Orbital names with occupancy, e.g. ``["1s2", "2s2", "2p2"]``."""
        self._require_valid()
        return [str(name) for name in self._ref.getOrbitalNames()]

    @property
    def isotopicMasses(self) -> Dict[int, float]:
        """Mapping of mass number to isotopic mass. Empty when untabulated."""
        self._require_valid()
        return {int(pair[0]): float(pair[1]) for pair in self._ref.getIsotopicMasses()}