app.py

import gradio as gr
import json, re, os

DATA_PATH = os.path.join(os.path.dirname(__file__), "elements.json")
with open(DATA_PATH) as f:
    ELEMENTS = json.load(f)

SYM_MAP = {e["symbol"]: e for e in ELEMENTS}
NAME_MAP = {e["name"].lower(): e for e in ELEMENTS}

CATEGORY_COLORS = {"Alkali Metal": "#FF6B6B", "Alkaline Earth": "#FFA07A", "Transition Metal": "#7EC8E3", "Post-Transition": "#98D8C8", "Metalloid": "#C3B1E1", "Noble Gas": "#FFFACD", "Halogen": "#F0E68C", "Nonmetal": "#90EE90", "Lanthanide": "#FFB6C1", "Actinide": "#DDA0DD"}


def lookup_element(query):
    query = query.strip()
    if not query:
        return "Please enter an element name, symbol, or atomic number."
    el = None
    if query.isdigit():
        num = int(query)
        for e in ELEMENTS:
            if e["number"] == num:
                el = e
                break
    if el is None:
        el = SYM_MAP.get(query) or SYM_MAP.get(query.capitalize())
    if el is None:
        el = NAME_MAP.get(query.lower())
    if el is None:
        return f"Element '{query}' not found. Try a name, symbol, or number."
    cc = CATEGORY_COLORS.get(el["category"], "#DDD")
    return f'''<div style="font-family:system-ui,sans-serif;max-width:680px;margin:0 auto"><div style="display:flex;align-items:center;gap:24px;margin-bottom:20px"><div style="background:{cc};border-radius:16px;width:120px;height:120px;display:flex;flex-direction:column;align-items:center;justify-content:center;box-shadow:0 4px 12px rgba(0,0,0,.15)"><span style="font-size:12px;opacity:.7">{el["number"]}</span><span style="font-size:42px;font-weight:700">{el["symbol"]}</span><span style="font-size:11px">{el["atomic_mass"]:.4f}</span></div><div><h2 style="margin:0 0 4px 0;font-size:28px">{el["name"]}</h2><span style="background:{cc};padding:3px 10px;border-radius:12px;font-size:13px">{el["category"]}</span></div></div><table style="width:100%;border-collapse:collapse;font-size:14px"><tr style="background:#f8f9fa"><td style="padding:8px 12px;font-weight:600">Group / Period</td><td style="padding:8px 12px">{el.get("group","N/A")} / {el["period"]}</td><td style="padding:8px 12px;font-weight:600">Block</td><td style="padding:8px 12px">{el["block"]}</td></tr><tr><td style="padding:8px 12px;font-weight:600">Phase</td><td style="padding:8px 12px">{el["phase"]}</td><td style="padding:8px 12px;font-weight:600">Atomic Mass</td><td style="padding:8px 12px">{el["atomic_mass"]:.4f} u</td></tr><tr style="background:#f8f9fa"><td style="padding:8px 12px;font-weight:600">Electron Config</td><td style="padding:8px 12px;font-family:monospace">{el.get("electron_config","N/A")}</td><td style="padding:8px 12px;font-weight:600">Electronegativity</td><td style="padding:8px 12px">{el.get("electronegativity","N/A")}</td></tr><tr><td style="padding:8px 12px;font-weight:600">Melting Point</td><td style="padding:8px 12px">{el.get("melting_point","N/A")} K</td><td style="padding:8px 12px;font-weight:600">Boiling Point</td><td style="padding:8px 12px">{el.get("boiling_point","N/A")} K</td></tr><tr style="background:#f8f9fa"><td style="padding:8px 12px;font-weight:600">Density</td><td style="padding:8px 12px">{el.get("density","N/A")}</td><td style="padding:8px 12px;font-weight:600">Ionization Energy</td><td style="padding:8px 12px">{el.get("ionization_energy","N/A")} kJ/mol</td></tr><tr><td style="padding:8px 12px;font-weight:600">Oxidation States</td><td style="padding:8px 12px" colspan="3">{el.get("oxidation_states","N/A")}</td></tr><tr style="background:#f8f9fa"><td style="padding:8px 12px;font-weight:600">Discovery</td><td style="padding:8px 12px" colspan="3">{el.get("discovered_by","Unknown")} ({el.get("discovery_year","?")})</td></tr></table></div>'''

def parse_formula(formula):
    tokens = re.findall(r'([A-Z][a-z]?)(\\d*)', formula)
    comp = {}
    for sym, count in tokens:
        if sym:
            comp[sym] = comp.get(sym, 0) + (int(count) if count else 1)
    return comp


def expand_formula(formula):
    try:
        return _parse_group(formula, 0, len(formula))
    except Exception:
        return parse_formula(formula) or None


def _parse_group(formula, start, end):
    comp = {}
    i = start
    while i < end:
        if formula[i] == '(':
            depth, j = 1, i + 1
            while j < end and depth > 0:
                if formula[j] == '(': depth += 1
                elif formula[j] == ')': depth -= 1
                j += 1
            inner = _parse_group(formula, i + 1, j - 1)
            k, num_str = j, ""
            while k < end and formula[k].isdigit(): num_str += formula[k]; k += 1
            mult = int(num_str) if num_str else 1
            for s, c in inner.items(): comp[s] = comp.get(s, 0) + c * mult
            i = k
        elif formula[i].isupper():
            sym = formula[i]
            i += 1
            while i < end and formula[i].islower(): sym += formula[i]; i += 1
            num_str = ""
            while i < end and formula[i].isdigit(): num_str += formula[i]; i += 1
            comp[sym] = comp.get(sym, 0) + (int(num_str) if num_str else 1)
        else:
            i += 1
    return comp


def calc_molar_mass(formula):
    formula = formula.strip()
    if not formula: return "Enter a formula like H2O, NaCl, or C6H12O6."
    expanded = expand_formula(formula)
    if expanded is None: return f"Could not parse: {formula}"
    rows, total = [], 0.0
    for sym, count in expanded.items():
        el = SYM_MAP.get(sym)
        if el is None: return f"Unknown element: {sym}"
        mass = el["atomic_mass"]
        sub = mass * count
        total += sub
        rows.append((sym, el["name"], count, mass, sub))
    trows = ''.join(f'<tr><td style="padding:6px 12px;font-weight:600">{s}</td><td style="padding:6px 12px">{n}</td><td style="padding:6px 12px;text-align:center">{c}</td><td style="padding:6px 12px;text-align:right">{m:.4f}</td><td style="padding:6px 12px;text-align:right">{su:.4f}</td><td style="padding:6px 12px;text-align:right">{su/total*100:.1f}%</td></tr>' for s,n,c,m,su in rows)
    return f'<div style="font-family:system-ui,sans-serif;max-width:600px;margin:0 auto"><div style="text-align:center;margin-bottom:16px"><span style="font-size:16px;color:#666">Molar Mass of</span> <span style="font-size:24px;font-weight:700">{formula}</span></div><div style="background:linear-gradient(135deg,#667eea,#764ba2);color:white;text-align:center;padding:20px;border-radius:12px;margin-bottom:16px"><div style="font-size:36px;font-weight:700">{total:.4f}</div><div style="font-size:14px;opacity:.9">g/mol</div></div><table style="width:100%;border-collapse:collapse;font-size:14px"><tr style="background:#f0f0f0;font-weight:600"><td style="padding:6px 12px">Symbol</td><td style="padding:6px 12px">Element</td><td style="padding:6px 12px;text-align:center">Count</td><td style="padding:6px 12px;text-align:right">Mass</td><td style="padding:6px 12px;text-align:right">Subtotal</td><td style="padding:6px 12px;text-align:right">%</td></tr>{trows}</table></div>'

def balance_equation(equation):
    from itertools import product as iprod
    equation = equation.strip()
    if not equation: return "Enter an equation like: Fe + O2 -> Fe2O3"
    sides = re.split(r'\s*(?:->|=)\s*', equation)
    if len(sides) != 2: return "Use -> or = to separate reactants and products."
    rstr = [s.strip() for s in sides[0].split('+')]
    pstr = [s.strip() for s in sides[1].split('+')]
    compounds = rstr + pstr
    nr = len(rstr)
    comps, elems = [], set()
    for c in compounds:
        p = expand_formula(c)
        if p is None: return f"Could not parse: {c}"
        comps.append(p)
        elems.update(p.keys())
    elems = sorted(elems)
    def check(coeffs):
        for e in elems:
            if sum(coeffs[i]*comps[i].get(e,0) for i in range(nr)) != sum(coeffs[nr+j]*comps[nr+j].get(e,0) for j in range(len(pstr))):
                return False
        return True
    found = None
    for coeffs in iprod(range(1, 21), repeat=len(compounds)):
        if check(coeffs):
            found = coeffs
            break
    if found is None: return "Could not balance this equation."
    lp = [f"{found[i]}{c}" if found[i]>1 else c for i,c in enumerate(rstr)]
    rp = [f"{found[nr+j]}{c}" if found[nr+j]>1 else c for j,c in enumerate(pstr)]
    bal = " + ".join(lp) + " -> " + " + ".join(rp)
    vrows = ''.join(f'<tr><td style="padding:4px 12px;font-weight:600">{e}</td><td style="padding:4px 12px;text-align:center">{sum(found[i]*comps[i].get(e,0) for i in range(nr))}</td><td style="padding:4px 12px;text-align:center">{sum(found[nr+j]*comps[nr+j].get(e,0) for j in range(len(pstr)))}</td></tr>' for e in elems)
    return f'<div style="font-family:system-ui,sans-serif;max-width:600px;margin:0 auto"><div style="background:linear-gradient(135deg,#11998e,#38ef7d);color:white;text-align:center;padding:20px;border-radius:12px;margin-bottom:16px"><div style="font-size:13px;opacity:.85;margin-bottom:4px">Balanced Equation</div><div style="font-size:22px;font-weight:700;font-family:monospace">{bal}</div></div><h4 style="margin:12px 0 8px">Atom Count Verification</h4><table style="width:100%;border-collapse:collapse;font-size:14px"><tr style="background:#f0f0f0;font-weight:600"><td style="padding:4px 12px">Element</td><td style="padding:4px 12px;text-align:center">Left</td><td style="padding:4px 12px;text-align:center">Right</td></tr>{vrows}</table></div>'

IONS = {
    "Monatomic Cations": [("H+","Hydrogen"),("Li+","Lithium"),("Na+","Sodium"),("K+","Potassium"),("Ag+","Silver"),("Mg2+","Magnesium"),("Ca2+","Calcium"),("Ba2+","Barium"),("Zn2+","Zinc"),("Al3+","Aluminum")],
    "Monatomic Anions": [("F-","Fluoride"),("Cl-","Chloride"),("Br-","Bromide"),("I-","Iodide"),("O2-","Oxide"),("S2-","Sulfide"),("N3-","Nitride")],
    "Polyatomic Ions": [("CO3 2-","Carbonate"),("NO3-","Nitrate"),("SO4 2-","Sulfate"),("PO4 3-","Phosphate"),("ClO3-","Chlorate"),("ClO4-","Perchlorate")],
    "Transition Metal Ions": [("Cu+","Copper(I)"),("Cu2+","Copper(II)"),("Fe2+","Iron(II)"),("Fe3+","Iron(III)"),("Pb2+","Lead(II)"),("MnO4-","Permanganate")],
    "Special Ions": [("NH4+","Ammonium"),("OH-","Hydroxide"),("HCO3-","Bicarbonate"),("CH3COO-","Acetate"),("CN-","Cyanide")],
}


def build_ions_html():
    colors = ["#FF6B6B","#FFA07A","#7EC8E3","#98D8C8","#C3B1E1"]
    html = '<div style="font-family:system-ui,sans-serif;max-width:700px;margin:0 auto">'
    for idx, (cat, ions) in enumerate(IONS.items()):
        c = colors[idx % len(colors)]
        html += f'<h3 style="margin:16px 0 8px;color:#333">{cat}</h3><div style="display:flex;flex-wrap:wrap;gap:8px;margin-bottom:12px">'
        for sym, name in ions:
            html += f'<div style="background:{c}22;border:2px solid {c};border-radius:10px;padding:8px 14px;text-align:center;min-width:80px"><div style="font-size:18px;font-weight:700">{sym}</div><div style="font-size:11px;color:#555">{name}</div></div>'
        html += '</div>'
    html += '</div>'
    return html


with gr.Blocks(title="Chemistry Toolkit", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# Chemistry Toolkit\n*Interactive chemistry reference inspired by [Zperiod](https://zperiod.app)*")
    with gr.Tab("Element Lookup"):
        gr.Markdown("Search by **name**, **symbol**, or **atomic number**.")
        with gr.Row():
            elem_input = gr.Textbox(label="Search", placeholder="e.g. Gold, Au, or 79", scale=3)
            elem_btn = gr.Button("Look Up", variant="primary", scale=1)
        elem_output = gr.HTML()
        elem_btn.click(lookup_element, inputs=elem_input, outputs=elem_output)
        elem_input.submit(lookup_element, inputs=elem_input, outputs=elem_output)
        gr.Examples(["Oxygen", "Fe", "79", "Carbon", "Cl"], inputs=elem_input)
    with gr.Tab("Molar Mass"):
        gr.Markdown("Enter a chemical formula to calculate its **molar mass**.")
        with gr.Row():
            mm_input = gr.Textbox(label="Chemical Formula", placeholder="e.g. H2O, CaCO3", scale=3)
            mm_btn = gr.Button("Calculate", variant="primary", scale=1)
        mm_output = gr.HTML()
        mm_btn.click(calc_molar_mass, inputs=mm_input, outputs=mm_output)
        mm_input.submit(calc_molar_mass, inputs=mm_input, outputs=mm_output)
        gr.Examples(["H2O", "NaCl", "C6H12O6", "Ca(OH)2", "H2SO4"], inputs=mm_input)
    with gr.Tab("Equation Balancer"):
        gr.Markdown("Enter an unbalanced equation using `->` or `=` to separate sides.")
        with gr.Row():
            eq_input = gr.Textbox(label="Unbalanced Equation", placeholder="e.g. Fe + O2 -> Fe2O3", scale=3)
            eq_btn = gr.Button("Balance", variant="primary", scale=1)
        eq_output = gr.HTML()
        eq_btn.click(balance_equation, inputs=eq_input, outputs=eq_output)
        eq_input.submit(balance_equation, inputs=eq_input, outputs=eq_output)
        gr.Examples(["Fe + O2 -> Fe2O3", "H2 + O2 -> H2O", "CH4 + O2 -> CO2 + H2O", "Na + Cl2 -> NaCl"], inputs=eq_input)
    with gr.Tab("Common Ions"):
        gr.Markdown("Quick reference for common ions.")
        gr.HTML(build_ions_html())

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