scripts/explorer/panels/feature.py

"""
Feature detail panel: stats header, top NSD MEI grid, cortical profile.

Pure display logic lives in explorer.feature_logic.

Exports:
  update_feature_display(feat)
  select_and_display(feat)
  status_div, stats_div, top_heatmap_div, brain_div
  add_steer_btn
"""

import threading

from bokeh.io import curdoc
from bokeh.models import Button, Div

from ..state import _S, active_ds
from ..brain import HAS_DYNADIFF, render_cortical_profile
from ..feature_logic import build_stats_html, build_top_mei_items
from ..rendering import status_html, make_image_grid_html
from .. import widgets

N_DISPLAY = 9

# ---------- Divs ----------

status_div      = Div(text=status_html('idle', 'Select a feature to begin.'), width=680)
stats_div       = Div(text="<h3>Select a feature to explore it</h3>", width=530)
top_heatmap_div = Div(text="", width=680)
brain_div       = Div(text="", width=680)

add_steer_btn = Button(
    label="+ Add to Steer",
    button_type="success",
    width=140,
    visible=HAS_DYNADIFF,
)


# ---------- Core display logic ----------

def update_feature_display(feature_idx):
    feat = int(feature_idx)
    _S.render_token += 1
    my_token = _S.render_token

    ds = active_ds()
    freq_val = ds['feature_frequency'][feat].item()
    feat_name = ds['feature_names'].get(feat, "")
    auto_name = ds['auto_interp_names'].get(feat, "")

    stats_div.text = build_stats_html(feat, freq_val, feat_name, auto_name)

    # Sync name-editing widget
    from .feature_list import name_input
    name_input.value = feat_name

    if freq_val == 0:
        status_div.text = status_html('dead', f'Feature {feat} never activated.')
        brain_div.text = render_cortical_profile(feat)
        top_heatmap_div.text = ""
        add_steer_btn.visible = False
        return

    add_steer_btn.visible = HAS_DYNADIFF
    status_div.text = status_html('loading', '&#x23F3; Loading...')

    doc = curdoc()

    def _render():
        if _S.render_token != my_token:
            return

        _ds = active_ds()
        try:
            zp = int(widgets.zoom_slider.value)
            ha = widgets.heatmap_alpha_slider.value
        except Exception:
            zp, ha = 16, 1.0
        top_infos, top_img_is, subset_label = build_top_mei_items(
            feat, _ds, n_display=N_DISPLAY,
            zoom_patches=zp,
            heatmap_alpha=ha,
        )

        top_heatmap_div.text = make_image_grid_html(
            top_infos, f"Top MEIs (feature {feat}){subset_label}",
            img_indices=top_img_is, cols=3)

        status_div.text = status_html('idle', '')

        # Render cortical profile in background thread so the UI stays responsive
        brain_div.text = ""

        def _render_brain():
            html = render_cortical_profile(feat)
            def _update_brain():
                if _S.render_token == my_token:
                    brain_div.text = html
            doc.add_next_tick_callback(_update_brain)

        threading.Thread(target=_render_brain, daemon=True).start()

    doc.add_next_tick_callback(_render)


def select_and_display(feat: int):
    """Display feature detail and highlight it."""
    widgets.feature_input.value = str(feat)
    update_feature_display(feat)


# ---------- Add-to-Steer callback ----------

def _on_add_steer():
    try:
        feat = int(widgets.feature_input.value)
    except ValueError:
        return
    from .steering import add_feature
    add_feature(feat)

add_steer_btn.on_click(_on_add_steer)


# ---------- Re-render on zoom/alpha changes ----------

def _rerender_current(attr, old, new):
    try:
        feat = int(widgets.feature_input.value)
        if 0 <= feat < active_ds()['d_model']:
            update_feature_display(feat)
    except ValueError:
        pass

widgets.zoom_slider.on_change('value', _rerender_current)
widgets.heatmap_alpha_slider.on_change('value', _rerender_current)