app.py

import math
from typing import Dict, List, Optional

import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from periodictable import elements

# -----------------------------
# Helpers
# -----------------------------
def to_float(x):
    """Coerce periodictable numeric (incl. uncertainties) to plain float, else NaN."""
    if x is None:
        return np.nan
    try:
        # uncertainties.UFloat has .nominal_value
        v = getattr(x, "nominal_value", x)
        return float(v)
    except Exception:
        try:
            return float(x)
        except Exception:
            return np.nan

# -----------------------------
# Data
# -----------------------------
NUMERIC_PROPS = [
    ("mass", "Atomic mass (u)"),
    ("density", "Density (g/cm^3)"),
    ("electronegativity", "Pauling electronegativity"),
    ("boiling_point", "Boiling point (K)"),
    ("melting_point", "Melting point (K)"),
    ("vdw_radius", "van der Waals radius (pm)"),
    ("covalent_radius", "Covalent radius (pm)"),
]

CURATED_FACTS: Dict[str, List[str]] = {
    "H": ["Lightest element; ~74% of the visible universe by mass is hydrogen in stars."],
    "He": ["Inert, used in cryogenics and balloons; second lightest element."],
    "Li": ["Batteries MVP: lithium-ion cells power phones and EVs."],
    "C": ["Backbone of life; diamond and graphite are pure carbon with wildly different properties."],
    "N": ["~78% of Earth's atmosphere is nitrogen (mostly N₂)."],
    "O": ["Essential for respiration; ~21% of Earth's atmosphere."],
    "Na": ["Sodium metal reacts violently with water—handle only under oil or inert gas."],
    "Mg": ["Burns with a bright white flame; used in flares and fireworks."],
    "Al": ["Light and strong; forms a protective oxide layer that resists corrosion."],
    "Si": ["Silicon is the basis of modern electronics—hello, semiconductors."],
    "Cl": ["Powerful disinfectant; elemental chlorine is toxic, compounds are widely useful."],
    "Ar": ["Argon is used to provide inert atmospheres for welding and 3D printing."],
    "Fe": ["Core of steel; iron is essential in hemoglobin for oxygen transport."],
    "Cu": ["Excellent electrical conductor; iconic blue-green patina (verdigris)."],
    "Ag": ["Highest electrical conductivity of all metals; historically used as currency."],
    "Au": ["Very unreactive ('noble'); prized for electronics and jewelry."],
    "Hg": ["Only metal that's liquid at room temperature; toxic—use with care."],
    "Pb": ["Dense and malleable; toxicity led to phase-out from gasoline and paints."],
    "U": ["Radioactive; used as nuclear reactor fuel (U-235)."],
    "Pu": ["Man-made in quantity; key in certain nuclear technologies."],
    "F": ["Most electronegative element; extremely reactive."],
    "Ne": ["Neon glows striking red-orange in discharge tubes—classic signs."],
    "Xe": ["Xenon makes bright camera flashes and high-intensity lamps."],
}

GROUP_FACTS = {
    "alkali": "Alkali metal: very reactive soft metal; forms +1 cations and reacts with water.",
    "alkaline-earth": "Alkaline earth metal: reactive (less than Group 1); forms +2 cations.",
    "transition": "Transition metal: often good catalysts, colorful compounds, multiple oxidation states.",
    "post-transition": "Post-transition metal: softer metals with lower melting points than transition metals.",
    "metalloid": "Metalloid: properties between metals and nonmetals; often semiconductors.",
    "nonmetal": "Nonmetal: tends to form covalent compounds; wide range of roles in biology and materials.",
    "halogen": "Halogen: very reactive nonmetals; form salts with metals and −1 oxidation state.",
    "noble-gas": "Noble gas: chemically inert under most conditions; monatomic gases.",
    "lanthanide": "Lanthanide: f-block rare earths; notable for magnets, lasers, and phosphors.",
    "actinide": "Actinide: radioactive f-block; includes nuclear fuel materials.",
}

def classify_category(el) -> str:
    try:
        if el.block == "s" and el.group == 1 and el.number != 1:
            return "alkali"
        if el.block == "s" and el.group == 2:
            return "alkaline-earth"
        if el.block == "p" and el.group == 17:
            return "halogen"
        if el.block == "p" and el.group == 18:
            return "noble-gas"
        if el.block == "d":
            return "transition"
        if el.block == "f" and 57 <= el.number <= 71:
            return "lanthanide"
        if el.block == "f" and 89 <= el.number <= 103:
            return "actinide"
        if el.block == "p" and not el.metallic:
            return "nonmetal"
        if el.block == "p" and el.metallic:
            return "post-transition"
    except Exception:
        pass
    return "post-transition" if getattr(el, "metallic", False) else "nonmetal"

def build_elements_df() -> pd.DataFrame:
    rows = []
    for Z in range(1, 119):
        el = elements[Z]
        if el is None:
            continue
        data = {
            "Z": el.number,
            "symbol": el.symbol,
            "name": el.name.title(),
            "period": getattr(el, "period", None),
            "group": getattr(el, "group", None),  # may be None for many
            "block": getattr(el, "block", None),
            "mass": to_float(getattr(el, "mass", None)),
            "density": to_float(getattr(el, "density", None)),
            "electronegativity": to_float(getattr(el, "electronegativity", None)),
            "boiling_point": to_float(getattr(el, "boiling_point", None)),
            "melting_point": to_float(getattr(el, "melting_point", None)),
            "vdw_radius": to_float(getattr(el, "vdw_radius", None)),
            "covalent_radius": to_float(getattr(el, "covalent_radius", None)),
            "category": classify_category(el),
            "is_radioactive": bool(getattr(el, "radioactive", False)),
        }
        rows.append(data)
    return pd.DataFrame(rows).sort_values("Z").reset_index(drop=True)

DF = build_elements_df()

# -----------------------------
# Build a robust grid (no reliance on group from the lib)
# Rules: s->groups 1-2, d->3..12, p->13..18; period 1 special (H at 1, He at 18)
# f-block shown separately.
# -----------------------------
MAX_GROUP = 18
MAX_PERIOD = 7
GRID: List[List[Optional[int]]] = [[None for _ in range(MAX_GROUP)] for _ in range(MAX_PERIOD)]

for period in range(1, MAX_PERIOD + 1):
    rows = DF[DF["period"] == period].sort_values("Z")
    s = rows[rows["block"] == "s"]["Z"].tolist()
    d = rows[rows["block"] == "d"]["Z"].tolist()
    p = rows[rows["block"] == "p"]["Z"].tolist()
    # Period 1 special case
    if period == 1:
        # Expect H then He
        if len(s) >= 1:
            GRID[0][0] = int(s[0])  # H -> group 1
        if len(p) >= 1:
            GRID[0][17] = int(p[-1])  # He -> group 18
        continue
    # s-block (usually 2)
    if len(s) >= 1:
        GRID[period - 1][0] = int(s[0])     # group 1
    if len(s) >= 2:
        GRID[period - 1][1] = int(s[1])     # group 2
    # d-block (10 wide), only in periods >= 4
    for i, z in enumerate(d):
        if i < 10:
            GRID[period - 1][2 + i] = int(z)  # groups 3..12
    # p-block (6 wide)
    for i, z in enumerate(p[-6:]):  # last 6 p-block in order
        GRID[period - 1][12 + i] = int(z)     # groups 13..18

# f-block lists (lanthanides/actinides)
LAN = [int(z) for z in DF["Z"] if 57 <= int(z) <= 71]
ACT = [int(z) for z in DF["Z"] if 89 <= int(z) <= 103]

# -----------------------------
# Plotting (Matplotlib -> gr.Plot)
# -----------------------------
def plot_trend(trend_df: pd.DataFrame, prop_key: str, Z: int, symbol: str):
    fig, ax = plt.subplots()
    ax.scatter(trend_df["Z"], trend_df[prop_key])
    sel = trend_df.loc[trend_df["Z"] == Z, prop_key]
    if not sel.empty and not pd.isna(sel.values[0]):
        ax.scatter([Z], [sel.values[0]], s=80)
        ax.text(Z, sel.values[0], symbol, ha="center", va="bottom")
    ax.set_xlabel("Atomic number (Z)")
    ax.set_ylabel(dict(NUMERIC_PROPS)[prop_key])
    ax.set_title(f"{dict(NUMERIC_PROPS)[prop_key]} across the periodic table")
    fig.tight_layout()
    return fig

def plot_heatmap(property_key: str):
    prop_label = dict(NUMERIC_PROPS)[property_key]
    grid_vals = np.full((MAX_PERIOD, MAX_GROUP), np.nan, dtype=float)
    for r in range(MAX_PERIOD):
        for c in range(MAX_GROUP):
            z = GRID[r][c]
            if z is None:
                continue
            val = DF.loc[DF["Z"] == z, property_key].values[0]
            if not pd.isna(val):
                grid_vals[r, c] = float(val)
    fig, ax = plt.subplots()
    im = ax.imshow(grid_vals, origin="upper", aspect="auto")
    ax.set_xticks(range(MAX_GROUP))
    ax.set_xticklabels([str(i) for i in range(1, MAX_GROUP + 1)])
    ax.set_yticks(range(MAX_PERIOD))
    ax.set_yticklabels([str(i) for i in range(1, MAX_PERIOD + 1)])
    ax.set_xlabel("Group")
    ax.set_ylabel("Period")
    ax.set_title(f"Periodic heatmap: {prop_label}")
    fig.colorbar(im, ax=ax, label=prop_label)
    fig.tight_layout()
    return fig

# -----------------------------
# Callbacks
# -----------------------------
def element_info(z_or_symbol: str):
    try:
        if z_or_symbol.isdigit():
            Z = int(z_or_symbol)
            _ = elements[Z]
        else:
            el = elements.symbol(z_or_symbol)
            Z = el.number
    except Exception:
        return f"Unknown element: {z_or_symbol}", None, None

    row = DF.loc[DF["Z"] == Z].iloc[0].to_dict()
    symbol = row["symbol"]

    facts = []
    facts.extend(CURATED_FACTS.get(symbol, []))
    facts.append(GROUP_FACTS.get(row["category"], None))
    facts = [f for f in facts if f]

    props_lines = [
        f"{row['name']} ({symbol}), Z = {Z}",
        f"Period {int(row['period']) if not pd.isna(row['period']) else '—'}, "
        f"Group {row['group'] if row['group'] is not None else '—'}, "
        f"Block {row['block']} | Category: {row['category'].replace('-', ' ').title()}",
        f"Atomic mass: {row['mass'] if not pd.isna(row['mass']) else '—'} u",
        f"Density: {row['density'] if not pd.isna(row['density']) else '—'} g/cm³",
        f"Electronegativity: {row['electronegativity'] if not pd.isna(row['electronegativity']) else '—'} (Pauling)",
        f"Melting point: {row['melting_point'] if not pd.isna(row['melting_point']) else '—'} K | "
        f"Boiling point: {row['boiling_point'] if not pd.isna(row['boiling_point']) else '—'} K",
        f"vdW radius: {row['vdw_radius'] if not pd.isna(row['vdw_radius']) else '—'} pm | "
        f"Covalent radius: {row['covalent_radius'] if not pd.isna(row['covalent_radius']) else '—'} pm",
        f"Radioactive: {'Yes' if row['is_radioactive'] else 'No'}",
    ]
    info_text = "\n".join(props_lines)
    facts_text = "\n• ".join(["Interesting facts:"] + facts) if facts else "No fact on file—still cool though!"

    prop_key = "electronegativity" if not pd.isna(row["electronegativity"]) else "mass"
    trend_df = DF[["Z", "symbol", prop_key]].dropna()
    fig = plot_trend(trend_df, prop_key, Z, symbol)

    return info_text, facts_text, fig

def handle_button_click(z: int):
    return element_info(str(z))

def search_element(query: str):
    query = (query or "").strip()
    if not query:
        return gr.update(), gr.update(), gr.update()
    return element_info(query)

# -----------------------------
# UI (Gradio 4.29.0 compatible)
# -----------------------------
with gr.Blocks(title="Interactive Periodic Table") as demo:
    gr.Markdown("# 🧪 Interactive Periodic Table\nClick an element or search by symbol/name/atomic number.")

    with gr.Row():
        # Inspector first so buttons can target these outputs
        with gr.Column(scale=1):
            gr.Markdown("### Inspector")
            search = gr.Textbox(label="Search (symbol/name/Z)", placeholder="e.g., C, Iron, 79")
            info = gr.Textbox(label="Properties", lines=10, interactive=False)
            facts = gr.Markdown("Select an element to see fun facts.")
            trend = gr.Plot()
            search.submit(search_element, inputs=[search], outputs=[info, facts, trend])

            gr.Markdown("### Trend heatmap")
            prop = gr.Dropdown(choices=[k for k, _ in NUMERIC_PROPS], value="electronegativity", label="Property")
            heat = gr.Plot()
            prop.change(lambda k: plot_heatmap(k), inputs=[prop], outputs=[heat])
            demo.load(lambda: plot_heatmap("electronegativity"), outputs=[heat])

        with gr.Column(scale=2):
            gr.Markdown("### Main Table")
            # Group headers (1..18)
            with gr.Row():
                for g in range(1, 19):
                    gr.Markdown(f"**{g}**")
            # Grid of element buttons
            for r in range(MAX_PERIOD):
                with gr.Row():
                    for c in range(MAX_GROUP):
                        z = GRID[r][c]
                        if z is None:
                            gr.Button("", interactive=False)
                        else:
                            sym = DF.loc[DF["Z"] == z, "symbol"].values[0]
                            btn = gr.Button(sym)
                            btn.click(handle_button_click, inputs=[gr.Number(z, visible=False)],
                                      outputs=[info, facts, trend])

            gr.Markdown("### f-block (lanthanides & actinides)")
            with gr.Row():
                for z in LAN:
                    sym = DF.loc[DF["Z"] == z, "symbol"].values[0]
                    btn = gr.Button(sym)
                    btn.click(handle_button_click, inputs=[gr.Number(z, visible=False)],
                              outputs=[info, facts, trend])
            with gr.Row():
                for z in ACT:
                    sym = DF.loc[DF["Z"] == z, "symbol"].values[0]
                    btn = gr.Button(sym)
                    btn.click(handle_button_click, inputs=[gr.Number(z, visible=False)],
                              outputs=[info, facts, trend])

if __name__ == "__main__":
    demo.launch()