scripts/explorer/panels/examples.py

"""
Examples panel: load preset feature combinations into the DynaDiff steering list.

Presets are defined in a JSON file passed via --examples-file:
  [
    {"name": "Faces + Scenes", "image": "nsd_22910", "features": [
      {"feat": 1234, "lam": 3.0, "threshold": 0.10},
      {"feat": 5678, "lam": 2.0, "threshold": 0.05}
    ]},
    ...
  ]

The optional "image" field is an NSD basename (e.g. "nsd_22910") or a dataset
image index. When present, clicking the preset also loads that image into the
patch explorer and sets it as the GT brain sample.

When --editable-examples is set, a "Save as Preset" button appears that appends
the current steering list to the examples JSON file (and pushes to HF).

Exports:
  examples_panel — column layout
"""

import json
import os
import threading

from bokeh.io import curdoc
from bokeh.layouts import column, row
from bokeh.models import Button, Div, TextInput

from ..args import args
from ..brain import HAS_DYNADIFF


# ---------- Load presets ----------

_presets = []
_examples_path = args.examples_file

if _examples_path and os.path.isfile(_examples_path):
    try:
        with open(_examples_path) as f:
            _presets = json.load(f)
        print(f"[Examples] Loaded {len(_presets)} presets from {_examples_path}")
    except Exception as e:
        print(f"[Examples] WARNING: could not load {_examples_path}: {e}")


# ---------- Build panel ----------

_header = Div(
    text=(
        '<div class="sae-card-header" style="border-bottom:none">'
        'Example Presets</div>'
        '<div style="font-size:11px;color:#9ca3af;margin-bottom:8px">'
        'Load pre-defined feature combinations into the steering list.</div>'
    ),
    width=200,
)

_buttons_container = column(width=200)
_save_status_div = Div(text="", width=200)


def _make_preset_callback(preset):
    def _on_click():
        from .steering import set_preset, load_patch_image
        set_preset(preset['features'], label=preset['name'])
        image = preset.get('image')
        if image:
            load_patch_image(str(image))
    return _on_click


_editable = bool(args.editable_examples and _examples_path)


def _remove_preset(idx):
    """Remove preset by index, save, rebuild, and push."""
    if 0 <= idx < len(_presets):
        removed = _presets.pop(idx)
        _save_examples_to_disk()
        _rebuild_buttons()
        _save_status_div.text = (
            f'<span style="color:#6b7280;font-size:11px">'
            f'Removed &ldquo;{removed["name"]}&rdquo;</span>')
        threading.Thread(target=_push_examples_to_hf, daemon=True).start()


def _make_remove_callback(idx):
    def _on_click():
        _remove_preset(idx)
    return _on_click


def _rebuild_buttons():
    """Rebuild the preset button list from _presets."""
    children = []
    for i, preset in enumerate(_presets):
        btn = Button(
            label=preset['name'],
            button_type='light',
            width=(155 if _editable else 190),
            height=30,
        )
        btn.on_click(_make_preset_callback(preset))
        if _editable:
            rm_btn = Button(
                label="✕", button_type="light",
                width=30, height=30,
            )
            rm_btn.on_click(_make_remove_callback(i))
            children.append(row(btn, rm_btn))
        else:
            children.append(btn)
    _buttons_container.children = children


# ---------- Save-as-preset ----------

_save_name_input = TextInput(
    placeholder="Preset name…", width=190, height=30,
    visible=False,
)
_save_btn = Button(
    label="+ Save as Preset", button_type="success", width=190, height=30,
    visible=False,
)


def _save_examples_to_disk():
    """Write _presets to the examples JSON file."""
    with open(_examples_path, 'w') as f:
        json.dump(_presets, f, indent=2)
    print(f"[Examples] Saved {len(_presets)} presets to {_examples_path}")


def _push_examples_to_hf():
    """Upload the examples file to HF dataset repo (blocking)."""
    hf_token = os.environ.get("HF_TOKEN")
    hf_repo = os.environ.get("HF_DATASET_REPO")
    if not (hf_token and hf_repo) or not _examples_path:
        return
    try:
        from huggingface_hub import upload_file
        upload_file(
            path_or_fileobj=_examples_path,
            path_in_repo=os.path.basename(_examples_path),
            repo_id=hf_repo, repo_type="dataset", token=hf_token,
            commit_message="Update example presets",
        )
        print(f"  Pushed {os.path.basename(_examples_path)} to HF dataset {hf_repo}")
    except Exception as e:
        print(f"  Warning: HF push of examples failed: {e}")


def _on_save_preset():
    from .steering import _dd_source, _Session

    feats = list(_dd_source.data['feat'])
    if not feats:
        _save_status_div.text = (
            '<span style="color:#dc2626;font-size:11px">'
            'Add features to steering first.</span>')
        return

    name = _save_name_input.value.strip()
    if not name:
        _save_status_div.text = (
            '<span style="color:#dc2626;font-size:11px">'
            'Enter a name for the preset.</span>')
        return

    lams = list(_dd_source.data['lam'])
    thrs = list(_dd_source.data['threshold'])
    features = [
        {"feat": int(f), "lam": float(l), "threshold": float(t)}
        for f, l, t in zip(feats, lams, thrs)
    ]
    preset = {"name": name, "features": features}
    if _Session.nsd_basename:
        preset["image"] = _Session.nsd_basename

    _presets.append(preset)
    _save_examples_to_disk()
    _rebuild_buttons()
    _save_name_input.value = ""
    _save_status_div.text = (
        f'<span style="color:#2563eb;font-size:11px">'
        f'Saved &ldquo;{name}&rdquo;</span>')

    # Push to HF in background
    threading.Thread(target=_push_examples_to_hf, daemon=True).start()


if not HAS_DYNADIFF:
    examples_panel = column(
        _header,
        Div(text='<i style="color:#aaa;font-size:11px">DynaDiff unavailable.</i>',
            width=200),
    )
else:
    _rebuild_buttons()

    _save_btn.visible = _editable
    _save_name_input.visible = _editable
    if _editable:
        _save_btn.on_click(_on_save_preset)

    examples_panel = column(
        _header,
        _buttons_container,
        _save_name_input,
        _save_btn,
        _save_status_div,
    )