utils/plots.py

"""Plotly figure factory for the LiB Simulation AI Engine dashboard."""

import numpy as np
import pandas as pd
import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

FF_COLORS = {
    "Yun et al.": "#3A86FF",
    "Wang et al.": "#4CAF50",
    "This Work": "#FF4757",
    "This Work (De Angelis 2024)": "#FF4757",
    "Yun et al. (2017)": "#3A86FF",
    "Wang et al. (2020)": "#4CAF50",
    "DFT (reference)": "#222",
    "DFT – Vacancy": "#8B5CF6",
    "DFT – Knock-off": "#F59E0B",
    "DFT – Direct-hopping": "#06B6D4",
    "DFT reference": "#999",
}

_DARK = dict(
    plot_bgcolor="rgba(0,0,0,0)",
    paper_bgcolor="rgba(0,0,0,0)",
    font=dict(color="#EEE"),
    legend=dict(bgcolor="rgba(0,0,0,0)"),
)
_GRID = dict(gridcolor="rgba(255,255,255,0.1)")


def _apply_dark(fig: go.Figure, height: int = 400) -> go.Figure:
    fig.update_layout(height=height, **_DARK, margin=dict(t=55, b=40))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


# ── Force Field Performance ────────────────────────────────────────────────────
def ff_performance_chart(df_perf: dict) -> go.Figure:
    ffs = df_perf["Force Field"]
    fig = make_subplots(rows=1, cols=2,
                        subplot_titles=["R² (higher is better)", "RMSE (lower is better)"],
                        horizontal_spacing=0.12)
    colors = [FF_COLORS.get(f, "#888") for f in ffs]
    for col, (label, key) in enumerate([("Energy R²", "Energy R²"), ("Force R²", "Force R²")]):
        for i, (ff, val) in enumerate(zip(ffs, df_perf[key])):
            fig.add_trace(go.Bar(name=ff, x=[label], y=[val], marker_color=colors[i],
                                 showlegend=(col == 0), legendgroup=ff,
                                 text=[f"{val:.3f}"], textposition="outside"), row=1, col=1)
    for label, key, scale in [("Energy RMSE (eV)", "Energy RMSE (eV)", 1),
                               ("Force RMSE ×10⁻³ (eV/Å)", "Force RMSE (eV/Å)", 1000)]:
        for i, (ff, val) in enumerate(zip(ffs, df_perf[key])):
            fig.add_trace(go.Bar(name=ff, x=[label], y=[val * scale], marker_color=colors[i],
                                 showlegend=False, legendgroup=ff,
                                 text=[f"{val*scale:.4f}"], textposition="outside"), row=1, col=2)
    fig.update_layout(barmode="group", height=440, **_DARK, margin=dict(t=60, b=40))
    fig.update_xaxes(showgrid=False)
    fig.update_yaxes(**_GRID)
    return fig


def loss_curve_chart(df: pd.DataFrame) -> go.Figure:
    ma = df["Loss (SSE)"].rolling(100, min_periods=1).mean()
    fig = go.Figure()
    fig.add_trace(go.Scatter(x=df["Iteration"], y=df["Loss (SSE)"], mode="lines",
                             name="Loss (SSE)", line=dict(color="#3A86FF", width=0.6), opacity=0.4))
    fig.add_trace(go.Scatter(x=df["Iteration"], y=ma, mode="lines",
                             name="Moving Avg (100 iter)", line=dict(color="#FF4757", width=2)))
    fig.add_vline(x=5000, line_dash="dash", line_color="#F59E0B",
                  annotation_text="bond → vdW phase", annotation_font_color="#F59E0B")
    fig.update_layout(xaxis_title="Optimization Iteration", yaxis_title="Loss (SSE)", **_DARK,
                      height=360, margin=dict(t=20, b=40))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


# ── Diffusion & Arrhenius ──────────────────────────────────────────────────────
def arrhenius_plot(df: pd.DataFrame, selected_system: str) -> go.Figure:
    fig = go.Figure()
    subset = df[df["System"] == selected_system]
    for ff in subset["Force Field"].unique():
        s = subset[subset["Force Field"] == ff]
        fig.add_trace(go.Scatter(x=s["1000/T (K⁻¹)"], y=s["log₁₀(D)"], mode="lines",
                                 name=ff, line=dict(color=FF_COLORS.get(ff, "#888"), width=2.5)))
    dft_df = df[df["Force Field"] == "DFT reference"]
    for mech_full in dft_df["System"].unique():
        s = dft_df[dft_df["System"] == mech_full]
        fig.add_trace(go.Scatter(x=s["1000/T (K⁻¹)"], y=s["log₁₀(D)"], mode="lines",
                                 name=mech_full, line=dict(color=FF_COLORS.get(mech_full, "#aaa"),
                                                            width=1.5, dash="dot")))
    for T_mark, label in [(300, "300 K"), (400, "400 K"), (500, "500 K")]:
        fig.add_vline(x=1000 / T_mark, line_color="rgba(255,255,255,0.15)", line_dash="dot",
                      annotation_text=label, annotation_font_color="rgba(255,255,255,0.5)",
                      annotation_font_size=10)
    fig.update_layout(xaxis_title="1000/T  (K⁻¹)", yaxis_title="log₁₀[D  (cm²/s)]",
                      title=f"Arrhenius Plot — {selected_system}", **_DARK,
                      height=420, margin=dict(t=50, b=40))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


def diffusion_bar_chart(df: pd.DataFrame, temperature: int) -> go.Figure:
    sub = df[df["Temperature (K)"] == temperature].copy()
    fig = px.bar(sub, x="System", y="D (cm²/s)", color="Force Field", barmode="group",
                 color_discrete_map={k: v for k, v in FF_COLORS.items()}, log_y=True,
                 labels={"D (cm²/s)": "Diffusion Coefficient (cm²/s)"},
                 title=f"Li Diffusion Coefficient at {temperature} K")
    return _apply_dark(fig, 420)


def activation_energy_comparison() -> go.Figure:
    from utils.data import ARRHENIUS_PARAMS
    systems = list(ARRHENIUS_PARAMS.keys())
    ffs = ["Yun et al.", "Wang et al.", "This Work"]
    fig = go.Figure()
    for ff in ffs:
        eas = [ARRHENIUS_PARAMS[s][ff]["Ea"] for s in systems]
        ea_errs = [ARRHENIUS_PARAMS[s][ff]["Ea_err"] for s in systems]
        fig.add_trace(go.Bar(name=ff, x=systems, y=eas,
                             error_y=dict(type="data", array=ea_errs, visible=True),
                             marker_color=FF_COLORS.get(ff, "#888")))
    fig.add_hline(y=24.1, line_dash="dot", line_color="#F59E0B",
                  annotation_text="CI-NEB knock-off (24.1 kJ/mol)", annotation_font_color="#F59E0B")
    fig.add_hline(y=63.7, line_dash="dot", line_color="#8B5CF6",
                  annotation_text="CI-NEB vacancy (63.7 kJ/mol)", annotation_font_color="#8B5CF6")
    fig.update_layout(barmode="group", xaxis_title="System", yaxis_title="Eₐ (kJ/mol)",
                      title="Activation Energy — All Systems & Force Fields", **_DARK,
                      height=450, margin=dict(t=60, b=60))
    fig.update_xaxes(**_GRID, tickangle=-10)
    fig.update_yaxes(**_GRID)
    return fig


def msd_plot(df: pd.DataFrame) -> go.Figure:
    fig = go.Figure()
    palette = px.colors.qualitative.Plotly
    for i, label in enumerate(df["Label"].unique()):
        s = df[df["Label"] == label]
        fig.add_trace(go.Scatter(x=s["Time (ps)"], y=s["MSD (Ų)"], mode="lines",
                                 name=label, line=dict(color=palette[i % len(palette)], width=1.8)))
    fig.update_layout(xaxis_title="Time (ps)", yaxis_title="MSD (Ų)",
                      title="Mean Square Displacement — Li Atoms", **_DARK,
                      height=400, margin=dict(t=50, b=40))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


# ── Crystal stability ──────────────────────────────────────────────────────────
def strain_energy_plot(df: pd.DataFrame) -> go.Figure:
    fig = go.Figure()
    dash_map = {"DFT (reference)": "solid", "Yun et al.": "dash",
                "Wang et al.": "dot", "This Work": "solid"}
    for method in df["Method"].unique():
        s = df[df["Method"] == method]
        fig.add_trace(go.Scatter(x=s["Strain ε₁₂"], y=s["Energy (eV/atom)"], mode="lines",
                                 name=method,
                                 line=dict(color=FF_COLORS.get(method, "#888"),
                                           dash=dash_map.get(method, "solid"), width=2)))
    fig.add_vline(x=0, line_color="rgba(255,255,255,0.3)", line_dash="dot")
    fig.update_layout(xaxis_title="Shear Strain ε₁₂", yaxis_title="Energy (eV/atom)",
                      title="LiF Energy–Strain (Shear Deformation)", **_DARK,
                      height=400, margin=dict(t=50, b=40))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


def eos_plot(df: pd.DataFrame) -> go.Figure:
    fig = go.Figure()
    for method in df["Method"].unique():
        s = df[df["Method"] == method]
        fig.add_trace(go.Scatter(x=s["V/V₀"], y=s["Energy (eV/atom)"], mode="lines",
                                 name=method,
                                 line=dict(color=FF_COLORS.get(method, "#888"), width=2)))
    fig.add_vline(x=1.0, line_color="rgba(255,255,255,0.3)", line_dash="dot",
                  annotation_text="V₀", annotation_font_color="rgba(255,255,255,0.5)")
    fig.update_layout(xaxis_title="V/V₀", yaxis_title="Energy (eV/atom)",
                      title="Murnaghan Equation of State — LiF", **_DARK,
                      height=400, margin=dict(t=50, b=40))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


def rdf_plot(df: pd.DataFrame, ff_label: str) -> go.Figure:
    times = sorted(df["Time (ps)"].unique())
    colorscale = px.colors.sequential.Plasma
    fig = go.Figure()
    for i, t in enumerate(times):
        s = df[df["Time (ps)"] == t]
        frac = i / max(len(times) - 1, 1)
        idx = int(frac * (len(colorscale) - 1))
        fig.add_trace(go.Scatter(x=s["r (Å)"], y=s["g(r)"], mode="lines", name=f"{t} ps",
                                 line=dict(color=colorscale[idx], width=1.8)))
    fig.add_vline(x=2.01, line_color="rgba(255,255,255,0.4)", line_dash="dash",
                  annotation_text="d(Li–F)=2.01 Å",
                  annotation_font_color="rgba(255,255,255,0.5)")
    fig.update_layout(xaxis_title="r  (Å)", yaxis_title="g(r)",
                      title=f"Li–F Radial Distribution Function — {ff_label}", **_DARK,
                      height=400, margin=dict(t=50, b=40))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


# ── SEI Analysis ───────────────────────────────────────────────────────────────
def sei_ranking_chart(df: pd.DataFrame) -> go.Figure:
    df_sorted = df.sort_values("Overall Score", ascending=True)
    colors = ["#FF4757" if c == "LiF (FEC-derived)" else "#3A86FF" for c in df_sorted["Component"]]
    fig = go.Figure(go.Bar(x=df_sorted["Overall Score"], y=df_sorted["Component"],
                           orientation="h", marker_color=colors,
                           text=df_sorted["Overall Score"].apply(lambda v: f"{v:.2f}"),
                           textposition="outside"))
    fig.update_layout(xaxis_title="Overall Score", title="SEI Component Ranking (Weighted Score)",
                      **_DARK, height=400, margin=dict(t=50, b=40, l=180))
    fig.update_xaxes(range=[0, 11], **_GRID)
    fig.update_yaxes(gridcolor=None)
    return fig


def sei_radar(df: pd.DataFrame, components: list) -> go.Figure:
    cats = ["Ionic Conductivity Score", "Electronic Insulation Score",
            "Mechanical Stability Score", "Thermal Stability Score"]
    fig = go.Figure()
    palette = px.colors.qualitative.Plotly
    for i, comp in enumerate(components):
        row = df[df["Component"] == comp].iloc[0]
        vals = [row[c] for c in cats]
        vals_closed = vals + [vals[0]]
        cats_closed = cats + [cats[0]]
        fig.add_trace(go.Scatterpolar(r=vals_closed, theta=cats_closed, fill="toself",
                                      name=comp, line=dict(color=palette[i % len(palette)]),
                                      opacity=0.7))
    fig.update_layout(polar=dict(radialaxis=dict(visible=True, range=[0, 10], color="#aaa"),
                                 bgcolor="rgba(0,0,0,0)"),
                      **_DARK, height=420, title="SEI Component Property Radar",
                      margin=dict(t=60, b=20))
    return fig


# ── ML model benchmark ─────────────────────────────────────────────────────────
def ml_model_scatter(df: pd.DataFrame) -> go.Figure:
    fig = px.scatter(df, x="Energy MAE (meV/atom)", y="Force MAE (meV/Å)",
                     size="Inference Speed (rel.)", color="Model",
                     hover_data=["R² Energy", "Training Data (DFT pts)"],
                     text="Model", size_max=40,
                     title="ML Force Field Accuracy vs Computational Speed",
                     color_discrete_sequence=px.colors.qualitative.Plotly)
    fig.update_traces(textposition="top center", textfont_size=10)
    return _apply_dark(fig, 440)


# ── Battery Property Predictor charts ─────────────────────────────────────────
def cycle_life_plot(df: pd.DataFrame) -> go.Figure:
    fig = px.line(df, x="Cycle", y="Capacity Retention (%)", color="SEI Type",
                  title="Capacity Retention vs. Cycle Number",
                  color_discrete_sequence=px.colors.qualitative.Plotly)
    fig.add_hline(y=80, line_dash="dash", line_color="rgba(255,200,0,0.6)",
                  annotation_text="80% EOL threshold",
                  annotation_font_color="rgba(255,200,0,0.8)")
    return _apply_dark(fig, 420)


def rate_capability_plot(df: pd.DataFrame) -> go.Figure:
    fig = px.line(df, x="C-rate", y="Discharge Capacity (mAh/g)", color="Anode / SEI",
                  log_x=True, title="Rate Capability — Discharge Capacity vs. C-Rate",
                  color_discrete_sequence=px.colors.qualitative.Plotly,
                  markers=True)
    return _apply_dark(fig, 420)


def thermal_safety_plot(df: pd.DataFrame) -> go.Figure:
    risk_color = {"Safe": "#4ade80", "Moderate": "#F59E0B", "High": "#f87171", "Critical": "#dc2626"}
    colors = [risk_color.get(r, "#888") for r in df["Risk Level"]]
    fig = go.Figure(go.Bar(
        x=df["Temperature (°C)"],
        y=df["Component / Event"],
        orientation="h",
        marker_color=colors,
        text=df["Temperature (°C)"].apply(lambda v: f"{v} °C"),
        textposition="outside",
        hovertext=df["Notes"],
    ))
    fig.add_vline(x=130, line_dash="dash", line_color="#dc2626",
                  annotation_text="Thermal runaway onset (130 °C)",
                  annotation_font_color="#dc2626")
    fig.update_layout(xaxis_title="Temperature (°C)",
                      title="Thermal Safety — Component Decomposition / Event Temperatures",
                      **_DARK, height=420, margin=dict(t=55, b=40, l=260))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(gridcolor=None)
    return fig


def electrolyte_ranking_chart(df: pd.DataFrame) -> go.Figure:
    df_s = df.sort_values("Overall Score", ascending=True)
    fig = go.Figure(go.Bar(
        x=df_s["Overall Score"], y=df_s["Electrolyte"], orientation="h",
        marker_color=["#FF4757" if "FEC" in e else "#3A86FF" for e in df_s["Electrolyte"]],
        text=df_s["Overall Score"].apply(lambda v: f"{v:.2f}"), textposition="outside",
    ))
    fig.update_layout(xaxis_title="Overall Score",
                      title="Electrolyte Ranking — Ionic Conductivity, Stability, Decomp. Risk",
                      **_DARK, height=400, margin=dict(t=55, b=40, l=240))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(gridcolor=None)
    return fig


def anode_ranking_chart(df: pd.DataFrame) -> go.Figure:
    df_s = df.sort_values("Overall Score", ascending=True)
    fig = go.Figure(go.Bar(
        x=df_s["Overall Score"], y=df_s["Anode Material"], orientation="h",
        marker_color=px.colors.qualitative.Plotly[:len(df_s)],
        text=df_s["Overall Score"].apply(lambda v: f"{v:.2f}"), textposition="outside",
    ))
    fig.update_layout(xaxis_title="Overall Score",
                      title="Anode Material Ranking — Capacity, Stability, SEI Compatibility",
                      **_DARK, height=380, margin=dict(t=55, b=40, l=200))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(gridcolor=None)
    return fig


def additive_ranking_chart(df: pd.DataFrame) -> go.Figure:
    df_s = df.sort_values("Overall Score", ascending=True)
    colors = ["#FF4757" if a == "FEC (fluoroethylene carbonate)" else "#3A86FF"
              for a in df_s["Additive"]]
    fig = go.Figure(go.Bar(
        x=df_s["Overall Score"], y=df_s["Additive"], orientation="h",
        marker_color=colors,
        text=df_s["Overall Score"].apply(lambda v: f"{v:.2f}"), textposition="outside",
    ))
    fig.update_layout(xaxis_title="Overall Score",
                      title="Electrolyte Additive Ranking — LiF SEI Enhancement & Cycle Life",
                      **_DARK, height=380, margin=dict(t=55, b=40, l=240))
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(gridcolor=None)
    return fig


def multi_objective_pareto(df_anode: pd.DataFrame) -> go.Figure:
    fig = px.scatter(
        df_anode,
        x="Volume Expansion (%)",
        y="Practical Capacity (mAh/g)",
        color="Anode Material",
        size="Cycle Stability Score",
        size_max=40,
        text="Anode Material",
        title="Multi-Objective Trade-off: Capacity vs. Volume Expansion (bubble = cycle stability)",
        color_discrete_sequence=px.colors.qualitative.Plotly,
    )
    fig.update_traces(textposition="top center", textfont_size=10)
    return _apply_dark(fig, 440)


def ion_mobility_plot(df: pd.DataFrame) -> go.Figure:
    colors = {"Inorganic (LiF-rich)": "#FF4757",
              "Mixed inorganic/organic": "#F59E0B",
              "Organic outer layer": "#3A86FF"}
    fig = go.Figure()
    for layer in df["SEI Layer"].unique():
        s = df[df["SEI Layer"] == layer]
        fig.add_trace(go.Scatter(
            x=s["SEI Depth (nm)"], y=s["Li⁺ Diffusivity (cm²/s)"],
            mode="lines", name=layer,
            line=dict(color=colors.get(layer, "#888"), width=2.5),
            fill="tozeroy", fillcolor=colors.get(layer, "#888").replace(")", ",0.08)").replace("rgb", "rgba"),
        ))
    fig.add_hline(y=3.44e-8, line_dash="dot", line_color="#FF4757",
                  annotation_text="LiF bulk D (paper: 3.44×10⁻⁸ cm²/s)",
                  annotation_font_color="#FF4757")
    fig.update_layout(
        xaxis_title="Depth from Anode Surface (nm)",
        yaxis_title="Li⁺ Diffusivity (cm²/s)",
        title="Ion Mobility Map — Li⁺ Diffusivity Through SEI Layers",
        yaxis_type="log",
        **_DARK, height=400, margin=dict(t=55, b=40),
    )
    fig.update_xaxes(**_GRID)
    fig.update_yaxes(**_GRID)
    return fig


def dendrite_risk_plot(df: pd.DataFrame) -> go.Figure:
    fig = px.line(df, x="Current Density (mA/cm²)", y="Overpotential (mV)",
                  color="SEI Condition", title="Dendrite Nucleation Risk — Overpotential vs. Current Density",
                  color_discrete_sequence=["#4ade80", "#F59E0B", "#f87171", "#dc2626"])
    fig.add_hline(y=50, line_dash="dot", line_color="rgba(255,255,255,0.3)",
                  annotation_text="Dendrite nucleation threshold (~50 mV)",
                  annotation_font_color="rgba(255,255,255,0.5)")
    return _apply_dark(fig, 400)


def reaction_pathway_plot(df: pd.DataFrame) -> go.Figure:
    fig = make_subplots(specs=[[{"secondary_y": True}]])
    fig.add_trace(go.Bar(x=df["Step"], y=df["ΔG (eV)"], name="ΔG (eV)",
                         marker_color=["#4ade80" if v < 0 else "#f87171" for v in df["ΔG (eV)"]],
                         text=df["Reaction"].apply(lambda r: r[:30] + "…" if len(r) > 30 else r),
                         textposition="outside"), secondary_y=False)
    fig.add_trace(go.Scatter(x=df["Step"], y=df["Barrier Ea (eV)"], mode="lines+markers",
                             name="Barrier Eₐ (eV)", line=dict(color="#F59E0B", width=2),
                             marker=dict(size=8)), secondary_y=True)
    fig.update_layout(title="FEC Decomposition → LiF: Reaction Pathway Energetics",
                      xaxis_title="Reaction Step", **_DARK, height=420, margin=dict(t=55, b=80))
    fig.update_yaxes(title_text="ΔG (eV)", **_GRID, secondary_y=False)
    fig.update_yaxes(title_text="Barrier Eₐ (eV)", secondary_y=True)
    fig.update_xaxes(**_GRID)
    return fig


def mechanical_spider(df: pd.DataFrame, materials: list) -> go.Figure:
    cats = ["Bulk Modulus (GPa)", "Shear Modulus (GPa)", "Young's Modulus (GPa)", "Fracture Toughness (MPa√m)"]
    fig = go.Figure()
    palette = px.colors.qualitative.Plotly
    for i, mat in enumerate(materials):
        row = df[df["Material"] == mat].iloc[0]
        # Normalise each column to 0-10
        norms = [row["Bulk Modulus (GPa)"] / 120 * 10,
                 row["Shear Modulus (GPa)"] / 60 * 10,
                 row["Young's Modulus (GPa)"] / 130 * 10,
                 row["Fracture Toughness (MPa√m)"] / 1.0 * 10]
        norms_c = norms + [norms[0]]
        cats_c = cats + [cats[0]]
        fig.add_trace(go.Scatterpolar(r=norms_c, theta=cats_c, fill="toself", name=mat,
                                      line=dict(color=palette[i % len(palette)]), opacity=0.7))
    fig.update_layout(polar=dict(radialaxis=dict(visible=True, range=[0, 10], color="#aaa"),
                                 bgcolor="rgba(0,0,0,0)"),
                      title="Mechanical Properties Radar (normalised)", **_DARK,
                      height=420, margin=dict(t=60, b=20))
    return fig


# ════════════════════════════════════════════════════════════════════════════════
# SODIUM-ION BATTERY (SIB) PLOTS
# ════════════════════════════════════════════════════════════════════════════════

_SIB_MAT_COLORS = {
    "NaFePO₄": "#3A86FF",
    "Na₂MnNiO₄": "#FF4757",
    "NaFe0.5Mn0.5PO₄ (predicted)": "#4ade80",
    "Na₂Mn0.5Co0.5O₄ (predicted)": "#F59E0B",
    "NaFe0.25Ni0.75PO₄ (predicted)": "#8B5CF6",
    "NaCoO₂ (reference)": "#aaa",
    "NaMnO₂ (reference)": "#ccc",
    "Na₃V₂(PO₄)₃ (reference)": "#999",
}

_ATOM_COLORS = {
    "Na": "#FFD700",
    "Fe": "#B7410E",
    "P":  "#FF6B35",
    "Mn": "#7B2D8B",
    "Ni": "#2E8B57",
    "O":  "#4FC3F7",
}


def sib_structure_overview(df: pd.DataFrame) -> go.Figure:
    """Grouped bar comparing total energy and formation energy per structure."""
    fig = make_subplots(rows=1, cols=2,
                        subplot_titles=["Total Energy (eV)", "Formation Energy (eV/atom)"],
                        horizontal_spacing=0.12)
    colors = [_SIB_MAT_COLORS.get(c, "#888") for c in df["Composition"]]

    fig.add_trace(go.Bar(x=df["Material"], y=df["Total Energy (eV)"],
                         marker_color=colors, name="Total Energy",
                         text=df["Total Energy (eV)"].apply(lambda v: f"{v:.2f}"),
                         textposition="outside"), row=1, col=1)
    fig.add_trace(go.Bar(x=df["Material"], y=df["Formation Energy (eV/atom)"],
                         marker_color=colors, name="Formation Energy",
                         text=df["Formation Energy (eV/atom)"].apply(lambda v: f"{v:.3f}"),
                         textposition="outside", showlegend=False), row=1, col=2)

    fig.update_layout(barmode="group", showlegend=False, **_DARK,
                      height=430, margin=dict(t=55, b=80))
    fig.update_xaxes(tickangle=-20, **_GRID)
    fig.update_yaxes(**_GRID)
    return fig


def sib_lattice_radar(df: pd.DataFrame, materials: list) -> go.Figure:
    """Radar of normalised lattice parameters for selected structures."""
    cats = ["a (Å)", "b (Å)", "c (Å)"]
    fig = go.Figure()
    palette = px.colors.qualitative.Plotly
    for i, mat in enumerate(materials):
        row = df[df["Material"] == mat].iloc[0]
        vals = [row["a (Å)"], row["b (Å)"], row["c (Å)"]]
        # normalise 0–40 range
        norms = [v / 40 * 10 for v in vals]
        norms_c = norms + [norms[0]]
        cats_c = cats + [cats[0]]
        fig.add_trace(go.Scatterpolar(r=norms_c, theta=cats_c, fill="toself", name=mat,
                                      line=dict(color=palette[i % len(palette)]),
                                      opacity=0.75))
    fig.update_layout(
        polar=dict(radialaxis=dict(visible=True, range=[0, 10], color="#aaa"),
                   bgcolor="rgba(0,0,0,0)"),
        title="Lattice Parameter Comparison (normalised to 0–10)",
        **_DARK, height=400, margin=dict(t=60, b=20),
    )
    return fig


def sib_bader_box(df: pd.DataFrame, material: str) -> go.Figure:
    """Box/strip plot of Bader charges per atom type for one material."""
    sub = df[df["Material"] == material]
    atoms = sub["Atom"].unique()
    fig = go.Figure()
    for atom in atoms:
        s = sub[sub["Atom"] == atom]
        fig.add_trace(go.Box(
            y=s["Charge (e)"], name=atom,
            marker_color=_ATOM_COLORS.get(atom, "#888"),
            boxpoints="all", jitter=0.4, pointpos=0,
            line=dict(width=1.5),
        ))
    fig.update_layout(
        yaxis_title="Bader Charge (e)",
        title=f"Bader Charge Distribution — {material}",
        **_DARK, height=420, margin=dict(t=55, b=40),
    )
    fig.update_xaxes(showgrid=False)
    fig.update_yaxes(**_GRID)
    return fig


def sib_bader_heatmap(df: pd.DataFrame) -> go.Figure:
    """Heatmap of average Bader charges: materials × atom types."""
    pivot = df.groupby(["Material", "Atom"])["Charge (e)"].mean().reset_index()
    pivot_wide = pivot.pivot(index="Material", columns="Atom", values="Charge (e)")
    # Order atoms sensibly
    atom_order = [a for a in ["Na", "Fe", "Mn", "Ni", "P", "O"] if a in pivot_wide.columns]
    pivot_wide = pivot_wide[atom_order]

    fig = go.Figure(go.Heatmap(
        z=pivot_wide.values,
        x=pivot_wide.columns.tolist(),
        y=pivot_wide.index.tolist(),
        colorscale="RdBu",
        zmid=0,
        text=[[f"{v:.3f}" for v in row] for row in pivot_wide.values],
        texttemplate="%{text}",
        textfont=dict(size=11),
        colorbar=dict(title="Charge (e)"),
    ))
    fig.update_layout(
        title="Average Bader Charges — All Structures (e)",
        xaxis_title="Atom Type",
        yaxis_title="Material",
        **_DARK, height=320, margin=dict(t=55, b=40),
    )
    return fig


def sib_charge_transfer_bar(df: pd.DataFrame) -> go.Figure:
    """Bar chart showing charge transfer from each cation to O for each structure."""
    # Compute charge neutrality check: sum of charges per structure
    summary = df.groupby(["Material", "Structure", "Atom"]).agg(
        total_charge=("Charge (e)", lambda x: x.sum()),
        n=("Charge (e)", "count"),
        avg_charge=("Charge (e)", "mean"),
    ).reset_index()
    cations = summary[summary["avg_charge"] > 0].copy()
    cations["label"] = cations["Material"] + "\n" + cations["Atom"]

    fig = px.bar(cations, x="Atom", y="avg_charge", color="Material",
                 barmode="group", text="avg_charge",
                 title="Average Cation Charges — Bader Analysis",
                 color_discrete_map=_SIB_MAT_COLORS,
                 labels={"avg_charge": "Average Charge (e)", "Atom": "Atom Type"})
    fig.update_traces(texttemplate="%{y:.3f}", textposition="outside")
    fig.update_layout(**_DARK, height=400, margin=dict(t=55, b=40))
    fig.update_xaxes(showgrid=False)
    fig.update_yaxes(**_GRID)
    return fig


def sib_formation_energy_chart(df: pd.DataFrame) -> go.Figure:
    """Bar chart of formation energies — DFT values + AI screened."""
    color_map = {"DFT": "#FF4757", "Literature": "#3A86FF",
                 "AI prediction": "#4ade80"}
    fig = go.Figure()
    for src in df["Source"].unique():
        sub = df[df["Source"] == src]
        fig.add_trace(go.Bar(
            x=sub["Material"], y=sub["Formation Energy (eV/atom)"],
            name=src, marker_color=color_map.get(src, "#888"),
            text=sub["Formation Energy (eV/atom)"].apply(lambda v: f"{v:.3f}"),
            textposition="outside",
        ))
    fig.add_hline(y=-2.38, line_dash="dot", line_color="#FF4757",
                  annotation_text="NaFePO₄: −2.38 eV/atom",
                  annotation_font_color="#FF4757")
    fig.add_hline(y=-1.542, line_dash="dot", line_color="#3A86FF",
                  annotation_text="Na₂MnNiO₄: −1.542 eV/atom",
                  annotation_font_color="#3A86FF")
    fig.update_layout(barmode="group",
                      title="Formation Energy — DFT vs AI-Screened Cathodes",
                      yaxis_title="Formation Energy (eV/atom)",
                      **_DARK, height=450, margin=dict(t=55, b=80))
    fig.update_xaxes(tickangle=-25, showgrid=False)
    fig.update_yaxes(**_GRID)
    return fig


def sib_cathode_ranking_chart(df: pd.DataFrame) -> go.Figure:
    """Horizontal bar of AI Score for all cathodes."""
    df_s = df.sort_values("AI Score (0–100)", ascending=True)
    colors = [
        "#FF4757" if s == "DFT" else
        "#4ade80" if s == "AI predicted" else
        "#3A86FF"
        for s in df_s["Source"]
    ]
    fig = go.Figure(go.Bar(
        x=df_s["AI Score (0–100)"],
        y=df_s["Material"],
        orientation="h",
        marker_color=colors,
        text=df_s["AI Score (0–100)"].apply(lambda v: f"{v:.1f}"),
        textposition="outside",
    ))
    fig.update_layout(
        title="AI Cathode Ranking (Stability 25% · Voltage 25% · Diffusion 20% · Capacity 20% · Safety 10%)",
        xaxis_title="AI Score (0–100)",
        **_DARK, height=430, margin=dict(t=60, b=40, l=240),
    )
    fig.update_xaxes(range=[0, 110], **_GRID)
    fig.update_yaxes(gridcolor=None)
    return fig


def sib_cathode_radar(df: pd.DataFrame, materials: list) -> go.Figure:
    """Radar of 5-property scores for selected cathodes."""
    cats = ["Stability Score", "Voltage Score", "Diffusion Score", "Capacity Score", "Safety Score"]
    fig = go.Figure()
    palette = px.colors.qualitative.Plotly
    for i, mat in enumerate(materials):
        row = df[df["Material"] == mat].iloc[0]
        vals = [row[c] for c in cats]
        fig.add_trace(go.Scatterpolar(
            r=vals + [vals[0]], theta=cats + [cats[0]],
            fill="toself", name=mat,
            line=dict(color=palette[i % len(palette)]), opacity=0.75,
        ))
    fig.update_layout(
        polar=dict(radialaxis=dict(visible=True, range=[0, 10], color="#aaa"),
                   bgcolor="rgba(0,0,0,0)"),
        title="Cathode Property Radar",
        **_DARK, height=430, margin=dict(t=60, b=20),
    )
    return fig


def sib_diffusion_bar(df: pd.DataFrame) -> go.Figure:
    """Migration barrier and diffusivity per mechanism."""
    fig = make_subplots(rows=1, cols=2,
                        subplot_titles=["Migration Barrier Eₐ (eV)", "Diffusivity D at 300K (cm²/s)"],
                        horizontal_spacing=0.14)
    colors = [_SIB_MAT_COLORS.get(m, "#888") for m in df["Material"]]
    fig.add_trace(go.Bar(
        x=df["Mechanism"], y=df["Migration Barrier Ea (eV)"],
        marker_color=colors, name="Eₐ",
        text=df["Migration Barrier Ea (eV)"].apply(lambda v: f"{v:.2f}"),
        textposition="outside",
    ), row=1, col=1)
    fig.add_trace(go.Bar(
        x=df["Mechanism"], y=np.log10(df["D at 300K (cm²/s)"]),
        marker_color=colors, name="log₁₀D",
        text=df["D at 300K (cm²/s)"].apply(lambda v: f"{v:.1e}"),
        textposition="outside", showlegend=False,
    ), row=1, col=2)
    fig.update_layout(**_DARK, height=430, margin=dict(t=55, b=80))
    fig.update_xaxes(tickangle=-20, showgrid=False)
    fig.update_yaxes(**_GRID)
    return fig


def sib_screened_bubble(df: pd.DataFrame) -> go.Figure:
    """Bubble chart: predicted voltage vs capacity, sized by |Eform|."""
    fig = px.scatter(
        df, x="Predicted Voltage (V)", y="Predicted Capacity (mAh/g)",
        size=df["Predicted Eform (eV/atom)"].abs(),
        color="Source",
        text="Material",
        size_max=40,
        title="AI-Screened Cathodes: Voltage vs Capacity (bubble = |Eform|)",
        color_discrete_map={"DFT": "#FF4757", "AI": "#4ade80"},
    )
    fig.update_traces(textposition="top center", textfont_size=9)
    return _apply_dark(fig, 460)


def sib_pipeline_status(df: pd.DataFrame) -> go.Figure:
    """Funnel / gantt-style view of the 8 AI pipeline stages."""
    status_colors = {"Done": "#4ade80", "Ready": "#3A86FF",
                     "Trained": "#F59E0B", "Predicted": "#8B5CF6", "Training": "#f87171"}
    colors = [status_colors.get(s, "#888") for s in df["Status"]]
    fig = go.Figure(go.Bar(
        x=df["Stage"], y=[1] * len(df),
        marker_color=colors, name="Stage",
        text=df["Name"], textposition="inside",
        textfont=dict(size=10, color="white"),
    ))
    for i, row in df.iterrows():
        fig.add_annotation(x=row["Stage"], y=1.05,
                           text=f"<b>{row['Status']}</b>",
                           showarrow=False,
                           font=dict(size=9, color=status_colors.get(row["Status"], "#888")))
    fig.update_layout(
        title="AI Pipeline Stages — Status",
        xaxis_title="Stage", yaxis=dict(showticklabels=False, showgrid=False),
        **_DARK, height=280, margin=dict(t=55, b=40),
    )
    fig.update_xaxes(tickvals=list(range(1, 9)),
                     ticktext=[f"S{i}" for i in range(1, 9)], showgrid=False)
    return fig