pages/visualizations.py

from __future__ import annotations

from pathlib import Path

import dash
import numpy as np
from dash import Input, Output, State, callback, clientside_callback, ctx, dcc, html, no_update

from analysis.viz.umap_plot import apply_feature_umap_highlight
from analysis.viz.vis_heatmaps_plotly import (
    TOGGLE_HEATMAP_OVERLAY_CLIENTSIDE,
    empty_heatmap_figure,
    heatmap_overlay_checklist_props,
    set_heatmap_overlay_visible,
)
from dashboard.image_utils import get_top_feature_overlays, overlay_show_heatmap_enabled
from dashboard.layout import page_back_nav
from dashboard.model_catalog import (
    dashboard_dataset_label,
    get_model_record,
    summarize_model_record,
    summarize_selector_cache,
)
from dashboard.umap_service import (
    AGGREGATION_MODE_LABELS,
    AGGREGATION_MODES,
    DEFAULT_AGGREGATION_MODE,
    DEFAULT_COLOR_MODE,
    FEATURE_COLOR_MODE,
    FEATURE_UMAP_FILTERED_HELP,
    FEATURE_UMAP_GARBAGE_HELP,
    IMAGE_COLOR_MODE_LABELS,
    IMAGE_COLOR_MODES,
    THUMB_SIZE,
    UMAP_DATASET,
    build_features_umap,
    build_images_umap,
    empty_umap_figure,
    format_feature_umap_status,
    format_umap_status,
)
from dashboard.viz_helpers import (
    UMAP_TOOLTIP_CLASSNAME,
    UMAP_TOOLTIP_ZINDEX,
    build_feature_detail_content,
    build_feature_umap_hover_tooltip,
    cached_hover_meta_line,
    umap_tooltip_clientside,
    cached_hover_thumb,
    feature_empty_state,
    feature_ids,
    feature_rows_for_selector,
    feature_unavailable_message,
    hover_image_idx,
    select_feature_id,
    selector_option_label,
    visualization_params,
)

ROOT = Path(__file__).resolve().parents[1]
FEATURE_TOP_IMAGE_SLOTS = 1
FEATURE_TOP_IMAGE_TOP_N = 5
HEATMAP_GRAPH_CONFIG = {
    "displayModeBar": False,
    "responsive": True,
    "scrollZoom": False,
    "doubleClick": "reset+autosize",
}

dash.register_page(__name__, path="/visualizations", name="Visualizations", title="SAE Explorer")


def _umap_images_cache_key(
    metric: str,
    model_key: str,
    aggregation_mode: str,
    color_mode: str,
    selected_category: str | None,
) -> str:
    category = "" if selected_category in {None, ""} else str(selected_category)
    return f"{metric}:{model_key}:{aggregation_mode}:{color_mode}:{category}"


_HYPERPARAM_HELP_ROWS: tuple[tuple[str, str], ...] = (
    ("layer_num", "IQA layer used as SAE input"),
    ("swin_num", "Swin stage (MANIQA only)"),
    # ("liqe_variant", "LIQE backbone variant"),
    ("lambda_param", "sparsity penalty in the SAE loss"),
    ("scaling_factor", "activation scaling before the SAE"),
    ("n_epochs", "number of training epochs"),
    ("sae_type", "sparse autoencoder architecture(sae - standard ReLU SAE, mp_sae - Matching Pursuit SAE)"),
    ("l0_loss", "mean active SAE features per validation image"),
)


def _label_desc_line(label: str, description: str, *, tag: str = "div") -> html.Div | html.Li:
    line = [
        html.Strong(label),
        html.Span(f" — {description}", className="help-desc"),
    ]
    if tag == "li":
        return html.Li(line)
    return html.Div(line, className="help-desc-line")


def _sae_hyperparams_help() -> html.Div:
    return html.Div(
        [
            html.Span("?", className="upload-help viz-sae-help-trigger", **{"aria-label": "SAE hyperparameters help"}),
            html.Div(
                [
                    html.Div("Hyperparameters (from training logs):", className="viz-help-popover-title"),
                    *[
                        _label_desc_line(label, desc)
                        for label, desc in _HYPERPARAM_HELP_ROWS
                    ],
                ],
                className="viz-sae-help-popover",
            ),
        ],
        className="viz-sae-help-wrap",
    )


def _umap_help() -> html.Div:
    return html.Div(
        [
            html.Span("?", className="upload-help viz-sae-help-trigger", **{"aria-label": "UMAP help"}),
            html.Div(
                [
                    html.Div("UMAP views", className="viz-help-popover-title"),
                    html.Div("Image UMAP", className="viz-help-popover-subtitle"),
                    html.Div(
                        f"Each point is one image from {UMAP_DATASET}: aggregated SAE activations "
                        "projected to 2D (UMAP).",
                        className="viz-help-popover-lead",
                    ),
                    _label_desc_line(
                        "Mean activation",
                        "aggregate patch activations per image with the mean over non-zero values.",
                    ),
                    _label_desc_line("Maximum", "aggregate with the per-image maximum activation."),
                    _label_desc_line("Sum", "aggregate with the per-image sum of activations."),
                    _label_desc_line(
                        "Distortion group",
                        "color points by distortion group.",
                    ),
                    _label_desc_line(
                        "Distortion type",
                        "color points by distortion type.",
                    ),
                    html.Div("Feature UMAP", className="viz-help-popover-subtitle"),
                    html.Div(
                        "Each point is one sparse SAE feature (latent direction) after filtering, "
                        "projected to 2D (UMAP).",
                        className="viz-help-popover-lead",
                    ),
                    _label_desc_line(
                        "Distortion group",
                        "color each feature by the distortion group with the strongest patch-level |correlation|.",
                    ),
                    _label_desc_line("filtered", FEATURE_UMAP_FILTERED_HELP),
                    _label_desc_line("garbage", FEATURE_UMAP_GARBAGE_HELP),
                    html.Div(
                        "filtered and garbage points are hidden by default; click their legend entries to show them.",
                        className="viz-help-popover-lead",
                    ),
                ],
                className="viz-sae-help-popover viz-umap-help-popover",
            ),
        ],
        className="viz-sae-help-wrap viz-umap-help-wrap",
    )


def _feature_workbench_help() -> html.Div:
    return html.Div(
        [
            html.Span(
                "?",
                className="upload-help viz-sae-help-trigger",
                **{"aria-label": "Feature workbench help"},
            ),
            html.Div(
                [
                    html.Div("Feature workbench", className="viz-help-popover-title"),
                    html.Div("Feature sort selector", className="viz-help-popover-subtitle"),
                    html.Div(
                        "Picks which precomputed ranking defines the ordered feature list "
                        "for this model and dataset (correlation, ROC-AUC, precision,"
                        "recall and similar metrics).",
                        className="viz-help-popover-lead",
                    ),
                    html.Div("Feature navigation", className="viz-help-popover-subtitle"),
                    _label_desc_line(
                        "Prev / Next",
                        "step through features in the active selector ranking.",
                    ),
                    _label_desc_line(
                        "Jump to feature id",
                        "go directly to a global SAE feature id (even if it is not top-ranked).",
                    ),
                    _label_desc_line(
                        "Feature UMAP click",
                        "in Features UMAP mode, clicking a point selects that feature id.",
                    ),
                    html.Div("Image ranking", className="viz-help-popover-subtitle"),
                    _label_desc_line(
                        "IoU",
                        "top images with the largest overlap between the feature heatmap and "
                        "dataset distortion masks.",
                    ),
                    _label_desc_line(
                        "Maximum activation",
                        "top images with the highest peak SAE activation for this feature.",
                    ),
                ],
                className="viz-sae-help-popover viz-feature-workbench-help-popover",
            ),
        ],
        className="viz-sae-help-wrap viz-feature-workbench-help-wrap",
    )


def _feature_filter_intro_block() -> html.Div:
    return html.Div(
        [
            html.P(
                "Cached activations are filtered before further analysis.",
                className="viz-filter-intro-lead",
            ),
            html.P("Available filters:", className="viz-filter-intro-lead"),
            html.Ul(
                [
                    _label_desc_line(
                        "nonzero_max",
                        "removes features whose maximum activation is 0 (they never fire on any image).",
                        tag="li",
                    ),
                    _label_desc_line(
                        "kruskal_wallis",
                        "keeps features with statistically significant differences across "
                        "distortion groups (Kruskal–Wallis test on activations).",
                        tag="li",
                    ),
                ],
                className="viz-filter-intro-list",
            ),
            html.P(
                "Each line below is one applied filter step and reports before / after / removed feature counts.",
                className="viz-filter-intro-lead",
            ),
        ],
        className="viz-filter-intro",
    )


def _format_filter_result_line(line: str) -> html.Div:
    if ":" not in line:
        return html.Div(line, className="viz-filter-line")
    name, rest = line.split(":", 1)
    return html.Div(
        [
            html.Strong(f"{name.strip()}:"),
            html.Span(rest, className="help-desc"),
        ],
        className="viz-filter-line",
    )


def _skeleton_lines(n: int = 3, short_last: bool = True) -> list[html.Div]:
    lines = []
    for index in range(n):
        class_name = "skeleton-line"
        if short_last and index == n - 1:
            class_name += " skeleton-line-short"
        lines.append(html.Div(className=class_name))
    return lines


def _feature_image_slot(slot_idx: int) -> html.Div:
    return html.Div(
        [
            dcc.Graph(
                id=f"feature-image-{slot_idx}",
                figure=empty_heatmap_figure("Loading..."),
                config=HEATMAP_GRAPH_CONFIG,
                className="heatmap-graph",
            ),
            html.Div(
                "Top images overlay is not available yet.",
                id=f"feature-image-caption-{slot_idx}",
                className="feature-image-caption",
            ),
        ],
        id=f"feature-image-card-{slot_idx}",
        className="feature-image-card",
    )


def _empty_feature_image_cache(message: str) -> tuple[object, str, str]:
    return None, message, "Top images overlay is not available yet."


layout = html.Div(
    [
        page_back_nav(),
        html.Div(
            [html.H1("SAE Explorer", className="app-title")],
            className="hero-card",
        ),
        html.Div(
            [
                html.Div(
                    [
                        html.Div("Selected SAE", className="panel-label"),
                        _sae_hyperparams_help(),
                        dcc.Loading(
                            id="viz-model-info-loading",
                            type="default",
                            color="#2563eb",
                            className="loading-block viz-model-info-loading",
                            children=html.Div(
                                id="viz-model-info",
                                className="mock-summary-lines",
                                children=_skeleton_lines(4),
                            ),
                        ),
                    ],
                    className="mock-card viz-sae-card",
                ),
                html.Div(
                    [
                        html.Div("UMAP SAE vector", className="panel-label"),
                        _umap_help(),
                        html.Div(
                            [
                                html.Div(
                                    [
                                        html.Div(
                                            [
                                                html.Div("UMAP mode", className="umap-mode-label"),
                                                dcc.RadioItems(
                                                    id="umap-mode-dropdown",
                                                    options=[
                                                        {"label": "Images", "value": "images"},
                                                        {"label": "Features", "value": "features"},
                                                    ],
                                                    value="images",
                                                    className="umap-mode-switch",
                                                    labelClassName="umap-mode-pill",
                                                    inputClassName="umap-mode-input",
                                                ),
                                            ],
                                            className="umap-mode-primary",
                                        ),
                                        html.Div(
                                            [
                                                html.Div(
                                                    [
                                                        html.Div("Aggregation mode", className="section-label"),
                                                        dcc.Dropdown(
                                                            id="aggregation-dropdown",
                                                            options=[
                                                                {
                                                                    "label": AGGREGATION_MODE_LABELS[mode],
                                                                    "value": mode,
                                                                }
                                                                for mode in AGGREGATION_MODES
                                                            ],
                                                            value=DEFAULT_AGGREGATION_MODE,
                                                            clearable=False,
                                                            className="dash-dropdown model-dropdown",
                                                        ),
                                                    ],
                                                    id="aggregation-control-block",
                                                    className="umap-control-block",
                                                ),
                                                html.Div(
                                                    [
                                                        html.Div("Color mode", className="section-label"),
                                                        dcc.Dropdown(
                                                            id="color-dropdown",
                                                            options=[
                                                                {
                                                                    "label": IMAGE_COLOR_MODE_LABELS[mode],
                                                                    "value": mode,
                                                                }
                                                                for mode in IMAGE_COLOR_MODES
                                                            ],
                                                            value=DEFAULT_COLOR_MODE,
                                                            clearable=False,
                                                            className="dash-dropdown model-dropdown",
                                                        ),
                                                    ],
                                                    id="color-control-block",
                                                    className="umap-control-block",
                                                ),
                                                html.Div(
                                                    [],
                                                    id="feature-color-source-control-block",
                                                    className="umap-control-block",
                                                    style={"display": "none"},
                                                ),
                                            ],
                                            className="umap-controls-row umap-controls-secondary",
                                        ),
                                    ],
                                    className="umap-controls",
                                ),
                                html.Div(
                                    [
                                        html.Div(id="umap-status", className="umap-status"),
                                        dcc.Loading(
                                            id="umap-graph-loading",
                                            type="circle",
                                            color="#2563eb",
                                            className="loading-block umap-graph-loading",
                                            children=dcc.Graph(
                                                id="umap-graph",
                                                figure=empty_umap_figure("Loading UMAP..."),
                                                clear_on_unhover=True,
                                                config={"displayModeBar": False, "responsive": True},
                                                className="umap-graph",
                                                style={"height": "min(58vh, 640px)", "minHeight": "430px"},
                                            ),
                                        ),
                                        dcc.Tooltip(
                                            id="umap-tooltip",
                                            className=UMAP_TOOLTIP_CLASSNAME,
                                            zindex=UMAP_TOOLTIP_ZINDEX,
                                            direction="right",
                                        ),
                                    ],
                                    className="umap-hover-shell",
                                ),
                                dcc.Store(id="category-filter-store", data=None),
                                dcc.Store(id="umap-figures-cache", data=None),
                            ],
                            className="mock-placeholder",
                        ),
                    ],
                    className="placeholder-card viz-umap-card",
                ),
            ],
            className="mock-grid",
        ),
        html.Div(
            [
                html.Div(
                    [
                        html.Div("Feature workbench", className="panel-label"),
                        _feature_workbench_help(),
                    ],
                    className="feature-workbench-title-row",
                ),
                html.Div(
                    [
                        html.Div(
                            [
                                html.Div(
                                    [
                                        html.Div("Feature sort selector", className="section-label"),
                                        dcc.Dropdown(
                                            id="feature-sort-selector",
                                            options=[],
                                            value=None,
                                            clearable=False,
                                            disabled=True,
                                            className="dash-dropdown model-dropdown",
                                        ),
                                    ],
                                    className="feature-workbench-controls feature-workbench-controls-sort",
                                ),
                                html.Div(
                                    [
                                        html.Div("Feature navigation", className="section-label"),
                                        html.Div(
                                            [
                                                html.Button(
                                                    "Prev", id="feature-prev", n_clicks=0, className="nav-button"
                                                ),
                                                dcc.Input(
                                                    id="feature-jump-input",
                                                    type="text",
                                                    inputMode="numeric",
                                                    pattern="[0-9]*",
                                                    debounce=True,
                                                    className="feature-jump-input",
                                                    value="",
                                                ),
                                                html.Button(
                                                    "Next", id="feature-next", n_clicks=0, className="nav-button"
                                                ),
                                            ],
                                            className="feature-nav-row",
                                        ),
                                    ],
                                    className="feature-workbench-controls feature-workbench-controls-nav",
                                ),
                                html.Div(
                                    [
                                        html.Div("Image ranking", className="section-label"),
                                        dcc.Dropdown(
                                            id="feature-image-ranking",
                                            options=[
                                                {"label": "IoU", "value": "iou"},
                                                {"label": "Maximum activation", "value": "activation"},
                                            ],
                                            value="iou",
                                            clearable=False,
                                            className="dash-dropdown feature-image-ranking-dropdown",
                                        ),
                                    ],
                                    className="feature-workbench-controls feature-workbench-controls-ranking",
                                ),
                            ],
                            className="feature-workbench-toolbar",
                        ),
                        html.Div(id="feature-status", className="feature-status"),
                    ],
                    className="feature-workbench-header",
                ),
                html.Div(
                    id="feature-detail-card",
                    className="feature-detail-card",
                    children=[
                        dcc.Loading(
                            id="feature-detail-body-loading",
                            type="default",
                            color="#2563eb",
                            className="loading-block feature-detail-body-loading",
                            children=html.Div(
                                id="feature-detail-body",
                                children=[
                                    html.Div("Loading feature details…", className="feature-empty"),
                                ],
                            ),
                        ),
                        html.Div(
                            [
                                html.Div(
                                    [
                                        html.Div("Top images", className="feature-placeholder-title"),
                                        dcc.Checklist(
                                            id="feature-show-heatmap-overlay",
                                            **heatmap_overlay_checklist_props(),
                                        ),
                                    ],
                                    className="feature-top-images-header",
                                ),
                                dcc.Store(id="feature-top-images-cache", data=None),
                                dcc.Loading(
                                    id="feature-images-loading",
                                    type="dot",
                                    color="#2563eb",
                                    className="loading-block feature-images-loading",
                                    children=[
                                        html.Div(
                                            "No feature selected.",
                                            id="feature-top-images-status",
                                            className="feature-empty",
                                        ),
                                        html.Div(
                                            id="feature-top-images",
                                            className="feature-top-images-grid",
                                            children=[_feature_image_slot(idx) for idx in range(FEATURE_TOP_IMAGE_SLOTS)],
                                        ),
                                    ],
                                ),
                            ],
                            className="feature-placeholder",
                        ),
                    ],
                ),
                dcc.Store(id="selected-feature-id", data=None),
            ],
            className="feature-workbench",
        ),
    ],
    className="app-shell",
)


@callback(
    Output("viz-model-info", "children"),
    Input("_pages_location", "search"),
)
def render_model_info(search: str | None):
    params = visualization_params(search)
    if params is None:
        return html.Div("Missing visualization parameters.", className="feature-empty")

    metric, selection_dataset, model_key = params
    record = get_model_record(metric, model_key)

    children = [
        html.Div([html.Span("Metric: "), html.Strong(metric)]),
        html.Div([html.Span("Dataset: "), html.Strong(dashboard_dataset_label(selection_dataset))]),
        html.Div([html.Span("Image UMAP embeddings: "), html.Strong(UMAP_DATASET)]),
    ]
    if record is None:
        children.append(html.Div("SAE checkpoint not found.", className="mock-summary-meta"))
        return children

    children.append(
        html.Div(summarize_model_record(record), className="mock-summary-hyperparams")
    )
    try:
        from dashboard.model_catalog import summarize_feature_filter_cache

        filter_lines = summarize_feature_filter_cache(record.checkpoint_path, selection_dataset)
    except Exception:
        filter_lines = ["Feature filter summary unavailable"]

    children.append(html.Hr(className="viz-summary-divider"))
    children.append(
        html.Div(
            [
                _feature_filter_intro_block(),
                *[_format_filter_result_line(line) for line in filter_lines],
            ],
            className="viz-filter-summary",
        )
    )
    return children


@callback(
    Output("feature-sort-selector", "options"),
    Output("feature-sort-selector", "value"),
    Output("feature-sort-selector", "disabled"),
    Input("_pages_location", "search"),
)
def init_feature_selector(search: str | None):
    params = visualization_params(search)
    if params is None:
        return [], None, True

    metric, selection_dataset, model_key = params
    record = get_model_record(metric, model_key)
    if record is None:
        return [], None, True

    summaries = summarize_selector_cache(record.checkpoint_path, selection_dataset)
    if not summaries:
        return [], None, True

    options = [
        {"label": selector_option_label(summary.selector_name), "value": summary.selector_name}
        for summary in summaries
    ]
    return options, summaries[0].selector_name, False


@callback(
    Output("aggregation-control-block", "style"),
    Output("color-control-block", "style"),
    Input("umap-mode-dropdown", "value"),
)
def toggle_umap_secondary_controls(umap_mode: str | None):
    if str(umap_mode or "images") == "features":
        return (
            {"display": "none"},
            {"display": "none"},
        )
    return ({}, {})


@callback(
    Output("category-filter-store", "data"),
    Input("umap-graph", "clickData"),
    State("category-filter-store", "data"),
    State("umap-mode-dropdown", "value"),
    prevent_initial_call=True,
)
def update_category_filter(click_data, current_category, umap_mode):
    if str(umap_mode or "images") != "images" or not click_data:
        return no_update

    point = click_data["points"][0]
    customdata = point.get("customdata", None)
    clicked_category = None
    if isinstance(customdata, (list, tuple, np.ndarray)) and len(customdata) >= 2:
        clicked_category = str(customdata[1])

    if not clicked_category:
        return no_update

    if current_category is not None and str(current_category) == clicked_category:
        return None

    return clicked_category


@callback(
    Output("umap-figures-cache", "data"),
    Input("_pages_location", "search"),
    Input("aggregation-dropdown", "value"),
    Input("color-dropdown", "value"),
    Input("category-filter-store", "data"),
    State("umap-figures-cache", "data"),
    running=[
        (Output("umap-status", "className"), "umap-status is-loading", "umap-status"),
        (Output("umap-status", "children"), "Loading UMAP embeddings…", ""),
    ],
)
def preload_umap_figures(search, aggregation_mode, color_mode, selected_category, prev_cache):
    params = visualization_params(search)
    if params is None:
        return None

    metric, _selection_dataset, model_key = params
    if not model_key or aggregation_mode is None:
        return no_update

    aggregation_mode = str(aggregation_mode)
    images_color_mode = str(color_mode or DEFAULT_COLOR_MODE)
    if images_color_mode not in IMAGE_COLOR_MODES:
        images_color_mode = DEFAULT_COLOR_MODE
    prev_cache = prev_cache or {}
    cache_key = _umap_images_cache_key(
        metric, model_key, aggregation_mode, images_color_mode, selected_category
    )
    prev_images = prev_cache.get("images") if isinstance(prev_cache.get("images"), dict) else None
    if prev_images and prev_images.get("cache_key") == cache_key:
        return no_update

    try:
        images_fig, images_info = build_images_umap(
            metric,
            model_key,
            UMAP_DATASET,
            aggregation_mode,
            color_mode=images_color_mode,
            selected_category=selected_category,
        )
        images_payload = {
            "figure": images_fig.to_plotly_json(),
            "status": format_umap_status(images_info, aggregation_mode, images_color_mode),
            "cache_key": cache_key,
        }
        return {
            "images": images_payload,
            "features": prev_cache.get("features") if ctx.triggered_id == "category-filter-store" else None,
        }
    except Exception as exc:
        error_fig = empty_umap_figure(f"UMAP unavailable: {exc}").to_plotly_json()
        message = f"UMAP unavailable: {exc}"
        error_payload = {"figure": error_fig, "status": message, "cache_key": cache_key}
        return {"images": error_payload, "features": prev_cache.get("features")}


@callback(
    Output("umap-figures-cache", "data", allow_duplicate=True),
    Input("umap-figures-cache", "data"),
    Input("umap-mode-dropdown", "value"),
    State("_pages_location", "search"),
    State("aggregation-dropdown", "value"),
    prevent_initial_call=True,
    running=[
        (Output("umap-status", "className"), "umap-status is-loading", "umap-status"),
        (Output("umap-status", "children"), "Preloading feature UMAP embeddings…", ""),
    ],
)
def preload_feature_umap(cache, umap_mode, search, aggregation_mode):
    # Eager preload: runs when images cache is ready (first page load), not only on mode switch.
    params = visualization_params(search)
    if params is None:
        return no_update

    metric, _selection_dataset, model_key = params
    if not model_key or aggregation_mode is None:
        return no_update

    if not cache or not cache.get("images"):
        return no_update

    if ctx.triggered_id == "umap-mode-dropdown" and str(umap_mode or "images") != "features":
        return no_update

    aggregation_mode = str(aggregation_mode)
    feature_cache = cache.get("features") if isinstance(cache.get("features"), dict) else {}
    if feature_cache.get("figure"):
        return no_update

    try:
        features_fig, features_info = build_features_umap(
            metric,
            model_key,
            UMAP_DATASET,
            aggregation_mode,
            hide_filtered_garbage=True,
        )
        features_payload = {
            "figure": features_fig.to_plotly_json(),
            "status": format_feature_umap_status(features_info, "group"),
        }
        next_cache = dict(cache)
        feature_cache = dict(feature_cache)
        feature_cache.update(features_payload)
        next_cache["features"] = feature_cache
        return next_cache
    except Exception as exc:
        error_fig = empty_umap_figure(f"UMAP unavailable: {exc}").to_plotly_json()
        message = f"UMAP unavailable: {exc}"
        next_cache = dict(cache)
        feature_cache = dict(feature_cache)
        feature_cache.update({"figure": error_fig, "status": message})
        return {
            "images": cache["images"],
            "features": feature_cache,
        }


def _parse_selected_feature_id(value) -> int | None:
    if value is None:
        return None
    try:
        return int(value)
    except (TypeError, ValueError):
        return None


@callback(
    Output("umap-graph", "figure"),
    Output("umap-status", "children"),
    Output("umap-status", "className"),
    Input("umap-mode-dropdown", "value"),
    Input("umap-figures-cache", "data"),
    Input("selected-feature-id", "data"),
)
def display_umap_from_cache(umap_mode, cache, selected_feature_id):
    if not cache:
        return empty_umap_figure("Loading UMAP..."), "Preloading UMAP embeddings...", "umap-status is-loading"

    mode = str(umap_mode or "images")
    if mode == "features":
        payload = cache.get("features")
        if payload is None:
            return (
                empty_umap_figure("Loading feature UMAP..."),
                "Preloading feature UMAP embeddings...",
                "umap-status is-loading",
            )
        if not payload:
            return empty_umap_figure("UMAP unavailable."), "UMAP unavailable.", "umap-status"
        figure = payload.get("figure")
        highlighted = apply_feature_umap_highlight(figure, _parse_selected_feature_id(selected_feature_id))
        return highlighted, payload.get("status", ""), "umap-status"

    payload = cache.get(mode)
    if payload is None and mode == "images":
        return (
            empty_umap_figure("Loading UMAP..."),
            "Preloading UMAP embeddings...",
            "umap-status is-loading",
        )
    if payload is None:
        return empty_umap_figure("UMAP unavailable."), "UMAP unavailable.", "umap-status"

    if not payload:
        return empty_umap_figure("UMAP unavailable."), "UMAP unavailable.", "umap-status"

    return payload["figure"], payload.get("status", ""), "umap-status"


clientside_callback(
    umap_tooltip_clientside("umap-graph"),
    Output("umap-tooltip", "show"),
    Output("umap-tooltip", "bbox"),
    Output("umap-tooltip", "direction"),
    Input("umap-graph", "hoverData"),
)


@callback(
    Output("umap-tooltip", "children"),
    Input("umap-graph", "hoverData"),
    State("_pages_location", "search"),
    State("umap-mode-dropdown", "value"),
)
def display_umap_hover_content(hover_data, search: str | None, umap_mode: str | None):
    if hover_data is None:
        return no_update

    umap_mode = str(umap_mode or "images")
    pt = hover_data["points"][0]
    customdata = pt.get("customdata", None)

    if umap_mode == "features":
        if not isinstance(customdata, (list, tuple, np.ndarray)) or len(customdata) < 4:
            return no_update
        try:
            fid = int(customdata[0])
            mean_val = float(customdata[1])
            std_val = float(customdata[2])
            nonzero = float(customdata[3])
        except (TypeError, ValueError):
            return no_update

        corr_label = str(pt.get("fullData", {}).get("name", ""))
        return build_feature_umap_hover_tooltip(
            fid,
            mean_val,
            std_val,
            nonzero,
            corr_label,
            corr_source_label="corr(group_patch)",
        )

    params = visualization_params(search)
    if params is None:
        return no_update

    metric, _selection_dataset, model_key = params
    if not model_key:
        return no_update

    img_idx = hover_image_idx(hover_data)
    if img_idx is None or img_idx < 0:
        return no_update

    meta_line = cached_hover_meta_line(metric, model_key, img_idx)
    if meta_line is None:
        return no_update

    b64 = cached_hover_thumb(metric, model_key, img_idx)
    if b64 is None:
        body = html.Div("Preview unavailable", style={"fontSize": "12px"})
    else:
        body = html.Img(
            src=f"data:image/jpeg;base64,{b64}",
            style={
                "width": f"{THUMB_SIZE}px",
                "height": f"{THUMB_SIZE}px",
                "objectFit": "cover",
                "border": "1px solid #CBD5E1",
            },
        )

    return html.Div(
        [
            html.Div(meta_line, style={"fontSize": "12px", "marginBottom": "8px"}),
            body,
        ],
        style={
            "width": f"{max(250, THUMB_SIZE + 20)}px",
            "whiteSpace": "normal",
            "padding": "10px",
            "backgroundColor": "#F8FAFC",
            "border": "1px solid #CBD5E1",
            "borderRadius": "8px",
        },
    )


@callback(
    Output("selected-feature-id", "data"),
    Output("feature-status", "children"),
    Output("feature-jump-input", "value"),
    Input("feature-sort-selector", "value"),
    Input("feature-prev", "n_clicks"),
    Input("feature-next", "n_clicks"),
    Input("feature-jump-input", "value"),
    Input("umap-graph", "clickData"),
    State("selected-feature-id", "data"),
    State("_pages_location", "search"),
    State("umap-mode-dropdown", "value"),
)
def update_selected_feature(
    selector_name: str | None,
    _prev_clicks: int,
    _next_clicks: int,
    jump_value,
    click_data,
    current_feature_id,
    search: str | None,
    umap_mode: str | None,
):
    params = visualization_params(search)
    if params is None:
        return None, "", ""

    metric, selection_dataset, model_key = params
    rows = feature_rows_for_selector(metric, selection_dataset, model_key, selector_name)

    trigger_id = ctx.triggered_id
    requested_feature_id: int | float | str | None = jump_value

    if trigger_id == "umap-graph":
        if str(umap_mode or "images") != "features" or not click_data:
            return no_update, no_update, no_update
        customdata = click_data["points"][0].get("customdata", None)
        if not isinstance(customdata, (list, tuple, np.ndarray)) or len(customdata) < 1:
            return no_update, no_update, no_update
        try:
            requested_feature_id = int(customdata[0])
        except (TypeError, ValueError):
            return no_update, no_update, no_update

    selected_feature_id, status = select_feature_id(
        feature_rows=rows,
        current_feature_id=current_feature_id,
        trigger_id=trigger_id,
        requested_feature_id=requested_feature_id,
    )
    if selected_feature_id is not None:
        unavailable = feature_unavailable_message(
            metric,
            model_key,
            int(selected_feature_id),
            activation_dataset=UMAP_DATASET,
        )
        if unavailable is not None:
            status = unavailable
    jump_display = str(selected_feature_id) if selected_feature_id is not None else ""
    return selected_feature_id, status, jump_display


@callback(
    Output("feature-detail-body", "children"),
    Input("selected-feature-id", "data"),
    State("_pages_location", "search"),
    State("feature-sort-selector", "value"),
)
def render_feature_detail(
    selected_feature_id,
    search: str | None,
    selector_name: str | None,
):
    params = visualization_params(search)
    if params is None:
        return html.Div("No selector cache files were found for this model.", className="feature-empty")

    metric, selection_dataset, model_key = params
    return build_feature_detail_content(
        metric,
        selection_dataset,
        model_key,
        selected_feature_id,
        selector_name,
    )


@callback(
    Output("feature-top-images-cache", "data"),
    Output("feature-top-images-status", "children"),
    Output("feature-image-caption-0", "children"),
    Input("selected-feature-id", "data"),
    Input("feature-image-ranking", "value"),
    State("_pages_location", "search"),
)
def load_feature_top_images_cache(
    selected_feature_id,
    ranking_mode: str | None,
    search: str | None,
):
    params = visualization_params(search)
    if params is None:
        return _empty_feature_image_cache("No images available.")

    metric, selection_dataset, model_key = params

    try:
        feature_id = int(selected_feature_id)
    except Exception:
        return _empty_feature_image_cache("No feature selected.")

    unavailable = feature_unavailable_message(
        metric,
        model_key,
        feature_id,
        activation_dataset=selection_dataset,
    )
    if unavailable is not None:
        return _empty_feature_image_cache(unavailable)

    ranking_mode = str(ranking_mode or "iou")
    try:
        overlays = get_top_feature_overlays(
            metric,
            selection_dataset,
            model_key,
            feature_id,
            top_n=FEATURE_TOP_IMAGE_TOP_N,
            ranking_mode=ranking_mode,
        )
    except Exception as exc:
        return _empty_feature_image_cache(f"Failed to build top images: {exc}")
    if not overlays:
        return _empty_feature_image_cache("No top images found for this feature.")

    figure = overlays[0].get("figure")
    if not isinstance(figure, dict):
        return _empty_feature_image_cache("No top images found for this feature.")

    caption = str(overlays[0].get("caption") or "Top image 1")
    return figure, "", caption


@callback(
    Output("feature-image-0", "figure"),
    Input("feature-top-images-cache", "data"),
    State("feature-show-heatmap-overlay", "value"),
)
def load_feature_top_images_figure(cache, show_heatmap_value):
    show_overlay = overlay_show_heatmap_enabled(show_heatmap_value)
    if not isinstance(cache, dict):
        return empty_heatmap_figure("Top images overlay is not available yet.").to_plotly_json()
    return set_heatmap_overlay_visible(cache, show_overlay)

clientside_callback(
    TOGGLE_HEATMAP_OVERLAY_CLIENTSIDE,
    Output("feature-image-0", "figure", allow_duplicate=True),
    Input("feature-show-heatmap-overlay", "value"),
    State("feature-image-0", "figure"),
    prevent_initial_call=True,
)