import gradio as gr
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from typing import Dict, List, Optional, Tuple
from periodictable import elements

# =========================
# Helpers & data utilities
# =========================
def to_float(x):
    """Coerce periodictable values (incl. uncertainties) to float; else NaN."""
    if x is None:
        return np.nan
    v = getattr(x, "nominal_value", x)
    try:
        return float(v)
    except Exception:
        return np.nan

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; dominant in stars."],
    "He": ["Inert; used in cryogenics and balloons."],
    "Li": ["Key in Li-ion batteries."],
    "C": ["Same element → diamond vs graphite (allotropy)."],
    "N": ["~78% of Earth’s atmosphere (N₂)."],
    "O": ["~21% of air; crucial for respiration."],
    "Na": ["Violently reacts with water."],
    "Mg": ["Burns with bright white flame."],
    "Si": ["Semiconductor backbone."],
    "Cl": ["Disinfectant; elemental Cl₂ is toxic."],
    "Fe": ["Steel & blood (heme) MVP."],
    "Cu": ["Great conductor; green patina."],
    "Ag": ["Highest electrical conductivity."],
    "Au": ["Very unreactive; great for electronics/jewelry."],
    "Hg": ["Liquid metal at room temp; toxic."],
    "Pb": ["Dense; toxicity drove phase-outs."],
    "U": ["Nuclear fuel (U-235)."],
    "Pu": ["Man-made in quantity; nuclear uses."],
    "F": ["Most electronegative; extremely reactive."],
    "Ne": ["Classic red-orange glow tubes."],
    "Xe": ["HID lamps & flashes."],
}

GROUP_FACTS = {
    "alkali": "Alkali metal: very reactive; forms +1; reacts with water.",
    "alkaline-earth": "Alkaline earth metal: reactive; forms +2.",
    "transition": "Transition metal: variable oxidation states; often colored compounds.",
    "post-transition": "Post-transition metal: softer; lower melting than transition metals.",
    "metalloid": "Metalloid: between metals and nonmetals; often semiconductors.",
    "nonmetal": "Nonmetal: covalent chemistry; key biological roles.",
    "halogen": "Halogen: ns²np⁵; gains 1e⁻; forms salts.",
    "noble-gas": "Noble gas: ns²np⁶; inert, monatomic.",
    "lanthanide": "Lanthanide: rare earths; magnets/lasers/phosphors.",
    "actinide": "Actinide: radioactive; nuclear 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 == "d":
            return "transition"
        if el.block == "p" and el.group == 17:
            return "halogen"
        if el.block == "p" and el.group == 18:
            return "noble-gas"
        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
        rows.append({
            "Z": el.number,
            "symbol": el.symbol,
            "name": el.name.title(),
            "period": getattr(el, "period", None),
            "group": getattr(el, "group", None),
            "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)),
        })
    return pd.DataFrame(rows).sort_values("Z").reset_index(drop=True)

DF = build_elements_df()

# =========================
# Hard-coded periodic layout
# =========================
# Periods 1–7, groups 1–18; La/Ac shown in group 3; f-block split below.
GRID = [
    [1, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, 2],
    [3, 4, None, None, None, None, None, None, None, None, None, None, 5, 6, 7, 8, 9, 10],
    [11, 12, None, None, None, None, None, None, None, None, None, None, 13, 14, 15, 16, 17, 18],
    [19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36],
    [37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54],
    [55, 56, 57, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86],
    [87, 88, 89, 104, 105, 106, 107, 108, 109, 110, 111, 112, 113, 114, 115, 116, 117, 118],
]
LAN = list(range(58, 72))   # Ce..Lu
ACT = list(range(90, 104))  # Th..Lr

def find_pos_in_grid(Z:int) -> Tuple[Optional[int], Optional[int]]:
    for r in range(len(GRID)):
        for c in range(len(GRID[0])):
            if GRID[r][c] == Z:
                return (r+1, c+1)  # human-friendly (period, group)
    return (None, None)

# =========================
# Explanations
# =========================
def valence_pattern(period:int, group:int, block:str) -> str:
    if period is None or group is None or block is None:
        return "Valence pattern unavailable."
    n = period
    if block == "s":
        return f"{n}s¹" if group == 1 else f"{n}s²"
    if block == "p" and 13 <= group <= 18:
        p_e = group - 12  # 1..6
        return f"{n}s²{n}p^{p_e}"
    if block == "d":
        return f"{n-1}d^(1–10){n}s^(0–2) (incomplete d-subshell)"
    if block == "f":
        return f"{n-2}f^(1–14){n-1}d^(0–1){n}s² (f-block)"
    return "Valence pattern unavailable."

def explain_element(row:dict, Z:int) -> str:
    period, group = find_pos_in_grid(Z)
    block = row["block"]
    cat = row["category"]
    en = row["electronegativity"]
    dens = row["density"]

    lines = []
    # Valence / block logic
    lines.append(f"**Valence & block:** {valence_pattern(period, group, block)}; {cat.replace('-', ' ')}.")
    # Reactivity / tendencies
    if group == 1:
        lines.append("**Reactivity:** Group 1 (ns¹) → easily loses 1 e⁻ (forms +1), reacts strongly with water.")
    elif group == 2:
        lines.append("**Reactivity:** Group 2 (ns²) → tends to lose 2 e⁻ (forms +2).")
    elif group == 17:
        lines.append("**Reactivity:** Halogen (ns²np⁵) → tends to gain 1 e⁻; oxidizing; reactivity decreases down the group.")
    elif group == 18:
        lines.append("**Reactivity:** Noble gas (ns²np⁶) → filled shell, minimal reactivity.")
    elif block == "d":
        lines.append("**d-block behavior:** Partially filled d-orbitals → multiple oxidation states; often colored complexes.")
    # Property tie-ins
    if not pd.isna(en) and not pd.isna(row["period"]):
        same_period = DF[(DF["period"] == row["period"]) & (~DF["electronegativity"].isna())]
        if len(same_period):
            med = same_period["electronegativity"].median()
            qual = "higher-than-average" if en > med else "lower-than-average"
            lines.append(f"**Electronegativity:** {en:.2f} ({qual} within period {int(row['period'])}).")
    if not pd.isna(dens):
        lines.append(f"**Density:** {dens:g} g/cm³ — linked to atomic mass and packing typical for its category.")

    return "### Why it behaves this way\n" + "\n".join(f"- {t}" for t in lines)

# =========================
# 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]
    max_period, max_group = len(GRID), len(GRID[0])
    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)

    if np.isnan(grid_vals).all():
        fig, ax = plt.subplots()
        ax.axis("off")
        ax.text(0.5, 0.5, f"No data for {prop_label}", ha="center", va="center", fontsize=12)
        fig.tight_layout()
        return fig

    masked = np.ma.masked_invalid(grid_vals)
    finite_vals = grid_vals[~np.isnan(grid_vals)]
    if finite_vals.size >= 2:
        vmin, vmax = np.nanpercentile(finite_vals, [5, 95])
    else:
        vmin, vmax = np.nanmin(finite_vals), np.nanmax(finite_vals)

    fig, ax = plt.subplots()
    im = ax.imshow(masked, origin="upper", aspect="auto", vmin=vmin, vmax=vmax)
    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

# =========================
# Core callbacks
# =========================
def compose_facts(row:dict, Z:int, show_expl:bool) -> str:
    symbol = row["symbol"]
    facts = []
    facts.extend(CURATED_FACTS.get(symbol, []))
    gf = GROUP_FACTS.get(row["category"], None)
    if gf:
        facts.append(gf)
    facts_text = "\n• ".join(["**Interesting facts:**"] + facts) if facts else ""
    if show_expl:
        expl = explain_element(row, Z)
        facts_text = (facts_text + "\n\n" if facts_text else "") + expl
    return facts_text if facts_text else "No fact on file—still cool though!"

def element_info(z_or_symbol: str, show_expl: bool):
    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}", "No data", None, None  # info, facts, fig, current_Z

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

    def show(v):
        return v if (v is not None and not pd.isna(v)) else "—"

    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: {show(row['mass'])} u",
        f"Density: {show(row['density'])} g/cm³",
        f"Electronegativity: {show(row['electronegativity'])} (Pauling)",
        f"Melting point: {show(row['melting_point'])} K | Boiling point: {show(row['boiling_point'])} K",
        f"vdW radius: {show(row['vdw_radius'])} pm | Covalent radius: {show(row['covalent_radius'])} pm",
        f"Radioactive: {'Yes' if row['is_radioactive'] else 'No'}",
    ]
    info_text = "\n".join(props_lines)

    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)

    facts_text = compose_facts(row, Z, show_expl)
    return info_text, facts_text, fig, Z

def handle_button_click(z: int, show_expl: bool):
    return element_info(str(z), show_expl)

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

def refresh_facts(current_Z: Optional[int], show_expl: bool):
    if current_Z is None:
        return gr.update()
    row = DF.loc[DF["Z"] == current_Z].iloc[0].to_dict()
    return compose_facts(row, int(current_Z), show_expl)

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

    with gr.Row():
        # Inspector & controls
        with gr.Column(scale=1):
            gr.Markdown("### Inspector")
            show_expl = gr.Checkbox(label="Show advanced explanation", value=False)
            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 facts and explanations.")
            trend = gr.Plot()
            current_Z = gr.State(value=None)

            search.submit(search_element, inputs=[search, show_expl], outputs=[info, facts, trend, current_Z])
            show_expl.change(refresh_facts, inputs=[current_Z, show_expl], outputs=[facts])

            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])

        # Main table
        with gr.Column(scale=2):
            gr.Markdown("### Main Table")
            with gr.Row():
                for g in range(1, 19):
                    gr.Markdown(f"**{g}**")
            for r in range(len(GRID)):
                with gr.Row():
                    for c in range(len(GRID[0])):
                        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), show_expl],
                                outputs=[info, facts, trend, current_Z],
                            )

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

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