"""GeoTransolver DrivAerML surface aero — interactive ZeroGPU demo.

Drop an STL of a car-like geometry; the model predicts surface pressure and wall
shear stress, rendered as a colored mesh.

Runtime contract:
- Model assets are pulled from a private HF repo (HF_TOKEN required as a Space secret).
- All CUDA work runs inside ``@spaces.GPU``; the wrapper is cached at module scope
  but instantiated lazily on the first GPU call (DistributedManager.initialize and
  load_checkpoint both touch CUDA, so they must run after ZeroGPU grants a device).
"""

from __future__ import annotations

import os
import re
import tempfile
import threading
import time
from pathlib import Path

import gradio as gr
import numpy as np
import plotly.graph_objects as go
import pyvista as pv
import spaces
import torch
import trimesh
from huggingface_hub import hf_hub_download, snapshot_download

# Multi-model registry. Each entry maps a UI label to a physicsnemo-cfd
# wrapper-registry name (resolved via ``get_model_wrapper`` + the AssetSpec
# defaults registered by ``register_builtin_model_packages``). Per-model load
# kwargs mirror the benchmark matrix in ``workflows/benchmarking/``.
MODELS: dict[str, dict] = {
    "GeoTransolver": {
        "name": "geotransolver_surface",
        "kwargs": {"batch_resolution": 60_000, "geometry_sampling": 300_000, "cuda_bf16_autocast": True},
    },
    "Transolver": {
        "name": "transolver_surface",
        "kwargs": {"batch_resolution": 60_000, "geometry_sampling": 300_000, "cuda_bf16_autocast": True},
    },
    "DoMINO": {
        "name": "domino_surface",
        "kwargs": {},
    },
    "XMGN": {
        "name": "xmgn_surface",
        "kwargs": {"max_points": 250_000, "interpolation_k": 4},
    },
    "FiGNet": {
        "name": "fignet_surface",
        "kwargs": {"max_points": 250_000, "interpolation_k": 4},
    },
}
MODEL_LABELS = list(MODELS.keys())

# Ground-truth drag forces from CFD simulation, one per DrivAerML validation
# run (keys are integer run IDs). Source: CFD simulation reference (Newtons,
# at the DrivAerML reference condition: 30 m/s, ρ = 1.205 kg/m³).
GROUND_TRUTH_DRAG: dict[int, float] = {
    4: 449.87, 17: 428.80, 18: 400.11, 27: 373.99, 62: 341.95, 71: 395.38,
    86: 490.02, 112: 483.59, 115: 588.99, 118: 512.57, 129: 587.64,
    145: 421.17, 149: 496.29, 159: 374.13, 171: 420.45, 183: 359.69,
    197: 474.08, 202: 482.34, 204: 422.69, 225: 372.67, 269: 347.67,
    270: 473.96, 271: 406.48, 298: 376.43, 305: 446.84, 320: 468.97,
    345: 447.11, 367: 494.45, 380: 420.77, 382: 501.06, 399: 475.53,
    409: 396.18, 417: 544.14, 419: 338.72, 421: 466.67, 424: 424.63,
    429: 504.16, 431: 579.65, 439: 590.83, 460: 532.03, 465: 512.24,
    468: 417.57, 469: 499.45, 478: 428.81, 483: 341.44, 489: 393.80,
    490: 460.81, 495: 586.43,
}


def _ground_truth_drag(run_label: str | None) -> float | None:
    """Look up ground-truth drag (N) for a ``run_<idx>`` label; None if unknown."""
    if not run_label:
        return None
    m = re.search(r"\d+", run_label)
    if m is None:
        return None
    return GROUND_TRUTH_DRAG.get(int(m.group(0)))

DRIVAERML_REPO = "neashton/drivaerml"
# DrivAerML validation run IDs (the 48 cases held out of training). Predictions on
# any other geometry are out-of-distribution — that's why the demo only exposes these.
VALIDATION_RUNS = [
    4, 17, 18, 27, 62, 71, 86, 112, 115, 118, 129, 145, 149, 159, 171, 183,
    197, 202, 204, 225, 269, 270, 271, 298, 305, 320, 345, 367, 380, 382,
    399, 409, 417, 419, 421, 424, 429, 431, 439, 460, 465, 468, 469, 478,
    483, 489, 490, 495,
]
RUN_CHOICES = [f"run_{i}" for i in VALIDATION_RUNS]

AIR_DENSITY = 1.205
# DrivAerML training velocity. The neural net predicts non-dimensional fields;
# velocity only sets the post-hoc scaling u² in pressure / shear-stress units.
# Hardcoding here keeps the UI focused on geometry → fields.
STREAM_VELOCITY = 30.0

FIELD_CHOICES = ["Pressure", "WSS magnitude", "WSS x", "WSS y", "WSS z"]

# One colormap across all fields for visual consistency. Diverging RdBu_r centers
# zero (meaningful for signed fields); for the strictly-positive magnitude it just
# uses the warm half of the colormap.
COLORSCALE = "RdBu_r"

FIELD_LABELS = {
    "Pressure": "Pressure  [Pa]",
    "WSS magnitude": "|τ_w|  [Pa]",
    "WSS x": "τ_w,x  [Pa]",
    "WSS y": "τ_w,y  [Pa]",
    "WSS z": "τ_w,z  [Pa]",
}

# Fixed colormap ranges so colors mean the same thing across geometries.
# Calibrated to typical DrivAerML magnitudes.
FIELD_RANGES = {
    "Pressure":      (-720.0, 500.0),
    "WSS magnitude": (   0.0,   5.0),
    "WSS x":         (  -5.0,   3.0),
    "WSS y":         (  -1.5,   1.5),
    "WSS z":         (  -2.0,   2.0),
}

BG_COLOR = "#0b0d12"
FG_COLOR = "#e7ebf2"

CUSTOM_CSS = """
/* Dark theme is the default. Light-mode visitors get the overrides in the
   media query below — same variable names, light-appropriate values, so
   the rest of the CSS works unchanged. */
:root {
    --bg: #0d1015;
    --card: #1a1d24;
    --card-hi: #232831;
    --border: #2a2f3a;
    --text: #f1f3f6;
    --text-muted: #9aa0ad;
    --accent: #76B900;
    --accent-soft: rgba(118, 185, 0, 0.12);
}

@media (prefers-color-scheme: light) {
    /* HF Spaces tends to force Gradio into dark mode (via ?__theme=dark or
       a .dark class on the container), which keeps Gradio's widget variables
       dark even when the user's browser is light. Override our palette AND
       Gradio's widget variables at every scope Gradio uses. */
    :root,
    .gradio-container,
    .gradio-container.dark,
    .dark,
    body.dark {
        --bg: #f5f7fb;
        --card: #ffffff;
        --card-hi: #f0f3f8;
        --border: #d6dae2;
        --text: #0d1015;
        --text-muted: #5a6270;
        --accent: #5C9300;
        --accent-soft: rgba(92, 147, 0, 0.10);
        /* Gradio's own variables — widget chrome reads these directly. */
        --body-background-fill: #f5f7fb;
        --background-fill-primary: #ffffff;
        --background-fill-secondary: #f0f3f8;
        --block-background-fill: #ffffff;
        --input-background-fill: #f0f3f8;
        --body-text-color: #0d1015;
        --body-text-color-subdued: #5a6270;
        --block-label-text-color: #0d1015;
        --block-title-text-color: #0d1015;
        --block-info-text-color: #5a6270;
        --input-text-color: #0d1015;
        --neutral-50: #ffffff;
        --neutral-100: #f5f7fb;
        --neutral-200: #f0f3f8;
        --neutral-300: #d6dae2;
        --checkbox-label-background-fill: #f0f3f8;
        /* Darker tint so the selected radio pill stands out against white. */
        --checkbox-label-background-fill-selected: rgba(92, 147, 0, 0.32);
    }
}

.gradio-container,
gradio-app,
.app,
.main,
.wrap,
.contain {
    max-width: 100% !important;
    width: 100% !important;
    margin: 0 auto !important;
    background: var(--bg) !important;
    color: var(--text) !important;
    color-scheme: light dark;
    padding-left: 1.2rem !important;
    padding-right: 1.2rem !important;
}

footer { display: none !important; }

.app-header {
    width: 100%;
    box-sizing: border-box;
    padding: 1.8rem 1.8rem 1.4rem;
    background: linear-gradient(180deg, var(--card-hi), var(--card));
    border: 1px solid var(--border);
    border-radius: 12px;
    /* Pull outer edges slightly outward so width matches the bottom row's
       combined card span (Gradio's default column padding is wider than our
       input-card / plot-card overrides, which made the header look inset). */
    /* Asymmetric: left edge already aligns; right edge needs extra outward
       push to match the bottom row's combined card span. */
    margin: 0 -1.6rem 1rem -0.9rem;
    position: relative;
    overflow: hidden;
}

.app-header::after {
    content: "";
    position: absolute;
    inset: auto 0 0 0;
    height: 1px;
    background: linear-gradient(90deg, transparent, rgba(118, 185, 0, 0.55), transparent);
}

.app-header h1 {
    color: var(--text) !important;
    font-size: 2rem !important;
    font-weight: 700 !important;
    letter-spacing: -0.015em !important;
    margin: 0 0 0.5rem !important;
}

.app-header .lede {
    color: var(--text-muted) !important;
    font-size: 0.95rem;
    line-height: 1.55;
    max-width: 900px;
    margin: 0;
}

.plot-card {
    background: var(--bg) !important;
    border: 1px solid var(--border) !important;
    border-radius: 14px !important;
    padding: 0.6rem !important;
}

.plot-card .panel-title {
    color: var(--text);
    font-weight: 700;
    font-size: 0.92rem;
    letter-spacing: 0.02em;
    margin: 0.15rem 0 0.5rem 0.3rem;
    text-transform: uppercase;
}

.input-card {
    background: var(--card) !important;
    border: 1px solid var(--border) !important;
    border-radius: 14px !important;
    padding: 0.8rem !important;
}

/* Single distinguished primary action — slightly muted NVIDIA-green tone.
   Saturated #76B900 reads as "loud" on the dark UI; this darker, softer
   variant keeps the brand cue without overpowering the rest of the page. */
button.primary, .primary {
    background: #5C9300 !important;
    background-image: none !important;
    color: #f1f3f6 !important;
    font-weight: 600 !important;
    font-size: 1rem !important;
    letter-spacing: 0.01em !important;
    padding: 0.85rem 1.3rem !important;
    border: 1px solid #3f6500 !important;
    border-radius: 8px !important;
    box-shadow: 0 2px 8px rgba(118, 185, 0, 0.15);
    transition: transform 0.12s ease, box-shadow 0.12s ease, background 0.12s ease;
}
button.primary:hover, .primary:hover {
    background: #6CA800 !important;
    transform: translateY(-1px);
    box-shadow: 0 4px 14px rgba(118, 185, 0, 0.25);
}

.footer-note {
    text-align: center;
    color: var(--text-muted) !important;
    font-size: 0.78rem;
    margin: 1.2rem 0 0.4rem;
    padding-top: 0.9rem;
    border-top: 1px solid var(--border);
}
.footer-note a { color: var(--accent) !important; text-decoration: none; }
.footer-note a:hover { text-decoration: underline; }

.coeff-row {
    display: flex;
    gap: 1rem;
    margin: 0.9rem 0 0.6rem;
    flex-wrap: wrap;
}
.coeff-card {
    flex: 1;
    min-width: 220px;
    padding: 1rem 1.3rem;
    background: var(--card);
    border: 1px solid var(--border);
    border-left: 3px solid var(--accent);
    border-radius: 8px;
    display: flex;
    align-items: baseline;
    justify-content: space-between;
    gap: 1rem;
}
.coeff-card .label {
    color: var(--text-muted) !important;
    font-size: 0.82rem;
    font-weight: 500;
    letter-spacing: 0.01em;
}
.coeff-card .value {
    color: var(--text) !important;
    font-size: 1.7rem;
    font-weight: 600;
    font-family: "JetBrains Mono", "SF Mono", "Menlo", "Consolas", monospace;
    letter-spacing: -0.01em;
    font-variant-numeric: tabular-nums;
}
.coeff-card .symbol {
    color: var(--text-muted) !important;
    font-size: 0.95rem;
    margin-right: 0.35rem;
    font-style: italic;
}
.coeff-card .unit {
    color: var(--text-muted) !important;
    font-size: 1rem;
    font-weight: 500;
    margin-left: 0.3rem;
    letter-spacing: 0;
}

.badge-row {
    display: flex;
    flex-wrap: wrap;
    gap: 0.4rem;
}

/* Big visible "running" banner during inference. ZeroGPU's own spinner is
   too small; we render a centered pulsing banner above the plot until the
   result lands. */
.run-banner {
    margin: 0.4rem 0 0.6rem;
    padding: 1rem 1.4rem;
    border: 1px solid rgba(118, 185, 0, 0.45);
    border-radius: 10px;
    background:
        radial-gradient(ellipse at center, rgba(118,185,0,0.18), rgba(118,185,0,0.04) 70%),
        var(--card);
    color: var(--accent) !important;
    text-align: center;
    font-weight: 600;
    font-size: 1rem;
    letter-spacing: 0.01em;
    animation: run-pulse 1.4s ease-in-out infinite;
}
.run-banner .spin {
    display: inline-block;
    vertical-align: middle;
    font-size: 2rem;
    line-height: 1;
    margin-right: 0.8rem;
    animation: run-spin 1.4s linear infinite;
    transform-origin: center;
}
@keyframes run-pulse {
    0%, 100% { opacity: 1; box-shadow: 0 0 24px rgba(118, 185, 0, 0.12); }
    50%      { opacity: 0.78; box-shadow: 0 0 36px rgba(118, 185, 0, 0.28); }
}
@keyframes run-spin {
    from { transform: rotate(0deg); }
    to   { transform: rotate(360deg); }
}

.results-table {
    margin-top: 0.9rem;
    width: 100%;
    border-collapse: collapse;
    background: var(--card);
    border: 1px solid var(--border);
    border-radius: 8px;
    overflow: hidden;
    font-size: 0.9rem;
}
.results-table th {
    text-align: left;
    padding: 0.55rem 0.95rem;
    background: var(--card-hi) !important;
    color: var(--text-muted) !important;
    font-weight: 600;
    font-size: 0.72rem;
    letter-spacing: 0.08em;
    text-transform: uppercase;
    border-bottom: 1px solid var(--border);
}
.results-table td {
    padding: 0.5rem 0.95rem;
    color: var(--text) !important;
    font-family: "JetBrains Mono", "SF Mono", "Menlo", monospace;
    font-variant-numeric: tabular-nums;
    border-bottom: 1px solid var(--border);
}
.results-table tr:last-child td { border-bottom: none; }
.results-table tr.current td {
    background: rgba(118, 185, 0, 0.06);
    border-left: 2px solid var(--accent);
}
.results-table .run-id {
    color: var(--accent) !important;
    font-weight: 600;
}
.results-table .order {
    color: var(--text-muted) !important;
    font-size: 0.8rem;
    width: 2.5rem;
}
.results-table .model {
    color: #d2b3ff !important;
    font-weight: 500;
}
.results-table .gt {
    color: var(--text-muted) !important;
    font-style: italic;
}
.results-table .err.good { color: #76B900 !important; }
.results-table .err.warn { color: #ffd28a !important; }
.results-table .err.bad  { color: #ff7a7a !important; }

/* Light-mode overrides: the pastel cell colors above (purple model, peach
   warn, salmon bad, light green good) wash out on a white card. Use darker
   saturated variants for adequate contrast. */
@media (prefers-color-scheme: light) {
    .results-table .model      { color: #5b21b6 !important; }
    .results-table .err.good   { color: #3f6500 !important; }
    .results-table .err.warn   { color: #b45309 !important; }
    .results-table .err.bad    { color: #b91c1c !important; }
}
.results-table .err.na, .results-table .gt.na {
    color: var(--text-muted) !important;
    opacity: 0.4;
}
"""

_WRAPPERS: dict[str, object] = {}  # spec_name -> loaded wrapper instance
_PREVIEW_CACHE: dict[int, str] = {}  # run_id -> path to decimated STL
# Guards the lazy first-load of each wrapper (double-checked locking pattern).
# Without this, two concurrent first-clicks on different models in Gradio's
# threaded queue could race the wrapper.load step.
_WRAPPER_LOCK = threading.Lock()


def _download_drivaerml_stl(run_id: int) -> str:
    """Download (or pull from cache) the multi-solid STL for a DrivAerML run,
    then expose it under a stable path with a real ``.stl`` extension.

    pyvista resolves symlinks before sniffing the extension, and the HF Hub
    cache stores the actual blob at an extensionless content-addressed path,
    so the canonical cached symlink ``run_4/drivaer_4.stl`` doesn't survive
    that resolution. We hardlink (or copy as fallback) into a known temp
    directory laid out the way physicsnemo-cfd's STL-finder expects:
    ``<tmp>/run_<i>/drivaer_<i>_single_solid.stl``."""
    blob_path = hf_hub_download(
        repo_id=DRIVAERML_REPO,
        repo_type="dataset",
        filename=f"run_{run_id}/drivaer_{run_id}.stl",
    )
    # Resolve the HF cache symlink to the real blob inode before linking;
    # the HF symlink target is *relative* (``../../../blobs/<hash>``), and
    # ``os.link`` on a relative-symlink source can produce a relative
    # symlink at the destination instead of a true hardlink — that symlink
    # then dereferences from the wrong base directory.
    real_blob = os.path.realpath(blob_path)
    stable_dir = os.path.join(tempfile.gettempdir(), f"drivaerml_run_{run_id}")
    os.makedirs(stable_dir, exist_ok=True)
    stable_path = os.path.join(stable_dir, f"drivaer_{run_id}_single_solid.stl")
    if not (os.path.exists(stable_path) and not os.path.islink(stable_path)):
        # Clean up any prior broken-symlink artifact from a botched run
        if os.path.lexists(stable_path):
            os.unlink(stable_path)
        try:
            os.link(real_blob, stable_path)  # hardlink — instant, zero disk
        except OSError:
            import shutil
            shutil.copyfile(real_blob, stable_path)
    return stable_path


def _build_preview(stl_path: str, run_id: int, target_faces: int = 500_000) -> str:
    """Decimate DrivAerML STL to GLB for the browser viewer. Rotates the mesh
    from CFD Z-up to glTF's Y-up convention before export (otherwise Three.js
    renders the car upside-down)."""
    cached = _PREVIEW_CACHE.get(run_id)
    if cached and os.path.isfile(cached):
        return cached
    mesh = trimesh.load(stl_path, force="mesh")
    mesh.merge_vertices()
    if len(mesh.faces) > target_faces:
        mesh = mesh.simplify_quadric_decimation(face_count=target_faces, aggression=2)
    # Z-up (CFD) → Y-up (glTF). Rotation of −90° about the longitudinal X axis
    # sends (x, y, z) → (x, z, −y), so the original "up" lands on +Y.
    mesh.apply_transform(trimesh.transformations.rotation_matrix(-np.pi / 2, [1, 0, 0]))
    out_path = os.path.join(tempfile.gettempdir(), f"preview_run_{run_id}.glb")
    mesh.export(out_path)
    _PREVIEW_CACHE[run_id] = out_path
    return out_path


def select_run(run_label: str | None):
    """Dropdown change handler: download full STL + build decimated preview.
    Returns (preview path for Model3D, full path for state, cleared prediction).

    Validates ``run_label`` against the closed set ``VALIDATION_RUNS`` so a
    bypassed-UI direct API call can't trigger an HF fetch for an arbitrary
    integer (would still hit a public dataset, but spawns temp dirs / 404s)."""
    if not run_label:
        return None, None, None
    m = re.search(r"\d+", run_label)
    if m is None:
        raise gr.Error(f"Invalid run label: {run_label!r}")
    run_id = int(m.group(0))
    if run_id not in VALIDATION_RUNS:
        raise gr.Error(
            f"run_{run_id} is not in the DrivAerML validation set; "
            f"pick one of the {len(VALIDATION_RUNS)} cases in the dropdown."
        )
    stl_path = _download_drivaerml_stl(run_id)
    preview_path = _build_preview(stl_path, run_id)
    return preview_path, stl_path, None


def _download_model_assets(spec_name: str) -> tuple[str, str, dict[str, str]]:
    """Download checkpoint dir + stats + companion files for a registered
    model, using the **pinned revision** from ``physicsnemo-cfd``'s built-in
    ``AssetSpec``.

    Uses ``snapshot_download`` (not single-file ``hf_hub_download``) so the
    full checkpoint directory materializes — physicsnemo's ``load_checkpoint``
    scans the parent directory and expects both the ``.mdlus`` and its ``.pt``
    companion to be present.

    Pinned SHA: each checkpoint was saved against a specific physicsnemo model
    SHA; the latest weights on ``main`` may not load against the physicsnemo
    SHA pinned in our ``requirements.txt`` (tensor shapes evolve with the
    model code). The AssetSpec roots ``hf://org/repo@sha`` encode the
    compatibility pair the benchmark team validated."""
    from physicsnemo.cfd.evaluation.assets.builtin_packages import (
        register_builtin_model_packages,
    )
    from physicsnemo.cfd.evaluation.assets.registry import get_default_asset

    register_builtin_model_packages()
    spec = get_default_asset(spec_name)
    if spec is None:
        raise RuntimeError(f"No registered AssetSpec for {spec_name!r}")
    if not spec.package_root.startswith("hf://"):
        raise RuntimeError(f"Unsupported package_root: {spec.package_root!r}")
    without_prefix = spec.package_root[len("hf://") :]
    if "@" in without_prefix:
        repo_id, revision = without_prefix.split("@", 1)
    else:
        repo_id, revision = without_prefix, None

    ck_parent = str(Path(spec.checkpoint_relpath).parent)
    if ck_parent in (".", ""):
        # Checkpoint at repo root (XMGN, FiGNet) — pull a tight filter of just
        # the files we actually reference, including any companion files.
        patterns = [spec.checkpoint_relpath, spec.stats_relpath]
        if spec.extra_resolve_relpaths:
            for _, rel_template in spec.extra_resolve_relpaths:
                patterns.append(rel_template.replace("{checkpoint_parent}", "."))
    else:
        # Pull the whole checkpoint subdirectory so load_checkpoint sees the
        # full set of sibling files (.mdlus + .pt + global_stats.json + ...).
        patterns = [f"{ck_parent}/*"]

    snapshot_dir = snapshot_download(
        repo_id=repo_id,
        revision=revision,
        allow_patterns=patterns,
    )
    ckpt_path = str(Path(snapshot_dir) / spec.checkpoint_relpath)
    stats_path = str(Path(snapshot_dir) / spec.stats_relpath)

    load_kw: dict[str, str] = {}
    if spec.extra_resolve_relpaths:
        for kw_name, rel_template in spec.extra_resolve_relpaths:
            rel = rel_template.replace("{checkpoint_parent}", ck_parent)
            load_kw[kw_name] = str(Path(snapshot_dir) / rel)
    return ckpt_path, stats_path, load_kw


def _ensure_wrapper(model_label: str, device: str):
    """Lazy-load a physicsnemo-cfd wrapper for ``model_label`` onto the GPU
    granted by ZeroGPU. One wrapper cached per model name for the lifetime of
    the worker process.

    Uses double-checked locking around the cache so concurrent first-loads of
    different models (or the same model from two queued sessions) don't race
    on the ``wrapper.load`` step."""
    if model_label not in MODELS:
        raise gr.Error(f"Unknown model: {model_label}")
    cfg = MODELS[model_label]
    spec_name = cfg["name"]
    cached = _WRAPPERS.get(spec_name)
    if cached is not None:
        return cached

    with _WRAPPER_LOCK:
        cached = _WRAPPERS.get(spec_name)
        if cached is not None:
            return cached

        import physicsnemo.cfd.evaluation.models.wrappers  # noqa: F401 register built-ins
        from physicsnemo.cfd.evaluation.models.model_registry import get_model_wrapper

        ckpt_path, stats_path, extras_kw = _download_model_assets(spec_name)
        wrapper_cls = get_model_wrapper(spec_name)
        wrapper = wrapper_cls()
        wrapper.load(
            checkpoint_path=ckpt_path,
            stats_path=stats_path,
            device=device,
            inference_domain="surface",
            **cfg["kwargs"],
            **extras_kw,
        )
        _WRAPPERS[spec_name] = wrapper
        return wrapper


def _ensure_placeholder_vtp(stl_path: str, mesh: pv.PolyData) -> str:
    """Save a sibling VTP carrying zero-filled placeholder surface fields.

    DoMINO's ``build_domin_surface_datadict`` reads cell-data fields from the
    file at ``case.mesh_path`` (a VTP) to populate the model's ground-truth
    supervision channel. The values are unused after the forward pass, so
    zeros are fine — but the **arrays must exist with the right names**.
    Field names mirror the DrivAerML DoMINO config:
    ``pMeanTrim`` (scalar) + ``wallShearStressMeanTrim`` (vector).

    Other wrappers ignore extra cell_data, so the synthesized VTP is benign
    everywhere."""
    vtp_path = stl_path.replace("_single_solid.stl", "_boundary.vtp")
    if os.path.exists(vtp_path):
        return vtp_path
    m = mesh.copy()
    n = m.n_cells
    m.cell_data["pMeanTrim"] = np.zeros(n, dtype=np.float32)
    m.cell_data["wallShearStressMeanTrim"] = np.zeros((n, 3), dtype=np.float32)
    m.save(vtp_path)
    return vtp_path


def _stl_to_polydata(stl_path: str) -> pv.PolyData:
    """Load STL via trimesh (validate + fix_normals), drop zero-area triangles,
    and return a PyVista PolyData with inward-pointing normals (the convention
    GeoTransolver was trained on for DrivAerML)."""
    mesh_tm = trimesh.load(stl_path, force="mesh")
    mesh_tm.process(validate=True)
    mesh_tm.fix_normals()
    # Invert face winding so the PolyData's winding-derived normals point
    # inward — matches the DrivAerML VTP convention the models were trained
    # on, and makes ``physicsnemo-cfd``'s ``compute_drag_and_lift`` return
    # positive drag directly (no post-hoc sign flip).
    mesh_tm.invert()

    areas = np.asarray(mesh_tm.area_faces, dtype=np.float32)
    valid = areas > 0.0

    vertices = np.asarray(mesh_tm.vertices, dtype=np.float32)
    faces = np.asarray(mesh_tm.faces, dtype=np.int64)[valid]
    cell_normals = np.asarray(mesh_tm.face_normals, dtype=np.float32)[valid]

    n_tri = faces.shape[0]
    pv_faces = np.empty((n_tri, 4), dtype=np.int64)
    pv_faces[:, 0] = 3
    pv_faces[:, 1:] = faces
    mesh = pv.PolyData(vertices, pv_faces.ravel())
    mesh.cell_data["Normals"] = cell_normals
    mesh.cell_data.active_normals_name = "Normals"
    return mesh


def _plot_mesh(pts: np.ndarray, faces: np.ndarray, point_values: np.ndarray, field_key: str) -> go.Figure:
    label = FIELD_LABELS[field_key]
    lo, hi = FIELD_RANGES[field_key]

    fig = go.Figure(
        go.Mesh3d(
            x=pts[:, 0], y=pts[:, 1], z=pts[:, 2],
            i=faces[:, 0], j=faces[:, 1], k=faces[:, 2],
            intensity=point_values,
            colorscale=COLORSCALE,
            cmin=float(lo), cmax=float(hi),
            colorbar=dict(
                title=dict(text=label, font=dict(color=FG_COLOR, size=13)),
                tickfont=dict(color=FG_COLOR, size=11),
                outlinewidth=0,
                thickness=18,
                len=0.7,
                x=0.98,
            ),
            lighting=dict(
                ambient=0.55, diffuse=0.7, specular=0.35,
                roughness=0.4, fresnel=0.2,
            ),
            lightposition=dict(x=200, y=200, z=300),
            hovertemplate=f"{label}: %{{intensity:.2f}}<extra></extra>",
            showscale=True,
        )
    )
    fig.update_layout(
        template="plotly_dark",
        scene=dict(
            aspectmode="data",
            xaxis=dict(visible=False, showbackground=False),
            yaxis=dict(visible=False, showbackground=False),
            zaxis=dict(visible=False, showbackground=False),
            bgcolor=BG_COLOR,
            camera=dict(eye=dict(x=1.6, y=-1.4, z=0.8)),
        ),
        paper_bgcolor=BG_COLOR,
        plot_bgcolor=BG_COLOR,
        margin=dict(l=0, r=0, t=0, b=0),
        font=dict(color=FG_COLOR, family="Inter, system-ui, sans-serif"),
        height=780,
    )
    return fig


def _cell_to_point(mesh: pv.PolyData, pressure: np.ndarray, wss: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
    """Average per-cell predictions onto vertices for smooth (Gouraud) shading."""
    m = mesh.copy()
    m.cell_data["_pressure"] = pressure
    m.cell_data["_wss"] = wss
    pm = m.cell_data_to_point_data()
    return (
        np.asarray(pm.point_data["_pressure"], dtype=np.float32),
        np.asarray(pm.point_data["_wss"], dtype=np.float32),
    )


def _compute_drag_lift(
    mesh: pv.PolyData, pressure: np.ndarray, wss: np.ndarray
) -> tuple[float, float]:
    """Integrate surface pressure + WSS to drag and lift forces (Newtons),
    delegating to the canonical ``compute_drag_and_lift`` from
    ``physicsnemo-cfd`` so the math is byte-identical to the benchmark.

    Our PolyData carries inward winding (from ``_stl_to_polydata``'s
    ``trimesh.invert()``), which matches the DrivAerML VTP convention; that
    lets the canonical formula return positive drag at stagnation directly,
    without any post-hoc sign flip."""
    from physicsnemo.cfd.postprocessing_tools.metrics.aero_forces import (
        compute_drag_and_lift,
    )

    m = mesh.copy()
    m.cell_data["__pressure"] = pressure
    m.cell_data["__wss"] = wss
    cd, _cd_p, _cd_f, cl, _cl_p, _cl_f = compute_drag_and_lift(
        m,
        pressure_field="__pressure",
        wss_field="__wss",
        coeff=1.0,
        drag_direction=[1.0, 0.0, 0.0],
        lift_direction=[0.0, 0.0, 1.0],
        dtype="cell",
    )
    return float(cd), float(cl)


@spaces.GPU(duration=240)
def run_inference(stl_path: str, model_label: str) -> dict:
    """GPU pass: load the selected model on first call, run the canonical
    physicsnemo-cfd wrapper flow (``prepare_inputs → predict → decode_outputs``),
    and return cacheable state."""
    if stl_path is None:
        raise gr.Error("Select a validation case first.")
    if model_label not in MODELS:
        raise gr.Error(f"Unknown model: {model_label}")
    device = "cuda:0"
    wrapper = _ensure_wrapper(model_label, device)

    from physicsnemo.cfd.evaluation.datasets.schema import CanonicalCase

    t0 = time.perf_counter()
    mesh = _stl_to_polydata(stl_path)
    # DoMINO reads cell_data fields straight from ``case.mesh_path``; point it
    # at a sibling VTP with placeholder zero fields. GeoTransolver/Transolver
    # ignore the VTP (they go through ``reference_geometry``), XMGN/FiGNet
    # don't read it at all, so this is safe across the board.
    vtp_path = _ensure_placeholder_vtp(stl_path, mesh)

    run_id_match = re.search(r"run_(\d+)", stl_path or "")
    run_id = f"run_{run_id_match.group(1)}" if run_id_match else "run_unknown"

    case = CanonicalCase(
        case_id=run_id,
        mesh_path=vtp_path,
        mesh_type="cell",
        inference_domain="surface",
        reference_geometry=mesh,
    )

    model_input = wrapper.prepare_inputs(case)
    raw_output = wrapper.predict(model_input)
    predictions = wrapper.decode_outputs(raw_output, case, model_input)

    pressure = np.asarray(predictions["pressure"], dtype=np.float32)
    wss = np.asarray(predictions["shear_stress"], dtype=np.float32)

    # Unify to cell-located values regardless of model output location. XMGN
    # predicts at vertices; the others predict at cells. Force integration and
    # smooth-shading both want a consistent input.
    if getattr(wrapper, "output_location", "cell") == "point":
        m = mesh.copy()
        m.point_data["_p"] = pressure
        m.point_data["_w"] = wss
        m = m.point_data_to_cell_data()
        pressure = np.asarray(m.cell_data["_p"], dtype=np.float32)
        wss = np.asarray(m.cell_data["_w"], dtype=np.float32)

    F_drag, F_lift = _compute_drag_lift(mesh, pressure, wss)
    point_pressure, point_wss = _cell_to_point(mesh, pressure, wss)
    elapsed = time.perf_counter() - t0

    gpu_name = torch.cuda.get_device_name(0) if torch.cuda.is_available() else "CPU"
    return {
        "run_id": run_id,
        "model": model_label,
        "pts": np.asarray(mesh.points, dtype=np.float32),
        "faces": np.asarray(mesh.regular_faces, dtype=np.int32),
        "point_pressure": point_pressure,
        "point_wss": point_wss,
        "n_cells": int(np.asarray(mesh.regular_faces).shape[0]),
        "elapsed_s": float(elapsed),
        "gpu_name": gpu_name,
        "F_drag": float(F_drag),
        "F_lift": float(F_lift),
    }


def update_history(state: dict | None, history: list[dict]) -> list[dict]:
    """Add or update the (model, run) entry in the session history. Dedupes by
    the ``(model, run)`` tuple so a given geometry can carry one row per model
    — that's how users compare architectures on the same case."""
    if state is None:
        return history
    new_entry = {
        "run": state.get("run_id", "unknown"),
        "model": state.get("model", "?"),
        "F_drag": state["F_drag"],
        "F_lift": state["F_lift"],
    }
    key = (new_entry["model"], new_entry["run"])
    updated = list(history)
    for i, entry in enumerate(updated):
        if (entry.get("model"), entry.get("run")) == key:
            updated[i] = new_entry
            return updated
    updated.append(new_entry)
    return updated


def _err_class(rel: float) -> str:
    """Color hint for the |relative error| column."""
    a = abs(rel)
    if a < 0.10:
        return "err good"
    if a < 0.25:
        return "err warn"
    return "err bad"


def render_history(
    history: list[dict],
    current_run_label: str | None,
    current_model_label: str | None,
) -> str:
    """Render the chronological history table. Each row carries the predicted
    drag and lift plus the CFD ground-truth drag and signed relative error.
    The row matching the currently selected (model, run) pair is highlighted."""
    if not history:
        return ""
    highlight_idx = -1
    for i, entry in enumerate(history):
        if entry.get("run") == current_run_label and entry.get("model") == current_model_label:
            highlight_idx = i
    rows = []
    for i, entry in enumerate(history):
        row_class = " class='current'" if i == highlight_idx else ""
        gt = _ground_truth_drag(entry.get("run"))
        if gt is None:
            gt_cell = "<td class='gt na'>—</td>"
            err_cell = "<td class='err na'>—</td>"
        else:
            rel = (entry["F_drag"] - gt) / gt
            sign = "+" if rel >= 0 else "−"
            err_cell = (
                f"<td class='{_err_class(rel)}'>{sign}{abs(rel) * 100:.1f}%</td>"
            )
            gt_cell = f"<td class='gt'>{gt:.1f}</td>"
        rows.append(
            f"<tr{row_class}>"
            f"<td class='order'>{i + 1}</td>"
            f"<td class='run-id'>{entry.get('run', '?')}</td>"
            f"<td class='model'>{entry.get('model', '?')}</td>"
            f"<td>{entry['F_drag']:.1f}</td>"
            f"{gt_cell}"
            f"{err_cell}"
            "</tr>"
        )
    return (
        "<table class='results-table'>"
        "<thead><tr>"
        "<th>#</th>"
        "<th>Run</th>"
        "<th>Model</th>"
        "<th>F<sub>x</sub> pred &nbsp;(N)</th>"
        "<th>F<sub>x</sub> CFD &nbsp;(N)</th>"
        "<th>Δ</th>"
        "</tr></thead>"
        f"<tbody>{''.join(rows)}</tbody>"
        "</table>"
    )


def render(state: dict | None, field_choice: str):
    """CPU-only re-render from cached prediction."""
    if state is None:
        return None, "Select a run and click **Run inference** to predict surface fields."
    if field_choice == "Pressure":
        values = state["point_pressure"]
    elif field_choice == "WSS magnitude":
        values = np.linalg.norm(state["point_wss"], axis=1)
    else:
        axis = {"WSS x": 0, "WSS y": 1, "WSS z": 2}[field_choice]
        values = state["point_wss"][:, axis]
    fig = _plot_mesh(state["pts"], state["faces"], values, field_choice)

    coeff_html = (
        "<div class='coeff-row'>"
        "<div class='coeff-card'>"
        "<div class='label'>Drag force<span class='symbol'> &nbsp;F<sub>x</sub></span></div>"
        f"<div class='value'>{state['F_drag']:.1f}<span class='unit'> N</span></div>"
        "</div>"
        "<div class='coeff-card'>"
        "<div class='label'>Lift force<span class='symbol'> &nbsp;F<sub>z</sub></span></div>"
        f"<div class='value'>{state['F_lift']:.1f}<span class='unit'> N</span></div>"
        "</div>"
        "</div>"
    )
    badges_html = (
        "<div class='badge-row'>"
        f"<span style='background:#1f2735;color:#a8d24a;padding:3px 9px;border-radius:14px;font-size:12px;margin-right:6px;font-weight:600'>{state.get('run_id', 'unknown')}</span>"
        f"<span style='background:#1f2735;color:#d2b3ff;padding:3px 9px;border-radius:14px;font-size:12px;margin-right:6px;font-weight:600'>{state.get('model', '?')}</span>"
        f"<span style='background:#1f2735;color:#7ee0c0;padding:3px 9px;border-radius:14px;font-size:12px;margin-right:6px'>{state['gpu_name']}</span>"
        f"<span style='background:#1f2735;color:#9bb5ff;padding:3px 9px;border-radius:14px;font-size:12px;margin-right:6px'>{state['n_cells']:,} cells</span>"
        f"<span style='background:#1f2735;color:#ffd28a;padding:3px 9px;border-radius:14px;font-size:12px;margin-right:6px'>{state['elapsed_s']:.1f}s</span>"
        f"<span style='background:#1f2735;color:#e0e0e0;padding:3px 9px;border-radius:14px;font-size:12px'>min {values.min():.1f} · max {values.max():.1f} Pa</span>"
        "</div>"
    )
    return fig, coeff_html + badges_html


# Base theme: NVIDIA-green primary, slate neutral, Inter typography. The
# primary-button visuals are overridden by ``CUSTOM_CSS`` (so the gradient /
# colour tokens are intentionally not set here).
theme = gr.themes.Soft(
    primary_hue="emerald",
    neutral_hue="slate",
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
).set(
    body_background_fill=BG_COLOR,
    body_background_fill_dark=BG_COLOR,
    block_background_fill="#141823",
    block_background_fill_dark="#141823",
    block_border_color="#222a39",
)

with gr.Blocks(title="GeoTransolver — DrivAerML surface aero") as demo:
    gr.HTML(f"<style>{CUSTOM_CSS}</style>")
    with gr.Row():
        with gr.Column():
            gr.HTML(
                '<div class="app-header">'
                '<h1>External Aerodynamics — Neural Surrogates on DrivAerML</h1>'
                '<div class="lede">Five published NVIDIA neural surrogates for external automotive '
                'aerodynamics, served through physicsnemo-cfd\'s benchmark wrappers. '
                'Pick a model and a held-out DrivAerML validation geometry; the demo predicts surface '
                'pressure, wall-shear-stress, and integrated drag and lift forces.</div>'
                '</div>'
            )
    with gr.Row():
        with gr.Column(scale=1, min_width=320, elem_classes="input-card"):
            model_dd = gr.Dropdown(
                choices=MODEL_LABELS,
                value=MODEL_LABELS[0],
                label="Model",
                info="Each architecture is loaded the first time you select it.",
            )
            run_dd = gr.Dropdown(
                choices=RUN_CHOICES,
                value=RUN_CHOICES[0],
                label="DrivAerML validation case",
                info="Pulled on demand from neashton/drivaerml.",
            )
            stl_view = gr.Model3D(
                label="STL geometry",
                display_mode="solid",
                height=360,
                # Side view: 90° around vertical axis, level horizon, auto distance.
                camera_position=(90, 0, None),
                # Pin the viewer backdrop dark in both light/dark modes so the
                # default near-white car mesh stays visible (in light mode it
                # would otherwise blend into a light Three.js background).
                clear_color=(0.10, 0.11, 0.14, 1.0),
            )
            btn = gr.Button("Run inference", variant="primary", size="lg")
            gr.Markdown(
                f"<sub>Fields in Pa, forces in N at the DrivAerML reference condition "
                f"(**{STREAM_VELOCITY:.0f} m/s**, ρ = {AIR_DENSITY} kg/m³).</sub>"
            )
        with gr.Column(scale=3, min_width=720, elem_classes="plot-card"):
            field = gr.Radio(
                FIELD_CHOICES,
                value=FIELD_CHOICES[0],
                label="Field",
                interactive=True,
            )
            inference_status = gr.HTML(value="", visible=False)
            plot = gr.Plot(label=None, show_label=False)
            badges = gr.HTML()
            results_table = gr.HTML()
            gr.Markdown(
                "<sub>For benchmark-grade drag/lift, see "
                "<a href='https://github.com/NVIDIA/physicsnemo-cfd' target='_blank'>"
                "physicsnemo-cfd</a>: forces here are integrated on the STL mesh "
                "used for inference, whereas the official benchmark integrates on "
                "the simulation mesh.</sub>"
            )

    with gr.Row():
        with gr.Column():
            gr.HTML(
                '<div class="footer-note">'
                'To use these architectures, see '
                '<a href="https://github.com/NVIDIA/physicsnemo" target="_blank">NVIDIA/physicsnemo</a>. '
                'To benchmark these checkpoints against your model, see '
                '<a href="https://github.com/NVIDIA/physicsnemo-cfd" target="_blank">NVIDIA/physicsnemo-cfd</a>.'
                '</div>'
            )

    state = gr.State(value=None)
    stl_path_state = gr.State(value=None)
    history_state = gr.State(value=[])

    # Run-inference button is disabled during async work (STL download/decimate,
    # initial page load) so a click while geometry is mid-load can't no-op or
    # race against stale state.
    def _btn_loading_geometry():
        return gr.update(interactive=False, value="Loading geometry…")

    def _btn_ready():
        return gr.update(interactive=True, value="Run inference")

    def _start_inference():
        # Disable button + show big visible "running" banner above the plot so
        # users get strong feedback while the @spaces.GPU function takes its
        # 10–30 seconds; ZeroGPU's own fan icon is too small to be obvious.
        return (
            gr.update(interactive=False, value="Running…"),
            gr.update(
                visible=True,
                value=(
                    "<div class='run-banner'>"
                    "<span class='spin'>⚙</span>"
                    "Running inference on ZeroGPU."
                    "</div>"
                ),
            ),
        )

    def _end_inference():
        return gr.update(interactive=True, value="Run inference"), gr.update(visible=False, value="")

    run_dd.change(_btn_loading_geometry, outputs=btn).then(
        select_run, inputs=run_dd, outputs=[stl_view, stl_path_state, state]
    ).then(
        render_history, inputs=[history_state, run_dd, model_dd], outputs=results_table
    ).then(
        _btn_ready, outputs=btn
    )
    model_dd.change(render_history, inputs=[history_state, run_dd, model_dd], outputs=results_table)
    btn.click(_start_inference, outputs=[btn, inference_status]).then(
        run_inference, inputs=[stl_path_state, model_dd], outputs=state
    ).then(
        update_history, inputs=[state, history_state], outputs=history_state
    ).then(
        render, inputs=[state, field], outputs=[plot, badges]
    ).then(
        render_history, inputs=[history_state, run_dd, model_dd], outputs=results_table
    ).then(
        _end_inference, outputs=[btn, inference_status]
    )
    field.change(render, inputs=[state, field], outputs=[plot, badges])
    # Pre-load the default selection on page open (button disabled until STL ready).
    demo.load(_btn_loading_geometry, outputs=btn).then(
        select_run, inputs=run_dd, outputs=[stl_view, stl_path_state, state]
    ).then(
        _btn_ready, outputs=btn
    )

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