periodic_table

"""Build a periodic-table toy dataset from PubChem's machine-readable table."""

from pathlib import Path

import pandas as pd


SOURCE = "https://pubchem.ncbi.nlm.nih.gov/rest/pug/periodictable/CSV"
OUT = Path(__file__).resolve().parent.parent / "csv" / "periodic_table.csv"


def period(atomic_number: int) -> int:
    for upper, value in ((2, 1), (10, 2), (18, 3), (36, 4), (54, 5), (86, 6), (118, 7)):
        if atomic_number <= upper:
            return value
    raise ValueError(f"Unsupported atomic number: {atomic_number}")


def group(atomic_number: int) -> int | None:
    rows = {
        1: [1, 18],
        2: [1, 2, 13, 14, 15, 16, 17, 18],
        3: [1, 2, 13, 14, 15, 16, 17, 18],
        4: list(range(1, 19)),
        5: list(range(1, 19)),
        6: [1, 2, 3, *range(4, 19)],
        7: [1, 2, 3, *range(4, 19)],
    }
    starts = {1: 1, 2: 3, 3: 11, 4: 19, 5: 37, 6: 55, 7: 87}
    p = period(atomic_number)
    offset = atomic_number - starts[p]
    if p in (6, 7) and offset > 2:
        # Ce–Lu and Th–Lr are drawn in the detached f-block, without assigning
        # the disputed group-3 position to every member.
        if offset <= 16:
            return None
        offset -= 14
    return rows[p][offset]


def layout(atomic_number: int) -> tuple[int, int]:
    if 58 <= atomic_number <= 71:
        return 8, atomic_number - 54
    if 90 <= atomic_number <= 103:
        return 9, atomic_number - 86
    return period(atomic_number), group(atomic_number) or 3


def main() -> None:
    data = pd.read_csv(SOURCE)
    if data["AtomicNumber"].tolist() != list(range(1, 119)):
        raise ValueError("PubChem periodic table is not the expected 1–118 sequence")

    data = data.rename(
        columns={
            "AtomicNumber": "atomic_number",
            "Symbol": "symbol",
            "Name": "name",
            "AtomicMass": "atomic_mass",
            "CPKHexColor": "cpk_hex_color",
            "ElectronConfiguration": "electron_configuration",
            "Electronegativity": "electronegativity",
            "AtomicRadius": "atomic_radius_pm",
            "IonizationEnergy": "ionization_energy_ev",
            "ElectronAffinity": "electron_affinity_ev",
            "OxidationStates": "oxidation_states",
            "StandardState": "standard_state",
            "MeltingPoint": "melting_point_k",
            "BoilingPoint": "boiling_point_k",
            "Density": "density_g_cm3",
            "GroupBlock": "category",
            "YearDiscovered": "year_discovered",
        }
    )
    data.insert(3, "period", data["atomic_number"].map(period))
    data.insert(4, "group", data["atomic_number"].map(group).astype("Int64"))
    positions = data["atomic_number"].map(layout)
    data.insert(5, "layout_row", positions.map(lambda value: value[0]))
    data.insert(6, "layout_column", positions.map(lambda value: value[1]))

    OUT.parent.mkdir(parents=True, exist_ok=True)
    data.to_csv(OUT, index=False)
    print(f"Done: {len(data)} elements saved to {OUT}")


if __name__ == "__main__":
    main()