"""Periodic-table / isotope registry with charge, spin, electron-count metadata (decision 42).

Data source: ase.data (atomic numbers, standard atomic masses, covalent radii; ASE is LGPL-2.1+).
Common oxidation states are a small curated table for screening only. Isotopes beyond the
standard-mass entry are not tabulated here: requests for specific isotopes are recorded as
metadata and flagged (`isotope_data: not_tabulated`) rather than given invented masses.
"""
from __future__ import annotations

from dataclasses import dataclass, field

from ase.data import atomic_masses, atomic_numbers, chemical_symbols, covalent_radii

# Valence electrons (main-group counting; transition metals: group number up to 12)
_MAIN_VALENCE = {
    "H": 1, "He": 2, "Li": 1, "Be": 2, "B": 3, "C": 4, "N": 5, "O": 6, "F": 7, "Ne": 8,
    "Na": 1, "Mg": 2, "Al": 3, "Si": 4, "P": 5, "S": 6, "Cl": 7, "Ar": 8,
    "K": 1, "Ca": 2, "Ga": 3, "Ge": 4, "As": 5, "Se": 6, "Br": 7, "Kr": 8,
    "Sc": 3, "Ti": 4, "V": 5, "Cr": 6, "Mn": 7, "Fe": 8, "Co": 9, "Ni": 10, "Cu": 11, "Zn": 12,
    "Ag": 11, "Au": 11, "Pd": 10, "Pt": 10, "Rb": 1, "Sr": 2, "In": 3, "Sn": 4, "Sb": 5, "Te": 6, "I": 7, "Xe": 8,
}
COMMON_OXIDATION_STATES = {
    "H": [-1, 1], "Li": [1], "Be": [2], "B": [3], "C": [-4, 2, 4], "N": [-3, 3, 5], "O": [-2], "F": [-1],
    "Na": [1], "Mg": [2], "Al": [3], "Si": [-4, 4], "P": [-3, 3, 5], "S": [-2, 2, 4, 6], "Cl": [-1, 1, 5, 7],
    "K": [1], "Ca": [2], "Ti": [2, 3, 4], "V": [2, 3, 4, 5], "Cr": [2, 3, 6], "Mn": [2, 3, 4, 7],
    "Fe": [2, 3], "Co": [2, 3], "Ni": [2], "Cu": [1, 2], "Zn": [2], "Ag": [1], "Au": [1, 3], "Pd": [2, 4], "Pt": [2, 4],
}
NOBLE_GASES = {"He", "Ne", "Ar", "Kr", "Xe", "Rn"}


@dataclass(frozen=True)
class ElementInfo:
    symbol: str
    Z: int
    mass_amu: float
    covalent_radius_ang: float
    valence_electrons: int | None
    common_oxidation_states: list[int] = field(default_factory=list)
    flags: list[str] = field(default_factory=list)


def element(symbol: str) -> ElementInfo:
    if symbol not in atomic_numbers or symbol == "X":
        raise ValueError(f"unknown element symbol: {symbol}")
    z = atomic_numbers[symbol]
    flags = []
    if z > 86:
        flags.append("heavy_element_relativistic_effects_not_handled_by_default_solvers")
    if z > 103:
        flags.append("superheavy_regime_unsupported")
    val = _MAIN_VALENCE.get(symbol)
    if val is None:
        flags.append("valence_count_not_tabulated")
    return ElementInfo(
        symbol=symbol,
        Z=z,
        mass_amu=float(atomic_masses[z]),
        covalent_radius_ang=float(covalent_radii[z]),
        valence_electrons=val,
        common_oxidation_states=list(COMMON_OXIDATION_STATES.get(symbol, [])),
        flags=flags,
    )


def all_symbols() -> list[str]:
    return [s for s in chemical_symbols[1:]]


def electron_count(symbols: list[str], charge: int = 0) -> int:
    return sum(atomic_numbers[s] for s in symbols) - int(charge)


def spin_parity_ok(n_electrons: int, multiplicity: int) -> bool:
    """2S+1 must have parity consistent with electron count: even N -> odd multiplicity."""
    return (n_electrons % 2) == ((multiplicity - 1) % 2)


def isotope_metadata(symbol: str, mass_number: int | None) -> dict:
    info = element(symbol)
    if mass_number is None:
        return {"symbol": symbol, "mass_amu": info.mass_amu, "isotope_data": "standard_atomic_weight"}
    return {"symbol": symbol, "mass_number": mass_number, "mass_amu": None, "isotope_data": "not_tabulated"}