Source code for vibesketch.chem.formula

"""Chemical formulas of solids that have no molecular graph (#142).

Oxides, salts, catalysts and scales are written as formulas, not drawn as
structures: Fe2O3, Pb3(CO3)2(OH)2, CuSO4·5H2O, a variable or doped
composition such as Fe3−xCrxO4 or CoAlxOy, a composite such as AlOx–Fe2O3,
with an optional polymorph (γ-Al2O3) and state ((s)). This module reads such
a formula into its composition, the canonical form that identifies it, and,
when every count is a number, its molar mass with the atomic weights RDKit
gives the species report.

Grammar, informally: ``[polymorph-]phase[–phase…][(state)]``, a phase being
groups and elements each with an optional count, hydrate parts after ``·``
(or ``.``/``*``) with a leading multiplier, and a count a number, a variable
(x, y, z or δ), or a number with a variable added or taken away (3−x).
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field

from vibesketch.elements import ATOMIC_NUMBER

__all__ = [
    "Formula",
    "FormulaError",
    "is_solid_text",
    "parse_formula",
]

#: Polymorph prefixes: Greek letters and their names.
_POLYMORPHS = {
    "α": "α", "β": "β", "γ": "γ", "δ": "δ", "η": "η", "θ": "θ", "κ": "κ", "χ": "χ",
    "alpha": "α", "beta": "β", "gamma": "γ", "delta": "δ", "eta": "η", "theta": "θ",
    "kappa": "κ", "chi": "χ",
}  # fmt: skip
_STATES = ("s", "l", "g", "aq", "ads")
_MINUS = "−–\\-"
_COUNT = re.compile(
    rf"\d+(?:\.\d+)?(?:[{_MINUS}+]\d*(?:\.\d+)?[xyzδ])?|\d*(?:\.\d+)?[xyzδ](?:[{_MINUS}+]\d+(?:\.\d+)?)?"
)
_HYDRATE = "·.*•"
#: Composite phases are joined by an en dash, a hyphen between phases, or a slash.
_COMPOSITE = re.compile(r"\s*[–/]\s*")




class FormulaError(ValueError):
    """Text that is not a formula this module reads; the message says where."""



Count = float | str


def _count(text: str | None) -> Count:
    if not text:
        return 1.0
    text = text.replace("−", "-").replace("–", "-")
    try:
        return float(text)
    except ValueError:
        return text




@dataclass(frozen=True)
class Formula:
    """A formula read: its phases' compositions, polymorph and state."""

    text: str
    phases: tuple[dict[str, Count], ...]
    polymorph: str | None = None
    state: str | None = None
    #: Each phase as written, for display.
    written: tuple[str, ...] = field(default=())



    @property
    def numeric(self) -> bool:
        return all(isinstance(c, float) for phase in self.phases for c in phase.values())




    @property
    def composition(self) -> dict[str, Count]:
        """The elements of all phases together; a variable count keeps its expression."""
        total: dict[str, Count] = {}
        for phase in self.phases:
            for element, count in phase.items():
                before = total.get(element)
                if before is None:
                    total[element] = count
                elif isinstance(before, float) and isinstance(count, float):
                    total[element] = before + count
                else:
                    total[element] = f"{_show(before)}+{_show(count)}"
        return total




    def canonical(self) -> str:
        """Each phase in Hill order, phases joined by a dash: what identifies the solid.

        A variable count is kept in parentheses, Fe3−xCrxO4 as Cr(x)Fe(3-x)O4.
        """
        return "–".join(_hill(phase) for phase in self.phases)




    def key(self) -> str:
        """The canonical formula with polymorph and state: γ-Al2O3(s)."""
        prefix = f"{self.polymorph}-" if self.polymorph else ""
        suffix = f"({self.state})" if self.state else ""
        return f"{prefix}{self.canonical()}{suffix}"




    def molar_mass(self) -> float | None:
        """g/mol from RDKit's atomic weights, or ``None`` with a variable count."""
        if not self.numeric:
            return None
        from rdkit import Chem

        table = Chem.GetPeriodicTable()
        return sum(
            float(table.GetAtomicWeight(element)) * float(count)
            for element, count in self.composition.items()
        )




    def monoisotopic_mass(self) -> float | None:
        """g/mol from each element's most common isotope, or ``None`` with a variable count."""
        if not self.numeric:
            return None
        from rdkit import Chem

        table = Chem.GetPeriodicTable()
        return sum(
            float(table.GetMostCommonIsotopeMass(element)) * float(count)
            for element, count in self.composition.items()
        )




    def counts(self) -> dict[str, float] | None:
        """Atoms per formula unit, or ``None`` with a variable count."""
        if not self.numeric:
            return None
        return {element: float(count) for element, count in self.composition.items()}




def _show(count: Count) -> str:
    if isinstance(count, float):
        return str(int(count)) if count == int(count) else f"{count:g}"
    return count


def _hill(phase: dict[str, Count]) -> str:
    present = set(phase)
    if "C" in present:
        order = ["C", *(["H"] if "H" in present else []), *sorted(present - {"C", "H"})]
    else:
        order = sorted(present)
    parts = []
    for element in order:
        count = phase[element]
        if isinstance(count, float):
            parts.append(element if count == 1 else f"{element}{_show(count)}")
        else:
            parts.append(f"{element}({count})")
    return "".join(parts)


def _add(target: dict[str, Count], element: str, count: Count, times: Count = 1.0) -> None:
    if isinstance(count, float) and isinstance(times, float):
        value: Count = count * times
    elif isinstance(times, float) and times == 1:
        value = count
    else:
        value = f"({_show(times)})*({_show(count)})"
    before = target.get(element)
    if before is None:
        target[element] = value
    elif isinstance(before, float) and isinstance(value, float):
        target[element] = before + value
    else:
        target[element] = f"{_show(before)}+{_show(value)}"


def _group(text: str, start: int, closing: str | None) -> tuple[dict[str, Count], int]:
    """Elements and groups from ``start`` up to ``closing`` (or the end)."""
    found: dict[str, Count] = {}
    index = start
    while index < len(text):
        char = text[index]
        if closing is not None and char == closing:
            return found, index + 1
        if char in "([":
            inner, index = _group(text, index + 1, ")" if char == "(" else "]")
            match = _COUNT.match(text, index)
            times = _count(match.group(0) if match else None)
            index = match.end() if match else index
            for element, count in inner.items():
                _add(found, element, count, times)
            continue
        symbol = None
        if char.isupper():
            two = text[index : index + 2]
            if len(two) == 2 and two[1].islower() and two in ATOMIC_NUMBER:
                symbol = two
            elif char in ATOMIC_NUMBER:
                symbol = char
        if symbol is None:
            raise FormulaError(f"{text!r}: no element at {text[index:]!r}")
        index += len(symbol)
        match = _COUNT.match(text, index)
        _add(found, symbol, _count(match.group(0) if match else None))
        index = match.end() if match else index
    if closing is not None:
        raise FormulaError(f"{text!r}: a bracket is not closed")
    return found, index


def _phase(text: str) -> dict[str, Count]:
    parts = re.split(rf"\s*[{_HYDRATE}]\s*", text)
    total: dict[str, Count] = {}
    for position, part in enumerate(parts):
        if not part:
            raise FormulaError(f"{text!r}: an empty part")
        times: Count = 1.0
        if position:
            match = re.match(r"\d+(?:\.\d+)?|[xyzn]", part)
            if match:
                times = _count(match.group(0))
                part = part[match.end() :]
        found, _ = _group(part, 0, None)
        for element, count in found.items():
            _add(total, element, count, times)
    if not total:
        raise FormulaError(f"{text!r}: no elements")
    return total




def parse_formula(text: str) -> Formula:
    """Read a formula; raises :class:`FormulaError` for anything else."""
    body = text.strip()
    if not body or len(body) > 120:
        raise FormulaError("a formula is 1 to 120 characters")
    state = None
    found = re.search(r"\(\s*(s|l|g|aq|ads)\s*\)\s*$", body)
    if found:
        state = found.group(1)
        body = body[: found.start()].rstrip()
    polymorph = None
    prefix = re.match(r"^([A-Za-zα-ω]+)\s*[-‐]\s*", body)
    if prefix and prefix.group(1) in _POLYMORPHS:
        polymorph = _POLYMORPHS[prefix.group(1)]
        body = body[prefix.end() :]
    written = tuple(p for p in _COMPOSITE.split(body) if p)
    if not written:
        raise FormulaError(f"{text!r}: no formula")
    return Formula(
        text=text.strip(),
        phases=tuple(_phase(p.replace(" ", "")) for p in written),
        polymorph=polymorph,
        state=state,
        written=written,
    )





def is_solid_text(text: str) -> Formula | None:
    """The formula a canvas text writes for a solid, or ``None``.

    Only a whole text that is a formula marked as a solid counts: with the
    state (s) or a polymorph, so that conditions such as NaOH or H2O at an
    arrow keep their meaning.
    """
    if "\n" in text:
        return None
    try:
        formula = parse_formula(text)
    except FormulaError:
        return None
    if formula.state == "s" or formula.polymorph is not None:
        return formula
    return None