app.py

"""
PoreGCN: Pore-Aware MOF Property Predictor
Gradio web interface with per-atom and per-pore XAI, iRASPA CIF export, and
4-tab output layout.

Author: Abdulmujeeb T. Onawole
Institution: Institute for Molecular Bioscience, The University of Queensland
"""

import os
import sys
import json
import warnings
import traceback
import tempfile
from typing import Dict, List, Optional, Tuple, Any

import numpy as np
import gradio as gr

warnings.filterwarnings("ignore")

# ==============================================================================
# BACKEND IMPORTS (written by parallel agent)
# ==============================================================================

try:
    from config import (
        DATASETS,
        DATASET_LABELS,
        PROPERTY_META,
        CV_THRESHOLD,
        AGREEMENT_THRESHOLD,
        DEVICE,
    )
    from xai_engine import (
        load_ensemble,
        ensemble_predict,
        compute_attributions,
        classify_scenario,
        substructure_breakdown,
    )
    from build_graph import cif_to_graph
    from visualize import (
        create_3d_visualization,
        export_iraspa_cif,
        export_attribution_csv,
    )
    BACKEND_AVAILABLE = True
except ImportError as _backend_err:
    BACKEND_AVAILABLE = False
    _backend_err_msg = str(_backend_err)
    # Graceful stubs so the UI still renders
    DATASETS = ["core_mof", "hmof_geometric", "hmof_gas"]
    DATASET_LABELS = {
        "core_mof": "CoRE MOF (2,737 MOFs, 7 properties)",
        "hmof_geometric": "hMOF Geometric (51,163 MOFs, 5 properties)",
        "hmof_gas": "hMOF Gas (51,163 MOFs, 20 properties)",
    }
    PROPERTY_META = {}
    CV_THRESHOLD = 0.10
    AGREEMENT_THRESHOLD = 0.70
    DEVICE = "cpu"

# ==============================================================================
# ZEO++ BINARY SETUP (compile-on-first-start for Linux / HF Spaces)
# Runs before ensemble loading so pore-node mode is available immediately.
# Falls back to atom-only mode if compilation is skipped or fails.
# ==============================================================================

try:
    from setup_zeopp import ensure_zeopp_binary as _ensure_zeopp
    _zeopp_ok = _ensure_zeopp()
    if _zeopp_ok:
        print("Zeo++ binary ready. Pore-node mode enabled.")
    else:
        print("Zeo++ binary not available. Running in atom-only mode (pore nodes disabled).")
except Exception as _zeopp_setup_err:
    print(f"Zeo++ setup error (non-fatal): {_zeopp_setup_err}")

# ==============================================================================
# PATHS
# ==============================================================================

HERE = os.path.dirname(os.path.abspath(__file__))
EXAMPLES_DIR = os.path.join(HERE, "example_cifs")
MODELS_DIR = os.path.join(HERE, "models")

# ==============================================================================
# ENSEMBLE LOADING AT STARTUP
# ==============================================================================

print("=" * 60)
print("PoreGCN MOF Predictor — Loading ensembles...")
print("=" * 60)

_ENSEMBLES: Dict[str, Any] = {}

if BACKEND_AVAILABLE:
    import json as _json
    from pathlib import Path as _Path
    HERE_PATH = _Path(__file__).parent
    for _ds in DATASETS:
        try:
            _models, _best, _norm, _props = load_ensemble(_ds, DEVICE)
            # Load per-ensemble property means (training-set means used by the
            # XAI agreement check to decide expected direction). Without this
            # the agreement check uses 0 as the reference and silently inverts
            # the expected-direction logic for properties whose prediction is
            # close to or below the training mean.
            _means_path = HERE_PATH / "models" / f"{_ds}_ensemble" / "property_means.json"
            _prop_means: Dict[str, float] = {}
            if _means_path.exists():
                try:
                    _prop_means = _json.loads(_means_path.read_text())
                except Exception as _me:
                    print(f"  [{_ds}] property_means.json parse error: {_me}")
            else:
                print(f"  [{_ds}] WARNING: property_means.json missing")
            _ENSEMBLES[_ds] = {
                "models": _models,
                "best_model": _best,
                "normalizer": _norm,
                "prop_names": _props,
                "property_means": _prop_means,
            }
            print(f"  [{_ds}] Loaded {len(_models)} models, {len(_props)} properties, "
                  f"means for {len(_prop_means)} props")
        except Exception as _e:
            print(f"  [{_ds}] Failed: {_e}")
            _ENSEMBLES[_ds] = None
else:
    print(f"  Backend unavailable: {_backend_err_msg}")

print("=" * 60)


# ==============================================================================
# CUSTOM CSS — Laboratory-notebook aesthetic
# ==============================================================================
# Typography: Syne (display, geometric, distinctive) + IBM Plex Mono (values)
# Palette:
#   --slate:   #0f172a  (primary dark)
#   --slate-m: #1e293b  (mid dark)
#   --stone:   #334155  (secondary text)
#   --mist:    #94a3b8  (tertiary)
#   --edge:    #e2e8f0  (borders)
#   --bg:      #f8fafc  (page background)
#   --teal:    #0891b2  (accent, pore channels)
#   --teal-lt: #cffafe  (teal light)
#   Scenario colors: A=green, B=amber, C=blue, D=red

CUSTOM_CSS = """
@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;500;600;700;800&family=IBM+Plex+Mono:wght@400;500;600&family=IBM+Plex+Sans:ital,wght@0,300;0,400;0,500;0,600;1,400&display=swap');

/* ---- Root tokens ---- */
:root {
    --slate:    #0f172a;
    --slate-m:  #1e293b;
    --stone:    #334155;
    --mist:     #94a3b8;
    --edge:     #e2e8f0;
    --bg:       #f8fafc;
    --white:    #ffffff;
    --teal:     #0891b2;
    --teal-d:   #0e7490;
    --teal-lt:  #e0f2fe;
    --green:    #16a34a;
    --green-lt: #dcfce7;
    --amber:    #d97706;
    --amber-lt: #fef3c7;
    --blue:     #2563eb;
    --blue-lt:  #dbeafe;
    --red:      #dc2626;
    --red-lt:   #fee2e2;
    --mono:     'IBM Plex Mono', 'JetBrains Mono', 'Fira Mono', monospace;
    --sans:     'IBM Plex Sans', -apple-system, BlinkMacSystemFont, sans-serif;
    --display:  'Syne', 'IBM Plex Sans', sans-serif;
}

/* ---- Container ---- */
.gradio-container {
    max-width: 1440px !important;
    margin: 0 auto !important;
    font-family: var(--sans) !important;
    background: var(--bg) !important;
}

/* Graph-paper background texture on body */
body {
    background-color: var(--bg) !important;
    background-image:
        linear-gradient(rgba(8,145,178,0.04) 1px, transparent 1px),
        linear-gradient(90deg, rgba(8,145,178,0.04) 1px, transparent 1px) !important;
    background-size: 28px 28px !important;
}

/* ---- Header ---- */
.poregcn-header {
    background: var(--slate) !important;
    padding: 40px 44px 32px !important;
    border-radius: 4px !important;
    margin-bottom: 20px !important;
    position: relative;
    overflow: hidden;
    border-left: 4px solid var(--teal) !important;
}
.poregcn-header::after {
    content: 'PORE';
    position: absolute;
    right: -12px;
    top: -12px;
    font-family: var(--display);
    font-size: 9em;
    font-weight: 800;
    color: rgba(8,145,178,0.06);
    letter-spacing: -0.05em;
    pointer-events: none;
    line-height: 1;
    user-select: none;
}
.poregcn-header h1 {
    font-family: var(--display) !important;
    font-size: 2.1em !important;
    font-weight: 800 !important;
    color: #ffffff !important;
    margin: 0 0 6px 0 !important;
    letter-spacing: -0.03em !important;
    line-height: 1.1 !important;
}
.poregcn-header h1 span {
    color: var(--teal) !important;
    font-weight: 800 !important;
}
.poregcn-header .subtitle {
    font-family: var(--sans) !important;
    color: #94a3b8 !important;
    font-size: 0.94em !important;
    line-height: 1.65 !important;
    margin: 0 0 20px 0 !important;
    max-width: 640px !important;
    font-weight: 300 !important;
}
.header-badges {
    display: flex;
    gap: 8px;
    flex-wrap: wrap;
}
.hbadge {
    font-family: var(--mono) !important;
    font-size: 0.70em !important;
    font-weight: 500 !important;
    padding: 3px 10px !important;
    border: 1px solid rgba(148,163,184,0.25) !important;
    color: #94a3b8 !important;
    letter-spacing: 0.04em !important;
    border-radius: 2px !important;
    background: rgba(255,255,255,0.04) !important;
}
.hbadge-accent {
    border-color: rgba(8,145,178,0.4) !important;
    color: #67e8f9 !important;
    background: rgba(8,145,178,0.08) !important;
}

/* ---- Input panel ---- */
.input-panel {
    background: var(--white) !important;
    border: 1px solid var(--edge) !important;
    border-radius: 4px !important;
    padding: 24px !important;
}
.panel-label {
    font-family: var(--display) !important;
    font-size: 0.70em !important;
    font-weight: 700 !important;
    text-transform: uppercase !important;
    letter-spacing: 0.12em !important;
    color: var(--mist) !important;
    margin: 0 0 14px 0 !important;
    display: block !important;
    border-bottom: 1px solid var(--edge) !important;
    padding-bottom: 8px !important;
}

/* Run button */
button.primary-btn, .run-btn button {
    background: var(--slate) !important;
    color: #ffffff !important;
    border: none !important;
    border-radius: 2px !important;
    font-family: var(--display) !important;
    font-weight: 700 !important;
    font-size: 0.95em !important;
    letter-spacing: 0.04em !important;
    padding: 14px 28px !important;
    transition: background 0.15s ease !important;
    cursor: pointer !important;
}
button.primary-btn:hover, .run-btn button:hover {
    background: var(--teal-d) !important;
}

/* Status textbox */
.status-box textarea {
    font-family: var(--mono) !important;
    font-size: 0.80em !important;
    color: var(--stone) !important;
    background: #f1f5f9 !important;
    border: 1px solid var(--edge) !important;
    border-radius: 2px !important;
}

/* ---- Scenario pills ---- */
.scenario-pill {
    display: inline-flex;
    align-items: center;
    gap: 5px;
    padding: 3px 10px;
    border-radius: 2px;
    font-family: var(--mono);
    font-size: 0.72em;
    font-weight: 600;
    letter-spacing: 0.06em;
    border: 1px solid;
}
.scenario-A { background: var(--green-lt); color: var(--green); border-color: #86efac; }
.scenario-B { background: var(--amber-lt); color: var(--amber); border-color: #fcd34d; }
.scenario-C { background: var(--blue-lt);  color: var(--blue);  border-color: #93c5fd; }
.scenario-D { background: var(--red-lt);   color: var(--red);   border-color: #fca5a5; }

/* ---- Predictions table ---- */
.pred-table {
    width: 100%;
    border-collapse: separate;
    border-spacing: 0;
    font-size: 0.875em;
    border: 1px solid var(--edge);
    border-radius: 4px;
    overflow: hidden;
}
.pred-table thead th {
    background: var(--slate-m) !important;
    color: #cbd5e1 !important;
    padding: 10px 14px;
    text-align: left;
    font-family: var(--display);
    font-weight: 600;
    font-size: 0.78em;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    border-bottom: 2px solid var(--teal);
}
.pred-table thead th:nth-child(2) { text-align: right; }
.pred-table thead th:nth-child(3) { text-align: center; }
.pred-table thead th:nth-child(4) { text-align: center; }
.pred-table tbody td {
    padding: 9px 14px;
    border-bottom: 1px solid var(--edge);
    color: var(--stone);
    vertical-align: middle;
}
.pred-table tbody tr:last-child td { border-bottom: none; }
.pred-table tbody tr:hover { background: #f1f5f9; }
.pred-table .prop-name {
    font-family: var(--sans);
    font-weight: 500;
    color: var(--slate);
}
.pred-table .prop-desc {
    font-family: var(--sans);
    font-size: 0.82em;
    color: var(--mist);
    display: block;
    margin-top: 1px;
}
.pred-table .pred-val {
    font-family: var(--mono);
    font-weight: 500;
    color: var(--slate);
    text-align: right;
    white-space: nowrap;
}
.pred-table .pred-std {
    font-family: var(--mono);
    font-size: 0.82em;
    color: var(--mist);
    display: block;
    text-align: right;
}
.pred-table .cv-val {
    font-family: var(--mono);
    font-size: 0.85em;
    text-align: center;
}
.cv-low  { color: var(--green); }
.cv-mid  { color: var(--amber); }
.cv-high { color: var(--red); }
.pred-table td.scenario-cell { text-align: center; }

/* ---- Summary bar ---- */
.summary-bar {
    background: var(--slate);
    color: #e2e8f0;
    border-radius: 2px;
    padding: 12px 18px;
    font-family: var(--mono);
    font-size: 0.82em;
    margin-top: 14px;
    line-height: 1.6;
    border-left: 3px solid var(--teal);
}
.summary-bar strong { color: #67e8f9; }

/* ---- Welcome placeholder ---- */
.welcome-state {
    border: 2px dashed var(--edge);
    border-radius: 4px;
    padding: 64px 40px;
    text-align: center;
    background: var(--white);
}
.welcome-state .welcome-title {
    font-family: var(--display);
    font-size: 1.15em;
    font-weight: 700;
    color: #64748b;
    margin: 0 0 8px 0;
}
.welcome-state p {
    font-family: var(--sans);
    color: var(--mist);
    font-size: 0.92em;
    margin: 0;
}

/* ---- Error / warning cards ---- */
.error-card {
    border: 1px solid var(--red);
    border-left: 4px solid var(--red);
    background: var(--red-lt);
    border-radius: 3px;
    padding: 16px 20px;
    color: #991b1b;
    font-family: var(--sans);
    font-size: 0.88em;
    line-height: 1.6;
}
.error-card .err-title {
    font-family: var(--display);
    font-weight: 700;
    font-size: 1.0em;
    margin-bottom: 6px;
}
.warning-card {
    border: 1px solid var(--amber);
    border-left: 4px solid var(--amber);
    background: var(--amber-lt);
    border-radius: 3px;
    padding: 12px 18px;
    color: #92400e;
    font-family: var(--sans);
    font-size: 0.85em;
    margin-bottom: 14px;
}

/* ---- 3D viewer info panel ---- */
.viewer-info-panel {
    background: var(--slate-m);
    color: #94a3b8;
    border-radius: 3px;
    padding: 10px 16px;
    font-family: var(--mono);
    font-size: 0.78em;
    line-height: 1.7;
    margin-bottom: 12px;
    border-left: 3px solid var(--teal);
}
.viewer-info-panel span.highlight { color: #67e8f9; font-weight: 600; }

.top-atoms-bar {
    background: var(--white);
    border: 1px solid var(--edge);
    border-radius: 3px;
    padding: 12px 16px;
    margin-top: 12px;
    font-family: var(--mono);
    font-size: 0.80em;
    color: var(--stone);
    line-height: 1.8;
}
.top-atoms-bar .ta-label {
    font-family: var(--display);
    font-size: 0.72em;
    text-transform: uppercase;
    letter-spacing: 0.1em;
    color: var(--mist);
    display: block;
    margin-bottom: 6px;
}
.atom-pos { color: var(--green); font-weight: 500; }
.atom-neg { color: var(--red); font-weight: 500; }

/* ---- Scenario legend grid ---- */
.scenario-grid {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 10px;
    margin-bottom: 18px;
}
.scenario-cell {
    border-radius: 3px;
    padding: 16px 18px;
    border: 1px solid;
}
.sc-A { background: var(--green-lt); border-color: #86efac; }
.sc-B { background: var(--amber-lt); border-color: #fcd34d; }
.sc-C { background: var(--blue-lt);  border-color: #93c5fd; }
.sc-D { background: var(--red-lt);   border-color: #fca5a5; }
.sc-letter {
    font-family: var(--display);
    font-size: 1.8em;
    font-weight: 800;
    line-height: 1;
    margin-bottom: 4px;
}
.sc-A .sc-letter { color: var(--green); }
.sc-B .sc-letter { color: var(--amber); }
.sc-C .sc-letter { color: var(--blue);  }
.sc-D .sc-letter { color: var(--red);   }
.sc-title {
    font-family: var(--display);
    font-weight: 700;
    font-size: 0.88em;
    margin-bottom: 4px;
}
.sc-desc {
    font-family: var(--sans);
    font-size: 0.80em;
    color: var(--stone);
    line-height: 1.5;
}
.sc-criteria {
    font-family: var(--mono);
    font-size: 0.74em;
    margin-top: 6px;
    opacity: 0.75;
}

/* ---- Substructure breakdown bars ---- */
.breakdown-section {
    background: var(--white);
    border: 1px solid var(--edge);
    border-radius: 3px;
    padding: 18px 20px;
}
.breakdown-title {
    font-family: var(--display);
    font-size: 0.75em;
    font-weight: 700;
    text-transform: uppercase;
    letter-spacing: 0.10em;
    color: var(--mist);
    margin: 0 0 14px 0;
    padding-bottom: 8px;
    border-bottom: 1px solid var(--edge);
}
.breakdown-row {
    margin-bottom: 14px;
}
.breakdown-row:last-child { margin-bottom: 0; }
.breakdown-label {
    display: flex;
    justify-content: space-between;
    align-items: baseline;
    margin-bottom: 5px;
}
.breakdown-name {
    font-family: var(--sans);
    font-weight: 500;
    font-size: 0.88em;
    color: var(--slate);
}
.breakdown-pct {
    font-family: var(--mono);
    font-weight: 600;
    font-size: 0.85em;
    color: var(--stone);
}
.breakdown-track {
    height: 8px;
    background: var(--edge);
    border-radius: 4px;
    overflow: hidden;
}
.breakdown-fill {
    height: 100%;
    border-radius: 4px;
    transition: width 0.4s ease;
}
.fill-metal  { background: #f87171; }
.fill-linker { background: #60a5fa; }
.fill-pore   { background: var(--teal); }

/* ---- Download tab ---- */
.iraspa-instructions {
    background: var(--white);
    border: 1px solid var(--edge);
    border-radius: 3px;
    padding: 20px 24px;
    font-family: var(--sans);
    font-size: 0.88em;
    line-height: 1.75;
    color: var(--stone);
}
.iraspa-instructions .step-label {
    font-family: var(--display);
    font-weight: 700;
    color: var(--slate);
    font-size: 0.88em;
    display: block;
    margin-top: 14px;
    margin-bottom: 2px;
}
.iraspa-instructions code {
    font-family: var(--mono);
    background: #f1f5f9;
    padding: 1px 6px;
    border-radius: 2px;
    font-size: 0.92em;
    color: var(--teal-d);
}

/* ---- Footer ---- */
.poregcn-footer {
    text-align: center;
    font-family: var(--sans);
    font-size: 0.76em;
    color: var(--mist);
    margin-top: 12px;
    padding: 14px 0 8px;
    border-top: 1px solid var(--edge);
    line-height: 1.8;
}
.poregcn-footer a {
    color: var(--teal);
    text-decoration: none;
}
.poregcn-footer a:hover { text-decoration: underline; }

/* ---- Responsive ---- */
@media (max-width: 900px) {
    .poregcn-header { padding: 24px 20px 20px !important; }
    .poregcn-header h1 { font-size: 1.5em !important; }
    .poregcn-header::after { display: none; }
    .scenario-grid { grid-template-columns: 1fr; }
}
"""


# ==============================================================================
# HTML GENERATOR FUNCTIONS
# ==============================================================================

def make_header_html() -> str:
    """Title bar with subtitle and metric badges."""
    return """
    <div class="poregcn-header">
        <h1>Pore<span>GCN</span> &thinsp; MOF Property Predictor</h1>
        <p class="subtitle">
            Heterogeneous graph neural network with Voronoi pore nodes for predicting
            geometric, gas adsorption, and thermal properties of metal-organic frameworks.
            Per-atom and per-pore XAI attributions with iRASPA CIF export.
        </p>
        <div class="header-badges">
            <span class="hbadge hbadge-accent">CoRE MOF &bull; 7 properties</span>
            <span class="hbadge hbadge-accent">hMOF Gas &bull; 20 properties</span>
            <span class="hbadge">Voronoi pore nodes</span>
            <span class="hbadge">Ensemble XAI</span>
            <span class="hbadge">Scenario A/B/C/D trustworthiness</span>
            <span class="hbadge">iRASPA attribution export</span>
        </div>
    </div>
    """


def make_welcome_html() -> str:
    """Placeholder shown before any prediction runs."""
    return """
    <div class="welcome-state">
        <p class="welcome-title">Ready for analysis</p>
        <p>Upload a MOF CIF file on the left, select the ensemble, and click <b>Run Prediction</b>.</p>
        <p style="margin-top: 10px; font-size: 0.85em;">
            Or load one of the example structures below the upload field.
        </p>
    </div>
    """


def _scenario_pill(letter: str, label: str) -> str:
    """Inline scenario pill HTML."""
    return (
        f'<span class="scenario-pill scenario-{letter}" '
        f'aria-label="Scenario {letter}: {label}">'
        f'{letter} &mdash; {label}</span>'
    )


def _cv_class(cv: float) -> str:
    if cv < 0.05:
        return "cv-low"
    if cv < 0.15:
        return "cv-mid"
    return "cv-high"


def make_predictions_table_html(
    predictions: Dict,
    scenarios: Dict,
    prop_names: List[str],
) -> str:
    """
    Sortable HTML table of predictions.
    predictions: {prop: {mean, std, cv}}
    scenarios:   {prop: (letter, label)}
    prop_names:  ordered list to control row order
    """
    if not predictions:
        return make_welcome_html()

    rows = ""
    best_prop, best_cv = None, 1e9
    for prop in prop_names:
        if prop not in predictions:
            continue
        info = predictions[prop]
        mean_val = info.get("mean", float("nan"))
        std_val  = info.get("std",  float("nan"))
        cv_val   = info.get("cv",   float("nan"))
        letter, label = scenarios.get(prop, ("D", "Unreliable"))

        # Track best (most trustworthy) for summary
        if letter == "A" and cv_val < best_cv:
            best_cv = cv_val
            best_prop = prop

        # Property display name from PROPERTY_META if available
        meta = PROPERTY_META.get(prop, {})
        display_name = meta.get("label", prop)
        unit_str = meta.get("unit", "")
        desc_str = meta.get("description", "")

        val_str = f"{mean_val:.4g}" if not (isinstance(mean_val, float) and np.isnan(mean_val)) else "—"
        std_str = f"&plusmn;{std_val:.3g}" if not (isinstance(std_val, float) and np.isnan(std_val)) else ""
        if unit_str:
            val_str += f" <span style='font-size:0.82em;color:#94a3b8;'>{unit_str}</span>"

        cv_pct = f"{cv_val:.1%}" if not (isinstance(cv_val, float) and np.isnan(cv_val)) else "—"
        cv_cls = _cv_class(cv_val) if not (isinstance(cv_val, float) and np.isnan(cv_val)) else ""

        pill = _scenario_pill(letter, label)

        rows += f"""
        <tr>
            <td class="prop-name">
                {display_name}
                {'<span class="prop-desc">' + desc_str + '</span>' if desc_str else ''}
            </td>
            <td class="pred-val">
                {val_str}
                <span class="pred-std">{std_str}</span>
            </td>
            <td class="cv-val {cv_cls}">{cv_pct}</td>
            <td class="scenario-cell">{pill}</td>
        </tr>
        """

    # Summary line
    if best_prop:
        meta = PROPERTY_META.get(best_prop, {})
        best_label = meta.get("label", best_prop)
        best_mean  = predictions[best_prop].get("mean", float("nan"))
        best_std   = predictions[best_prop].get("std",  float("nan"))
        summary_html = (
            f'<div class="summary-bar">'
            f'Most trustworthy: <strong>{best_label}</strong> = '
            f'<strong>{best_mean:.4g} &plusmn; {best_std:.3g}</strong> '
            f'[Scenario A, CV = {best_cv:.1%}]'
            f'</div>'
        )
    else:
        # No Scenario A predictions
        summary_html = (
            '<div class="summary-bar" style="border-color:#f59e0b;">'
            'No Scenario A predictions. All results carry elevated uncertainty — '
            'physical validation is recommended.'
            '</div>'
        )

    return f"""
    <table class="pred-table" role="table" aria-label="Prediction results">
        <thead>
            <tr>
                <th scope="col">Property</th>
                <th scope="col" style="text-align:right;">Prediction (mean &plusmn; std)</th>
                <th scope="col" style="text-align:center;">CV</th>
                <th scope="col" style="text-align:center;">Scenario</th>
            </tr>
        </thead>
        <tbody>{rows}</tbody>
    </table>
    {summary_html}
    """


def make_viewer_info_html(prop_name: str, has_pores: bool = True) -> str:
    """Small info panel above the 3D plot."""
    meta = PROPERTY_META.get(prop_name, {})
    label = meta.get("label", prop_name)
    pore_note = (
        "Pores: translucent yellow void mesh. Cavity attribution: small "
        "orange (positive) or blue (negative) bead at each cavity centre, "
        "labelled with its signed contribution (size scales with magnitude). "
        if has_pores else
        "Voronoi unavailable &mdash; atom contributions only. "
    )
    return (
        f'<div class="viewer-info-panel">'
        f'Atoms colored by attribution to <span class="highlight">{label}</span>. '
        f'{pore_note}'
        f'Blue = negative contribution &bull; Red = positive contribution &bull; '
        f'White = near-zero'
        f'</div>'
    )


def make_top_atoms_html(per_atom, structure_symbols: List[str], n: int = 6) -> str:
    """Top-N atoms by absolute attribution magnitude."""
    if per_atom is None or len(per_atom) == 0 or not structure_symbols:
        return ""

    # per_atom is a list of floats indexed by atom index
    try:
        indexed = [(i, float(a)) for i, a in enumerate(per_atom)]
        indexed.sort(key=lambda x: abs(x[1]), reverse=True)
        top = indexed[:n]
    except Exception:
        return ""

    parts = []
    for idx, val in top:
        sym = structure_symbols[idx] if idx < len(structure_symbols) else "?"
        sign = "+" if val >= 0 else ""
        cls  = "atom-pos" if val >= 0 else "atom-neg"
        parts.append(f'<span class="{cls}">{sym}[{idx}] {sign}{val:.3f}</span>')

    return (
        f'<div class="top-atoms-bar">'
        f'<span class="ta-label">Top contributing atoms (by magnitude)</span>'
        + " &ensp;|&ensp; ".join(parts)
        + f'</div>'
    )


def make_scenario_legend_html() -> str:
    """4-cell colored grid explaining A/B/C/D framework."""
    cv_pct = f"{CV_THRESHOLD:.0%}"
    agr_pct = f"{AGREEMENT_THRESHOLD:.0%}"
    return f"""
    <div class="scenario-grid" role="list" aria-label="Scenario definitions">
        <div class="scenario-cell sc-A" role="listitem">
            <div class="sc-letter">A</div>
            <div class="sc-title">Trustworthy</div>
            <p class="sc-desc">
                Ensemble agrees and XAI attributions align with the physical
                direction of the property.
            </p>
            <div class="sc-criteria">
                CV &lt; {cv_pct} &bull; agreement &ge; {agr_pct}
            </div>
        </div>
        <div class="scenario-cell sc-B" role="listitem">
            <div class="sc-letter">B</div>
            <div class="sc-title">Overconfident</div>
            <p class="sc-desc">
                Ensemble agrees on the value, but attributions do not align
                with physical expectation. Interpret with caution.
            </p>
            <div class="sc-criteria">
                CV &lt; {cv_pct} &bull; agreement &lt; {agr_pct}
            </div>
        </div>
        <div class="scenario-cell sc-C" role="listitem">
            <div class="sc-letter">C</div>
            <div class="sc-title">Underconfident</div>
            <p class="sc-desc">
                XAI attributions are consistent but ensemble variance is high.
                The model is uncertain despite coherent explanations.
            </p>
            <div class="sc-criteria">
                CV &ge; {cv_pct} &bull; agreement &ge; {agr_pct}
            </div>
        </div>
        <div class="scenario-cell sc-D" role="listitem">
            <div class="sc-letter">D</div>
            <div class="sc-title">Unreliable</div>
            <p class="sc-desc">
                High ensemble variance and incoherent attributions. Physical
                validation required before acting on this prediction.
            </p>
            <div class="sc-criteria">
                CV &ge; {cv_pct} &bull; agreement &lt; {agr_pct}
            </div>
        </div>
    </div>
    """


def make_substructure_breakdown_html(breakdown: Dict, prop_name: str = "") -> str:
    """Three horizontal bars: metal % | linker % | pore %."""
    if not breakdown:
        return ""

    meta = PROPERTY_META.get(prop_name, {})
    label = meta.get("label", prop_name) if prop_name else "selected property"

    metal_frac  = breakdown.get("metal_frac",  0.0)
    linker_frac = breakdown.get("linker_frac", 0.0)
    pore_frac   = breakdown.get("pore_frac",   0.0)

    def _bar(frac: float, css_class: str, name: str) -> str:
        pct = max(0, min(100, int(frac * 100)))
        return f"""
        <div class="breakdown-row">
            <div class="breakdown-label">
                <span class="breakdown-name">{name}</span>
                <span class="breakdown-pct">{pct}%</span>
            </div>
            <div class="breakdown-track" role="progressbar"
                 aria-valuenow="{pct}" aria-valuemin="0" aria-valuemax="100"
                 aria-label="{name} attribution {pct}%">
                <div class="breakdown-fill {css_class}"
                     style="width: {pct}%;"></div>
            </div>
        </div>
        """

    title = f"Attribution breakdown &mdash; {label}" if label else "Attribution breakdown"

    return f"""
    <div class="breakdown-section">
        <div class="breakdown-title">{title}</div>
        {_bar(metal_frac,  "fill-metal",  "Metal nodes")}
        {_bar(linker_frac, "fill-linker", "Linker / organic nodes")}
        {_bar(pore_frac,   "fill-pore",   "Pore nodes (Voronoi)")}
    </div>
    """


def make_iraspa_instructions_html() -> str:
    """Step-by-step iRASPA workflow for the Download tab."""
    return """
    <div class="iraspa-instructions">
        <p>
            The CIF file below encodes per-atom XAI attributions in the
            <code>_atom_site_B_iso_or_equiv</code> column (B-factor field).
            Values are scaled from 1 (most negative attribution) through 50 (neutral)
            to 99 (most positive attribution).
        </p>
        <span class="step-label">Step 1 &mdash; Open in iRASPA</span>
        Drag the downloaded CIF into iRASPA, or use
        <code>File &rarr; Open&hellip;</code>.

        <span class="step-label">Step 2 &mdash; Color by attribution</span>
        In the right-hand panel: <code>Appearance &rarr; Atoms &rarr; Color by &rarr;
        Temperature Factor</code>. Select a blue-white-red colormap for clearest contrast.

        <span class="step-label">Step 3 &mdash; Other viewers</span>
        The same B-factor column is read by Mercury, OVITO, ChimeraX, VESTA, and any
        CIF viewer that supports B-factor coloring. In VESTA, use
        <code>Edit &rarr; Color Settings &rarr; Isosurface / B-factors</code>.

        <span class="step-label">Interpretation</span>
        Blue atoms drive the property toward lower values; red atoms drive it higher.
        White atoms contribute near-zero. For properties like void fraction and surface
        area, atoms with high positive attribution are adjacent to the pore interior
        and are strong candidates for chemical modification.
    </div>
    """


def make_ensembles_explainer_html() -> str:
    """About-Ensembles tab content. Explains the three trained ensembles,
    which kind of MOF each is best for, and why hMOF Gas is the default
    recommendation. Reuses the iraspa-instructions/step-label CSS so the
    visual style matches the Download tab."""
    return """
    <div class="iraspa-instructions">
        <p>
            PoreGCN ships three trained ensembles. Each is a 5-fold cross-validation
            top-k selection trained with multi-task learning and inverse-std loss
            weighting. Pick the ensemble whose training distribution most closely
            matches your CIF.
        </p>

        <span class="step-label">hMOF Gas &mdash; default, broadest target list</span>
        Trained on 51,163 hypothetical MOFs from the Wilmer hMOF database.
        Predicts five geometric properties plus 14 gas adsorption capacities
        (CO<sub>2</sub>, N<sub>2</sub>, CH<sub>4</sub>, H<sub>2</sub> across multiple
        pressures and temperatures) plus log<sub>10</sub>(CO<sub>2</sub>/N<sub>2</sub>)
        selectivity. Best for new MOFs whose target use case is gas separation,
        carbon capture, or hydrogen storage. Example MOF in this app:
        <b>HKUST-1</b> (Cu paddlewheel, supports the CO<sub>2</sub>/N<sub>2</sub>
        selectivity narrative).

        <span class="step-label">hMOF Geometric &mdash; lighter, geometry only</span>
        Same training set as hMOF Gas (51,163 hypothetical MOFs) but predicts only
        the five geometric properties: void fraction (VF), gravimetric surface area
        (GSA), accessible surface area (ASA), largest cavity diameter (LCD), and
        pore-limiting diameter (PLD). Use this when you only need pore geometry
        for a screening campaign and want the smaller, faster model. Example MOF
        in this app: <b>MOF-5</b> (canonical high-VF Zn IRMOF, the geometric
        benchmark in most MOF reviews).

        <span class="step-label">CoRE MOF &mdash; experimental structures, stability</span>
        Trained on 2,737 EXPERIMENTAL MOFs from the CoRE MOF 2019 database. Predicts
        ASA, GSA, VF, LCD, PLD, thermal stability, and density. Use this when you
        have a synthesized MOF and care about thermodynamic descriptors that the
        hypothetical-MOF ensembles do not predict. The smaller training set means
        predictions are grounded in real experimental structures, but the chemical
        space coverage is narrower than the hMOF ensembles. Two example MOFs in
        this app are paired with this ensemble. <b>UiO-66</b> is the canonical
        experimental benchmark for stability studies; it lands Scenario A on
        most CoRE MOF properties. <b>Tb-MOF-CrystEngComm2023</b> is a CoRE MOF
        validation example selected from a dual-XAI-method search across all
        2,737 training entries: it lands Scenario A on 5 of 7 properties under
        BOTH attribution methods (signed occlusion and gradient x input), with
        sub-5% relative error on each, while the remaining 2 properties (GSA,
        ASA) are correctly flagged Scenario C and the predictions are 20-41%
        off the ground-truth values. The example concretely demonstrates the
        trustworthiness framework as a working filter, regardless of which
        XAI method the live tool runs.

        <span class="step-label">Why hMOF Gas is the default</span>
        It has the broadest target list and is trained on the largest dataset.
        For a new user uploading a CIF without a specific property in mind, the
        gas-adsorption ensemble returns the most useful per-prediction information.
        Switch to hMOF Geometric if your question is purely geometric and you want
        faster inference; switch to CoRE MOF if your structure is experimental and
        you need stability or density.

        <span class="step-label">A caution on the trustworthiness scenario</span>
        The Trustworthiness tab classifies each prediction A/B/C/D based on
        ensemble agreement and XAI directional consistency. Only Scenario A
        predictions are recommended for downstream screening without further
        validation. The choice of ensemble does not change this classification,
        but a CIF that sits well outside an ensemble's training distribution
        is more likely to fall into Scenario C or D.

        <span class="step-label">XAI method &mdash; Fast vs Slow attribution</span>
        The "XAI method" radio in the input panel selects how per-atom and
        per-pore attributions are computed.
        <br><br>
        <b>The predicted values are the same regardless of method.</b> The
        five-model ensemble produces the same forward pass either way, so
        the property numbers in the table do not change when you flip the
        radio. What changes is how each atom's contribution to the prediction
        is estimated, and that contribution then feeds the Scenario A/B/C/D
        classification. Two methods can therefore agree on the predicted
        value while disagreeing on whether to flag the prediction as
        trustworthy.
        <br><br>
        <b>Fast (default, ~35 sec) &mdash; gradient x input, one pass per property.</b>
        For each property, the model is run forward and backward once.
        Each atom's contribution is read off the gradient of the prediction
        with respect to that atom's input features. This is the model's
        local sensitivity to each atom. Cheap, snappy, and good enough for
        most use cases. Each row in the table gets its own classification.
        <br><br>
        <b>Slow (~3 min on the target property) &mdash; signed occlusion.</b>
        Each atom is removed (its features zeroed) one at a time, the
        model is re-run, and the recorded change in prediction is that
        atom's contribution. Same procedure for each Voronoi pore vertex.
        Slower because the model is re-evaluated once per atom, but the
        contribution is measured directly from how the prediction actually
        changes rather than estimated from a gradient. The target property
        selected in the XAI target dropdown is computed this way; the
        other rows in the table fall back to the Fast surrogate so the
        table stays complete (running Slow across all properties would
        take twenty-plus minutes per click).
        <br><br>
        <b>Why Fast and Slow can disagree on the Scenario column.</b>
        Fast asks "how sensitive is the prediction to each atom's input
        features?". Slow asks "how does the prediction change if this atom
        is not there?". The two questions are related but not identical,
        and they disagree most on properties whose value depends on a
        balance of contributions (e.g. pore-geometry properties where many
        atoms collectively define the cavity shape rather than one atom
        type dominating). When they disagree, Slow is the more conservative
        measurement because it interrogates the prediction directly rather
        than through a gradient approximation.
        <br><br>
        <b>How best to use the tool, in practice.</b>
        <ul style="margin:6px 0 0 18px;padding:0;">
          <li><b>Everyday exploration and presentations:</b> Fast mode.
              All scenarios independently classified in under a minute.</li>
          <li><b>Rigorous classification on a specific property:</b>
              Slow mode, with that property selected in the XAI target
              dropdown. Other rows in the table use Fast to keep the
              table complete.</li>
        </ul>
    </div>
    """


def make_error_html(title: str, detail: str, suggestion: str = "") -> str:
    """Red-bordered error card — no Python traceback exposed."""
    sug_html = (
        f'<p style="margin-top:8px; font-size:0.9em;"><b>Suggestion:</b> {suggestion}</p>'
        if suggestion else ""
    )
    return (
        f'<div class="error-card" role="alert">'
        f'<div class="err-title">{title}</div>'
        f'<p>{detail}</p>{sug_html}'
        f'</div>'
    )


def make_warning_html(message: str) -> str:
    """Amber warning card."""
    return (
        f'<div class="warning-card" role="alert">'
        f'<b>Note:</b> {message}'
        f'</div>'
    )


def make_footer_html() -> str:
    """Author, affiliation, contact, citation, GitHub, license, plus an
    optional Goatcounter visitor-tracking pixel.

    To activate visitor analytics:
      1. Sign up at https://www.goatcounter.com (free, no cookies, GDPR-friendly).
      2. Note the code at the start of your dashboard URL (e.g. for
         https://mujeebonawole.goatcounter.com the code is "mujeebonawole").
      3. On the HF Space, go to Settings -> Variables and Secrets, add a
         **Variable** (not secret) named GOATCOUNTER_CODE with that value.
      4. The pixel below activates on the next page load. Goatcounter dashboard
         shows daily/weekly/monthly visit counts and country breakdown.
    """
    gc_code = os.environ.get('GOATCOUNTER_CODE', '').strip()
    if gc_code:
        # No-JS counting pixel; works even if Gradio sanitises inline scripts
        gc_pixel = (
            f'<img src="https://{gc_code}.goatcounter.com/count'
            f'?p=/poregcn-hf" alt="" referrerpolicy="no-referrer-when-downgrade" '
            f'style="position:absolute;width:1px;height:1px;opacity:0;">'
        )
    else:
        gc_pixel = ''
    return f"""
    <div class="poregcn-footer">
        PoreGCN v1.0 &bull;
        Abdulmujeeb T. Onawole &bull;
        The University of Queensland &bull;
        <a href="mailto:atonawole@gmail.com">atonawole@gmail.com</a> /
        <a href="mailto:a.onawole@uq.edu.qa">a.onawole@uq.edu.qa</a> &bull;
        Trained on CoRE MOF, hMOF Geometric, hMOF Gas &bull;
        <a href="https://github.com/MujeebOnawole/PoreGCN" target="_blank"
           rel="noopener noreferrer">GitHub</a> &bull;
        Onawole, A. T. (2026). <em>PoreGCN: Pore-Aware Graph Neural Network
        for MOF Property Prediction with Explainability.</em> &bull;
        <a href="https://creativecommons.org/licenses/by/4.0/" target="_blank"
           rel="noopener noreferrer">MIT License</a>
        {gc_pixel}
    </div>
    """


# ==============================================================================
# EXAMPLES LOADER
# ==============================================================================

def get_examples() -> List[List]:
    """
    Build gr.Examples input list from example_cifs/ at runtime.
    Each entry: [cif_path, dataset_key]

    Each example is paired with the ensemble it is best suited to, so a new
    user sees one canonical MOF per ensemble before they have to choose:
      - HKUST-1 with hmof_gas: Cu paddlewheel + CO2/N2 selectivity narrative
      - MOF-5 with hmof_geometric: canonical high-VF Zn IRMOF, geometric benchmark
      - UiO-66 with core_mof: experimental, stable, dense Zr MOF
      - Tb-MOF-CrystEngComm2023 with core_mof: CoRE MOF validation example
        selected from a dual-XAI-method search (signed occlusion + gradient
        x input). 5 of 7 properties land Scenario A under both XAI methods
        with sub-5% error on each; the remaining 2 (GSA, ASA) are correctly
        flagged Scenario C and the predictions are 20-41% off the
        ground-truth values. Demonstrates the trustworthiness framework
        as a working filter, regardless of whether the live tool runs the
        Fast (gradient x input) or Full (signed occlusion) attribution.
    Unknown CIFs fall through to alphabetical cycling across the three ensembles.
    """
    if not os.path.isdir(EXAMPLES_DIR):
        return []

    cifs = sorted(f for f in os.listdir(EXAMPLES_DIR) if f.lower().endswith(".cif"))
    if not cifs:
        return []

    KNOWN_MOFS = {
        "HKUST-1.cif": "hmof_gas",
        "MOF-5.cif": "hmof_geometric",
        "UiO-66.cif": "core_mof",
        "Tb-MOF-CrystEngComm2023.cif": "core_mof",
        "ZIF-8.cif": "hmof_gas",
        "Mg-MOF-74.cif": "hmof_gas",
    }
    ds_cycle = ["hmof_gas", "core_mof", "hmof_geometric"]
    examples = []
    fallback_idx = 0
    for cif_name in cifs:
        if cif_name in KNOWN_MOFS:
            ds = KNOWN_MOFS[cif_name]
        else:
            ds = ds_cycle[fallback_idx % len(ds_cycle)]
            fallback_idx += 1
        examples.append([os.path.join(EXAMPLES_DIR, cif_name), ds])
    return examples


# ==============================================================================
# PROPERTY LIST HELPER
# ==============================================================================

def properties_for_dataset(dataset: str) -> List[str]:
    """Return ordered property list for a dataset key."""
    if not BACKEND_AVAILABLE or dataset not in _ENSEMBLES or _ENSEMBLES[dataset] is None:
        # Fallback from static config
        from config import DATASET_PRESETS
        return list(DATASET_PRESETS.get(dataset, {}).get("property_names", []))
    return list(_ENSEMBLES[dataset]["prop_names"])


# ==============================================================================
# CORE PREDICTION FUNCTION
# ==============================================================================

def run_prediction(cif_file, dataset: str, target_prop: str, xai_method: str = "Fast (gradient x input, ~5 sec)", progress=gr.Progress()):
    """
    Main handler: CIF -> graph -> ensemble predict -> XAI -> HTML outputs.

    Returns tuple matching the gr.Blocks outputs list:
        (predictions_html, viewer_info_html, top_atoms_html,
         viewer_html_3d, trust_legend_html, breakdown_html,
         iraspa_path, status_text, warning_html)
    """
    empty = (
        make_welcome_html(), "", "",
        "",
        make_scenario_legend_html(), "",
        None, "Upload a CIF to begin.", ""
    )

    if cif_file is None:
        return empty

    if not BACKEND_AVAILABLE:
        err = make_error_html(
            "Backend unavailable",
            f"The inference modules could not be imported: {_backend_err_msg}",
            "Check that config.py, xai_engine.py, build_graph.py, and visualize.py "
            "are present in the Space root."
        )
        return (err, "", "", "", make_scenario_legend_html(), "", None, None,
                "Error: backend unavailable.", "")

    ens_info = _ENSEMBLES.get(dataset)
    if ens_info is None:
        err = make_error_html(
            "Ensemble not loaded",
            f"The '{dataset}' ensemble failed to load at startup.",
            "Check that model checkpoints are present in models/ "
            "and that the checkpoint manifest is valid."
        )
        return (err, "", "", "", make_scenario_legend_html(), "", None, None,
                "Error: ensemble not loaded.", "")

    warning_parts = []

    try:
        cif_path = cif_file if isinstance(cif_file, str) else str(cif_file)
        progress(0.05, desc="Parsing CIF...")

        # Build graph
        try:
            graph = cif_to_graph(cif_path, dataset)
        except Exception as e:
            err_msg = str(e)
            if "element" in err_msg.lower() or "vocab" in err_msg.lower():
                return (
                    make_error_html(
                        "Unknown element",
                        f"The CIF contains an element not in the vocabulary: {err_msg}",
                        "Try a MOF with common elements (Zr, Cu, Zn, Fe, Co, Al, C, N, O, H)."
                    ),
                    "", "", "", make_scenario_legend_html(), "", None, None,
                    f"Error: {err_msg}", ""
                )
            if "voronoi" in err_msg.lower() or "zeo" in err_msg.lower():
                warning_parts.append(
                    "Voronoi unavailable — running in atom-only mode. "
                    "Pore-dependent properties may be less accurate."
                )
                graph = cif_to_graph(cif_path, dataset)  # retry without pore nodes
            else:
                return (
                    make_error_html(
                        "CIF parse error",
                        f"Could not parse the CIF file: {err_msg}",
                        "Verify the file is a valid CIF (try opening in Mercury or VESTA). "
                        "Consider one of the provided example structures."
                    ),
                    "", "", "", make_scenario_legend_html(), "", None, None,
                    f"Error: {err_msg}", ""
                )

        progress(0.20, desc="Running ensemble...")

        models    = ens_info["models"]
        best_model = ens_info["best_model"]
        normalizer = ens_info["normalizer"]
        prop_names = ens_info["prop_names"]

        # Ensemble prediction
        predictions = ensemble_predict(graph, models, normalizer, DEVICE)

        progress(0.45, desc="Computing XAI attributions...")

        # Default to first property if target not in list
        if target_prop not in prop_names:
            target_prop = prop_names[0]

        # Per-property means for direction classification. Loaded once at
        # startup from the ensemble's property_means.json (training-set means).
        # The XAI agreement check uses these to determine the expected sign
        # of attributions: if prediction > mean, expect positive attributions
        # on contributing atoms, otherwise expect negative. Defaulting to 0
        # silently inverts the check for properties below the training mean.
        prop_means: Dict[str, float] = {}
        if BACKEND_AVAILABLE and dataset in _ENSEMBLES and _ENSEMBLES[dataset]:
            prop_means = _ENSEMBLES[dataset].get("property_means", {}) or {}

        # Pick the XAI method from the radio.
        #   "Fast" -- gradient x input, run for ALL properties so every
        #            scenario in the table is classified on its own XAI
        #            (~5 seconds per property, ~35 seconds total).
        #   "Full" -- signed occlusion (manuscript method) on the TARGET
        #            property only (~3 minutes); the other six properties
        #            still get scenario classifications, computed via the
        #            fast surrogate so the table is complete.
        # Radio label may read "Slow (signed occlusion, ...)" or
        # "Fast (gradient x input, ...)". Anything starting with slow/full
        # routes to the signed-occlusion branch; default is fast.
        _xm_lc = (xai_method or "").lower().strip()
        method_key = "full" if _xm_lc.startswith(("slow", "full")) else "fast"

        n_props = len(prop_names)

        # Per-property XAI cache: prop -> dict from compute_attributions
        per_prop_xai: Dict[str, Dict] = {}

        if method_key == "full":
            # Target property: full signed occlusion
            def _xai_progress(step, total, msg):
                frac = 0.30 + 0.40 * (step / max(total, 1))
                progress(min(frac, 0.70), desc=msg)
            progress(0.30, desc="Full XAI on target (signed occlusion, ~3 minutes)...")
            target_res = compute_attributions(
                graph, best_model, normalizer, prop_means, target_prop, DEVICE,
                xai_method='full', progress_callback=_xai_progress,
            )
            per_prop_xai[target_prop] = target_res
            # Other properties: fast surrogate so the rest of the table is filled in
            other_props = [p for p in prop_names if p != target_prop]
            for i, prop in enumerate(other_props, 1):
                progress(0.70 + 0.10 * i / max(len(other_props), 1),
                         desc=f"Fast XAI for {prop} ({i}/{len(other_props)})...")
                per_prop_xai[prop] = compute_attributions(
                    graph, best_model, normalizer, prop_means, prop, DEVICE,
                    xai_method='fast',
                )
        else:
            # Fast mode: gradient x input for every property
            for i, prop in enumerate(prop_names, 1):
                progress(0.30 + 0.40 * i / max(n_props, 1),
                         desc=f"Fast XAI for {prop} ({i}/{n_props})...")
                per_prop_xai[prop] = compute_attributions(
                    graph, best_model, normalizer, prop_means, prop, DEVICE,
                    xai_method='fast',
                )

        # 3D viewer needs the target property's per-atom and per-pore
        target_xai = per_prop_xai.get(target_prop, {})
        per_atom = target_xai.get("per_atom", [])
        per_pore = target_xai.get("per_pore", [])
        exp_dir = target_xai.get("expected_direction", "+")

        progress(0.80, desc="Classifying scenarios...")

        # Each property's scenario uses ITS OWN XAI agreement values. This
        # matches the manuscript's offline xai.py pipeline (which runs
        # check_xai_agreement once per property) instead of recycling the
        # target's agreement_frac across the whole table as the live tool
        # used to do.
        scenarios: Dict[str, Tuple[str, str]] = {}
        for prop in prop_names:
            pinfo = predictions.get(prop, {})
            prop_mean = prop_means.get(prop, float("nan"))
            p_xai = per_prop_xai.get(prop, {})
            letter, s_label = classify_scenario(
                pinfo.get("mean", float("nan")),
                pinfo.get("std", float("nan")),
                float(p_xai.get("agreement_frac", 0.0)),
                prop_mean,
                mean_signed=float(p_xai.get("mean_signed", 0.0)),
                expected_positive=bool(p_xai.get("expected_positive", True)),
            )
            scenarios[prop] = (letter, s_label)
        agr_frac = float(target_xai.get("agreement_frac", 0.0))

        progress(0.70, desc="Building predictions table...")

        predictions_html = make_predictions_table_html(predictions, scenarios, prop_names)
        warning_html     = "".join(make_warning_html(w) for w in warning_parts)

        progress(0.78, desc="Computing substructure breakdown...")

        breakdown    = substructure_breakdown(graph, per_atom, per_pore)
        breakdown_html = make_substructure_breakdown_html(breakdown, target_prop)

        progress(0.84, desc="Rendering 3D visualization...")

        # Extract structure metadata for visualizer
        structure = graph.get("structure")   # ASE Atoms or pymatgen Structure
        pore_pos  = graph.get("pore_positions", [])
        pore_radii = graph.get("pore_radii",   [])
        has_pores  = len(pore_pos) > 0

        try:
            viewer_html_3d = create_3d_visualization(
                structure, per_atom, per_pore, pore_pos, pore_radii, target_prop
            )
        except Exception:
            viewer_html_3d = ""

        viewer_info_html = make_viewer_info_html(target_prop, has_pores)

        # Top contributing atoms
        symbols = []
        try:
            if hasattr(structure, "get_chemical_symbols"):
                symbols = structure.get_chemical_symbols()
            elif hasattr(structure, "species"):
                symbols = [str(s) for s in structure.species]
        except Exception:
            pass
        top_atoms_html = make_top_atoms_html(per_atom, symbols)

        progress(0.92, desc="Exporting iRASPA CIF...")

        iraspa_path = None
        try:
            tmp_dir = tempfile.mkdtemp()
            formula = "mof"
            if structure is not None:
                try:
                    if hasattr(structure, "get_chemical_formula"):
                        formula = structure.get_chemical_formula()
                    elif hasattr(structure, "composition"):
                        formula = structure.composition.reduced_formula
                except Exception:
                    pass
            out_path = os.path.join(tmp_dir, f"{formula}_attribution.cif")
            iraspa_path = export_iraspa_cif(structure, per_atom, out_path)
        except Exception:
            iraspa_path = None

        # CSV export of per-atom and per-pore attributions for downstream analysis
        try:
            csv_path = os.path.join(tmp_dir, f"{formula}_{target_prop}_attribution.csv")
            csv_out_path = export_attribution_csv(
                structure=structure,
                per_atom_attrs=per_atom,
                pore_positions=np.asarray(graph.get("pore_positions", np.zeros((0, 3)))),
                pore_radii=np.asarray(graph.get("pore_radii", np.zeros(0))),
                per_pore_attrs=per_pore if per_pore is not None else np.zeros(0),
                property_name=target_prop,
                output_path=csv_path,
            )
        except Exception:
            csv_out_path = None

        progress(1.0, desc="Done.")

        # Status summary
        n_atoms = len(per_atom) if per_atom is not None else 0
        n_pores = len(per_pore) if per_pore is not None else 0
        scenario_A_props = [p for p, (l, _) in scenarios.items() if l == "A"]
        status = (
            f"{n_atoms} atoms, {n_pores} pores | "
            f"{len(prop_names)} properties | "
            f"Scenario A: {len(scenario_A_props)}/{len(prop_names)} | "
            f"XAI target: {target_prop}"
        )

        # Trust legend (always shown, not dependent on result)
        trust_legend_html = make_scenario_legend_html()

        return (
            predictions_html,
            viewer_info_html,
            top_atoms_html,
            viewer_html_3d,
            trust_legend_html,
            breakdown_html,
            iraspa_path,
            csv_out_path,
            status,
            warning_html,
        )

    except Exception as e:
        traceback.print_exc()
        return (
            make_error_html(
                "Prediction failed",
                f"An unexpected error occurred during prediction.",
                "Check the structure is a valid MOF CIF. "
                "Try one of the provided example structures."
            ),
            "", "", "",
            make_scenario_legend_html(), "",
            None, None,
            f"Error: {type(e).__name__}",
            ""
        )


# ==============================================================================
# DATASET -> PROPERTY DROPDOWN UPDATE
# ==============================================================================

def update_property_dropdown(dataset: str) -> gr.Dropdown:
    """Refresh the property dropdown when the dataset changes."""
    props = properties_for_dataset(dataset)
    first = props[0] if props else None
    return gr.Dropdown(choices=props, value=first)


# ==============================================================================
# GRADIO INTERFACE
# ==============================================================================

_dataset_choices = [(DATASET_LABELS.get(k, k), k) for k in DATASETS]
_default_dataset = "hmof_gas"
_default_props   = properties_for_dataset(_default_dataset)
_default_prop    = _default_props[0] if _default_props else None

with gr.Blocks(
    css=CUSTOM_CSS,
    title="PoreGCN: MOF Property Predictor",
    theme=gr.themes.Base(
        font=[gr.themes.GoogleFont("IBM Plex Sans"), "sans-serif"],
        font_mono=[gr.themes.GoogleFont("IBM Plex Mono"), "monospace"],
    ),
) as demo:

    # --- Header ---
    gr.HTML(make_header_html())

    # --- Main two-column layout ---
    with gr.Row(equal_height=False):

        # ---- LEFT: inputs ----
        with gr.Column(scale=1, min_width=300, elem_classes=["input-panel"]):

            gr.HTML('<span class="panel-label">Input</span>')

            cif_input = gr.File(
                label="MOF CIF file",
                file_types=[".cif"],
                type="filepath",
            )

            dataset_dropdown = gr.Dropdown(
                choices=_dataset_choices,
                value=_default_dataset,
                label="Ensemble",
                info="Select the trained ensemble to use for prediction.",
            )

            property_dropdown = gr.Dropdown(
                choices=_default_props,
                value=_default_prop,
                label="XAI target property",
                info="Property for which atom-level attributions are computed.",
            )

            xai_method_radio = gr.Radio(
                choices=[
                    "Fast (gradient x input, ~35 sec)",
                    "Slow (signed occlusion, ~3 min on target property)",
                ],
                value="Fast (gradient x input, ~35 sec)",
                label="XAI method",
                info=(
                    "Predicted values are the same in both modes. "
                    "Only the per-atom attributions and resulting Scenario "
                    "classifications differ. Fast: snappy, suitable for "
                    "everyday use. Slow: more rigorous on the selected "
                    "target property."
                ),
            )

            run_btn = gr.Button(
                "Run Prediction",
                variant="primary",
                size="lg",
                elem_classes=["run-btn"],
            )

            status_box = gr.Textbox(
                label="Status",
                value="Upload a CIF to begin.",
                interactive=False,
                elem_classes=["status-box"],
                lines=2,
            )

            # Warning area (Voronoi missing, etc.)
            warning_out = gr.HTML(value="")

            # Examples — generated at runtime
            _examples = get_examples()
            if _examples:
                gr.HTML(value=(
                    '<div style="font-size:0.85em; color:#475569; margin-top:8px; '
                    'padding:8px 12px; background:#f8fafc; border-left:3px solid #0891b2;">'
                    '<b>Tutorial:</b> click an example below to load a famous MOF, then press <b>Run Prediction</b>. '
                    'HKUST-1 (Cu paddlewheel) is the recommended starting point and reproduces the '
                    'manuscript\'s Cu-attribution finding for CO₂/N₂ selectivity.'
                    '</div>'
                ))
                gr.Examples(
                    examples=_examples,
                    inputs=[cif_input, dataset_dropdown],
                    label="Example MOF structures (click to load)",
                    examples_per_page=5,
                )

        # ---- RIGHT: output tabs ----
        with gr.Column(scale=3):

            with gr.Tabs():

                # ---- TAB 1: Predictions ----
                with gr.TabItem("Predictions", id="tab_predictions"):
                    predictions_out = gr.HTML(value=make_welcome_html())

                # ---- TAB 2: 3D Viewer ----
                with gr.TabItem("3D Attribution Viewer", id="tab_viewer"):
                    viewer_info_out  = gr.HTML(value="")
                    viewer_plot_out  = gr.HTML(value="")
                    top_atoms_out    = gr.HTML(value="")

                # ---- TAB 3: Trustworthiness ----
                with gr.TabItem("Trustworthiness", id="tab_trust"):
                    gr.HTML(
                        '<p style="font-family:\'IBM Plex Sans\',sans-serif;'
                        'font-size:0.88em;color:#334155;margin:0 0 14px 0;'
                        'line-height:1.7;">'
                        'Each prediction is classified by two criteria: '
                        '(1) ensemble coefficient of variation (CV) measures model agreement; '
                        '(2) XAI directional agreement measures whether per-atom attributions '
                        'align with the expected physical direction of the property.'
                        '</p>'
                    )
                    trust_legend_out    = gr.HTML(value=make_scenario_legend_html())
                    breakdown_out       = gr.HTML(value="")

                # ---- TAB: About Ensembles and XAI ----
                with gr.TabItem("About Ensembles and XAI", id="tab_ensembles"):
                    gr.HTML(make_ensembles_explainer_html())

                # ---- TAB 4: Download ----
                with gr.TabItem("Download", id="tab_download"):
                    gr.HTML(make_iraspa_instructions_html())
                    iraspa_out = gr.File(
                        label="iRASPA-ready CIF (B-factor = attribution)",
                        visible=True,
                    )
                    gr.HTML(
                        '<div style="margin-top:12px;padding:10px 14px;background:#f8fafc;'
                        'border-left:3px solid #0891b2;font-family:\'IBM Plex Sans\',sans-serif;'
                        'font-size:0.88em;line-height:1.6;color:#334155;">'
                        '<b>Per-atom and per-pore attribution (CSV):</b> '
                        'tabular export of every atom and Voronoi pore vertex with its '
                        'Cartesian coordinates and signed attribution to the selected property. '
                        'Use this for downstream analysis (ranking high-attribution motifs, '
                        'feeding into MD or DFT on selected atom subsets, or as input to '
                        'design tools that propose linker substitutions).'
                        '</div>'
                    )
                    csv_out = gr.File(
                        label="Attribution CSV (per-atom + per-pore)",
                        visible=True,
                    )

    # --- Footer ---
    gr.HTML(make_footer_html())

    # ==============================================================================
    # EVENT WIRING
    # ==============================================================================

    # Dataset change -> refresh property dropdown
    dataset_dropdown.change(
        fn=update_property_dropdown,
        inputs=[dataset_dropdown],
        outputs=[property_dropdown],
    )

    # Run button
    run_btn.click(
        fn=run_prediction,
        inputs=[cif_input, dataset_dropdown, property_dropdown, xai_method_radio],
        outputs=[
            predictions_out,   # tab 1
            viewer_info_out,   # tab 2
            top_atoms_out,     # tab 2
            viewer_plot_out,   # tab 2
            trust_legend_out,  # tab 3
            breakdown_out,     # tab 3
            iraspa_out,        # tab 4
            csv_out,           # tab 4 (attribution CSV)
            status_box,        # left panel
            warning_out,       # left panel (warning card)
        ],
    )

    # Also trigger on CIF file upload (auto-run when example is clicked)
    cif_input.change(
        fn=run_prediction,
        inputs=[cif_input, dataset_dropdown, property_dropdown],
        outputs=[
            predictions_out,
            viewer_info_out,
            top_atoms_out,
            viewer_plot_out,
            trust_legend_out,
            breakdown_out,
            iraspa_out,
            status_box,
            warning_out,
        ],
    )


# ==============================================================================
# LAUNCH
# ==============================================================================

if __name__ == "__main__":
    demo.launch(
        share=False,
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
    )