brainEncoding/visualization/serve_ic_label_gui.py

#!/usr/bin/env python
"""Serve a small local GUI for approving feature labels."""

from __future__ import annotations

import argparse
import csv
import json
import os
import re
import sys
import threading
import tempfile
from datetime import datetime, timezone
from http import HTTPStatus
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from io import BytesIO
from pathlib import Path
from typing import Any
from urllib.parse import parse_qs, urlparse


CONFIDENCE_VALUES = ("high", "medium", "low", "unclear")
INTERPRETATION_TYPES = (
    "form",
    "word",
    "phrase",
    "sentence",
    "long-range",
    "global",
    "position",
    "sophisticated",
    "unclear",
)
SUBJECT = "UTS03"
N_VERTICES = 20484
HEMISPHERES = ("lh", "rh")
BRAIN_ENCODING = Path(__file__).resolve().parents[1]
DEFAULT_SCRATCH_ROOT = Path(os.environ.get("BRAINENCODING_SCRATCH_ROOT", "/storage/scratch1/8/whuang409"))
DEFAULT_PYCORTEX_DB = DEFAULT_SCRATCH_ROOT / "visionLM_data/ds003020/derivatives/pycortex-db"
DEFAULT_FREESURFER_SUBJECTS = DEFAULT_SCRATCH_ROOT / "visionLM_data/ds003020/derivatives/freesurfer_subjdir"
DEFAULT_FSAVERAGE5 = Path("/storage/project/r-aivanova7-0/shared/.env/freesurfer/subjects/fsaverage5")
DEFAULT_NOISE_CEILING = (
    Path("/storage/project/r-aivanova7-0/whuang409/loopedEncoding/results/noise_ceiling_fsaverage5")
    / "UTS03_wheretheressmoke_noise_ceiling_fsaverage5.npz"
)
DEFAULT_ICA_PREDICTIVITY = (
    BRAIN_ENCODING
    / "results/gemma-2-2b_layer13_lookback128_hook_resid_post_ica_lasso_fixed_a0.02207_ridge_sweep/predictivity/"
    "UTS03_gemma-2-2b_ica_lasso_fixed_lookback_128_hook_resid_post_layer_13_a0.02207_ridge_sweep_predictivity.npz"
)
DEFAULT_SAE_PREDICTIVITY = (
    BRAIN_ENCODING
    / "results/gemma-2-2b_layer13_lookback128_hook_resid_post_sae_dense_alpha_sweep/predictivity/"
    "UTS03_gemma-2-2b_sae_dense_lookback_128_hook_resid_post_layer_13_alpha_sweep_predictivity.npz"
)
STOPWORDS = {
    "a",
    "an",
    "and",
    "are",
    "as",
    "at",
    "be",
    "by",
    "for",
    "from",
    "in",
    "is",
    "it",
    "of",
    "on",
    "or",
    "that",
    "the",
    "this",
    "to",
    "with",
}
SEMANTIC_GROUPS = (
    (
        "visual",
        "visual vision see saw seen look looks looking watch watching image picture photo camera light color dark bright face body scene object eye eyes draw drawing art mirror",
    ),
    (
        "social",
        "social person people family mother father mom dad brother sister aunt child children boy girl woman man wife husband friend partner relationship married kinship",
    ),
    (
        "mental",
        "mental mind think thought know knew believe realize understand remember memory worry hope fear afraid ashamed self identity confidence reflection inner",
    ),
    (
        "speech",
        "speech say said tell told ask asked answer quote quoted dialogue conversation voice talk called utterance instruction command",
    ),
    (
        "motion",
        "motion move moving walk run drive fly fall jump landing turn pull push go come arrive leave direction forward back toward",
    ),
    (
        "time",
        "time date year years day night morning early late week month birthday age first start opening beginning after before",
    ),
    (
        "number",
        "number count counted counting one two three four five seven ten twenty seventy eighty ninety hundred thousand million digit math",
    ),
    (
        "body",
        "body hand hands arm leg head eye eyes mouth face skin blood hurt pain sick illness disease epilepsy doctor medical hospital",
    ),
    (
        "place",
        "place location room house home school college university office stage car taxi boat store camp prison church city",
    ),
    (
        "emotion",
        "emotion feel feeling felt love joy happy sad grief cry crying shame fear worried anxious angry regret relief",
    ),
    (
        "religion",
        "religion religious church god soul souls prayer pray temple faith sin spiritual exorcism offering",
    ),
    (
        "academic",
        "academic school college university class student teacher phd mit harvard princeton exam study science database actuarial",
    ),
    (
        "object",
        "object thing stuff phone key keys card box door car bag clothes furniture couch table book paper photograph",
    ),
    (
        "function",
        "function grammar article determiner preposition conjunction pronoun deictic demonstrative modal infinitive filler pause um uh and the for",
    ),
)
SEMANTIC_EXPANSIONS: dict[str, set[str]] = {}
SEMANTIC_GROUP_TOKENS: dict[str, set[str]] = {}
for _group_name, _words in SEMANTIC_GROUPS:
    _tokens = {_group_name, *re.findall(r"[a-z0-9]+", _words.lower())}
    SEMANTIC_GROUP_TOKENS[_group_name] = _tokens
    for _token in _tokens:
        SEMANTIC_EXPANSIONS.setdefault(_token, set()).update(_tokens)


def main() -> int:
    args = parse_args()
    packets = load_packets(args.packets)
    labels = load_labels(args.labels)
    highlights = load_highlights(args.highlights)
    server = LabelServer(
        (args.host, int(args.port)),
        Handler,
        packets=packets,
        packet_paths=args.packets,
        labels=labels,
        highlights=highlights,
        labels_path=args.labels,
        highlights_path=args.highlights,
        stage_a_root=args.stage_a_root,
        layer=args.layer,
        encoding_dir=args.encoding_dir,
        ica_contribution_dir=args.ica_contribution_dir,
        sae_coeff_dir=args.sae_coeff_dir,
        sae_erf_csv=args.sae_erf_csv,
        noise_ceiling_npz=args.noise_ceiling_npz,
        ica_predictivity_npz=args.ica_predictivity_npz,
        sae_predictivity_npz=args.sae_predictivity_npz,
        ncsnr_threshold=args.ncsnr_threshold,
        mask_reliable_vertices=args.mask_reliable_vertices,
        roi_dir=args.roi_dir,
        brainmap_cache=args.brainmap_cache,
        pycortex_db=args.pycortex_db,
        freesurfer_subjects=args.freesurfer_subjects,
        fsaverage5=args.fsaverage5,
        flatmap_height=args.flatmap_height,
    )
    print(f"serving http://{args.host}:{args.port}", flush=True)
    print(f"packets={len(packets)} labels={args.labels}", flush=True)
    print(f"highlights={args.highlights}", flush=True)
    server.serve_forever()
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    root = Path(__file__).resolve().parents[1]
    label_root = root / "results/story_ica_labels_UTS03_layer13"
    encoding_root = root / "results/story_ica_encoding_UTS03_layer13/encoding/layer_13/UTS03"
    parser.add_argument("--packets", type=Path, nargs="+", default=[label_root / "ic_label_packets.jsonl"])
    parser.add_argument("--labels", type=Path, default=label_root / "ic_labels.json")
    parser.add_argument("--highlights", type=Path, default=label_root / "ic_highlights.json")
    parser.add_argument("--stage-a-root", type=Path, default=root / "results/story_ica_stage_a_UTS03_layer13")
    parser.add_argument("--encoding-dir", type=Path, default=encoding_root)
    parser.add_argument("--ica-contribution-dir", type=Path, default=None)
    parser.add_argument("--sae-coeff-dir", type=Path, default=None)
    parser.add_argument("--sae-erf-csv", type=Path, default=None)
    parser.add_argument("--noise-ceiling-npz", type=Path, default=DEFAULT_NOISE_CEILING)
    parser.add_argument("--ica-predictivity-npz", type=Path, default=DEFAULT_ICA_PREDICTIVITY)
    parser.add_argument("--sae-predictivity-npz", type=Path, default=DEFAULT_SAE_PREDICTIVITY)
    parser.add_argument("--ncsnr-threshold", type=float, default=0.20)
    parser.add_argument("--mask-reliable-vertices", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--roi-dir", type=Path, default=root / "masks/requested_roi_groups_20484")
    parser.add_argument("--brainmap-cache", type=Path, default=label_root / "ic_brainmaps")
    parser.add_argument("--pycortex-db", type=Path, default=DEFAULT_PYCORTEX_DB)
    parser.add_argument("--freesurfer-subjects", type=Path, default=DEFAULT_FREESURFER_SUBJECTS)
    parser.add_argument("--fsaverage5", type=Path, default=DEFAULT_FSAVERAGE5)
    parser.add_argument("--flatmap-height", type=int, default=1024)
    parser.add_argument("--layer", type=int, default=13)
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=8765)
    return parser.parse_args()


def load_packets(paths: list[Path]) -> list[dict[str, Any]]:
    packets = []
    for path in paths:
        if not path.exists():
            continue
        with path.open(encoding="utf-8") as handle:
            for line in handle:
                if line.strip():
                    packet = json.loads(line)
                    packet.setdefault("feature_source", "ica")
                    packet.setdefault("feature_label", "IC" if packet["feature_source"] == "ica" else str(packet["feature_source"]).upper())
                    packet["_search_evidence_text"] = search_evidence_text(packet)
                    packet["_search_evidence_tokens"] = semantic_tokens(packet["_search_evidence_text"])
                    packets.append(packet)
    return sorted(
        packets,
        key=lambda packet: (
            1 if packet.get("feature_source", "ica") == "ica" else 0,
            int(packet.get("encoding_summary", {}).get("selected_voxels") or 0),
            int(packet.get("encoding_summary", {}).get("activation_count") or 0),
            int(packet.get("encoding_summary", {}).get("selected_rows") or 0),
            -int(packet.get("component") or 0),
            1 if packet.get("side") == "positive" else 0,
        ),
        reverse=True,
    )


def load_labels(path: Path) -> dict[str, Any]:
    if path.exists():
        return json.loads(path.read_text(encoding="utf-8"))
    return {"metadata": {"created_utc": datetime.now(timezone.utc).isoformat()}, "labels": {}}


def load_highlights(path: Path) -> dict[str, Any]:
    if path.exists():
        payload = json.loads(path.read_text(encoding="utf-8"))
    else:
        payload = {"metadata": {"created_utc": datetime.now(timezone.utc).isoformat()}, "components": [], "features": []}
    payload.setdefault("metadata", {})
    payload["components"] = sorted({int(component) for component in payload.get("components", [])})
    legacy = {f"ica:{component}" for component in payload["components"]}
    payload["features"] = sorted({str(feature) for feature in payload.get("features", [])} | legacy)
    return payload


class LabelServer(ThreadingHTTPServer):
    def __init__(
        self,
        *args: Any,
        packets: list[dict[str, Any]],
        packet_paths: list[Path],
        labels: dict[str, Any],
        highlights: dict[str, Any],
        labels_path: Path,
        highlights_path: Path,
        stage_a_root: Path,
        layer: int,
        encoding_dir: Path,
        ica_contribution_dir: Path | None,
        sae_coeff_dir: Path | None,
        sae_erf_csv: Path | None,
        noise_ceiling_npz: Path | None,
        ica_predictivity_npz: Path | None,
        sae_predictivity_npz: Path | None,
        ncsnr_threshold: float,
        mask_reliable_vertices: bool,
        roi_dir: Path,
        brainmap_cache: Path,
        pycortex_db: Path,
        freesurfer_subjects: Path,
        fsaverage5: Path,
        flatmap_height: int,
        **kwargs: Any,
    ) -> None:
        super().__init__(*args, **kwargs)
        self.packets = packets
        self.packet_paths = packet_paths
        self.packet_mtimes = self.current_packet_mtimes()
        self.labels = labels
        self.highlights = highlights
        self.labels_path = labels_path
        self.highlights_path = highlights_path
        self.labels_mtime = labels_path.stat().st_mtime if labels_path.exists() else 0.0
        self.highlights_mtime = highlights_path.stat().st_mtime if highlights_path.exists() else 0.0
        self.stage_a_root = stage_a_root
        self.layer = int(layer)
        self.encoding_dir = encoding_dir
        self.ica_contribution_dir = ica_contribution_dir
        self.sae_coeff_dir = sae_coeff_dir
        self.sae_erf_csv = sae_erf_csv
        self.noise_ceiling_npz = noise_ceiling_npz
        self.ica_predictivity_npz = ica_predictivity_npz
        self.sae_predictivity_npz = sae_predictivity_npz
        self.ncsnr_threshold = float(ncsnr_threshold)
        self.mask_reliable_vertices = bool(mask_reliable_vertices)
        self.roi_dir = roi_dir
        self.brainmap_cache = brainmap_cache
        self.pycortex_db = pycortex_db
        self.freesurfer_subjects = freesurfer_subjects
        self.fsaverage5 = fsaverage5
        self.flatmap_height = int(flatmap_height)
        self._brainmap_context: ICBrainmapContext | None = None
        self._brainmap_lock = threading.Lock()
        self._highlight_lock = threading.Lock()

    def current_packet_mtimes(self) -> dict[str, float]:
        return {str(path): path.stat().st_mtime for path in self.packet_paths if path.exists()}

    def reload_packets_if_changed(self) -> None:
        mtimes = self.current_packet_mtimes()
        if mtimes != self.packet_mtimes:
            self.packets = load_packets(self.packet_paths)
            self.packet_mtimes = mtimes

    def reload_labels_if_changed(self) -> None:
        if not self.labels_path.exists():
            return
        mtime = self.labels_path.stat().st_mtime
        if mtime > self.labels_mtime:
            self.labels = load_labels(self.labels_path)
            self.labels_mtime = mtime

    def reload_highlights_if_changed(self) -> None:
        if not self.highlights_path.exists():
            return
        mtime = self.highlights_path.stat().st_mtime
        if mtime > self.highlights_mtime:
            self.highlights = load_highlights(self.highlights_path)
            self.highlights_mtime = mtime

    def brainmap_context(self) -> "ICBrainmapContext":
        if self._brainmap_context is not None:
            return self._brainmap_context
        with self._brainmap_lock:
            if self._brainmap_context is None:
                self._brainmap_context = ICBrainmapContext(
                    layer=self.layer,
                    encoding_dir=self.encoding_dir,
                    ica_contribution_dir=self.ica_contribution_dir,
                    sae_coeff_dir=self.sae_coeff_dir,
                    noise_ceiling_npz=self.noise_ceiling_npz,
                    ica_predictivity_npz=self.ica_predictivity_npz,
                    sae_predictivity_npz=self.sae_predictivity_npz,
                    ncsnr_threshold=self.ncsnr_threshold,
                    mask_reliable_vertices=self.mask_reliable_vertices,
                    stage_a_root=self.stage_a_root,
                    roi_dir=self.roi_dir,
                    output_dir=self.brainmap_cache,
                    pycortex_db=self.pycortex_db,
                    freesurfer_subjects=self.freesurfer_subjects,
                    fsaverage5=self.fsaverage5,
                    flatmap_height=self.flatmap_height,
                )
        return self._brainmap_context


class Handler(BaseHTTPRequestHandler):
    server: LabelServer

    def do_GET(self) -> None:  # noqa: N802
        parsed = urlparse(self.path)
        if parsed.path == "/":
            self.send_html(INDEX_HTML)
            return
        if parsed.path == "/api/state":
            query = parse_qs(parsed.query)
            self.server.reload_packets_if_changed()
            self.server.reload_labels_if_changed()
            self.server.reload_highlights_if_changed()
            self.send_json(self.state(query))
            return
        if parsed.path == "/api/labels":
            self.server.reload_labels_if_changed()
            self.send_json(self.server.labels)
            return
        if parsed.path == "/api/highlights":
            self.server.reload_highlights_if_changed()
            self.send_json(self.server.highlights)
            return
        if parsed.path in {"/api/ic-brainmap", "/api/feature-brainmap"}:
            self.send_ic_brainmap(parse_qs(parsed.query))
            return
        if parsed.path == "/api/erf-brainmap":
            self.send_erf_brainmap(parse_qs(parsed.query))
            return
        self.send_error(HTTPStatus.NOT_FOUND)

    def do_POST(self) -> None:  # noqa: N802
        parsed = urlparse(self.path)
        if parsed.path == "/api/label":
            self.save_label()
            return
        if parsed.path == "/api/highlight":
            self.save_highlight()
            return
        self.send_error(HTTPStatus.NOT_FOUND)

    def read_json_payload(self) -> dict[str, Any]:
        length = int(self.headers.get("content-length") or 0)
        return json.loads(self.rfile.read(length).decode("utf-8"))

    def save_label(self) -> None:
        self.server.reload_labels_if_changed()
        payload = self.read_json_payload()
        packet_id = str(payload["packet_id"])
        packet = next((item for item in self.server.packets if str(item.get("packet_id")) == packet_id), {})
        item = {
            "packet_id": packet_id,
            "component": int(payload["component"]),
            "feature_source": str(payload.get("feature_source") or "ica"),
            "side": str(payload["side"]),
            "working_label": str(payload.get("working_label", "")).strip(),
            "confidence": normalize_choice(payload.get("confidence"), CONFIDENCE_VALUES),
            "interpretation_type": normalize_choice(payload.get("interpretation_type"), INTERPRETATION_TYPES),
            "approved": bool(payload.get("approved", False)),
            "notes": str(payload.get("notes", "")).strip(),
            "updated_utc": datetime.now(timezone.utc).isoformat(),
        }
        item["semantic_tags"] = normalize_semantic_tags(payload.get("semantic_tags")) or broad_semantic_tags(packet, item)
        self.server.labels.setdefault("labels", {})[packet_id] = item
        save_json_atomic(self.server.labels_path, self.server.labels)
        self.server.labels_mtime = self.server.labels_path.stat().st_mtime
        self.send_json({"ok": True, "label": item})

    def save_highlight(self) -> None:
        self.server.reload_highlights_if_changed()
        payload = self.read_json_payload()
        component = int(payload["component"])
        feature_source = str(payload.get("feature_source") or "ica")
        feature_key = make_feature_key(feature_source, component)
        highlighted = bool(payload.get("highlighted", False))
        with self.server._highlight_lock:
            components = {int(value) for value in self.server.highlights.get("components", [])}
            features = {str(value) for value in self.server.highlights.get("features", [])}
            if highlighted:
                features.add(feature_key)
                if feature_source == "ica":
                    components.add(component)
            else:
                features.discard(feature_key)
                if feature_source == "ica":
                    components.discard(component)
            self.server.highlights["components"] = sorted(components)
            self.server.highlights["features"] = sorted(features)
            self.server.highlights.setdefault("metadata", {})["updated_utc"] = datetime.now(timezone.utc).isoformat()
            save_json_atomic(self.server.highlights_path, self.server.highlights)
            self.server.highlights_mtime = self.server.highlights_path.stat().st_mtime
        self.send_json({"ok": True, "component": component, "feature_source": feature_source, "highlighted": highlighted})

    def state(self, query: dict[str, list[str]]) -> dict[str, Any]:
        packets = self.server.packets
        labels = self.server.labels.get("labels", {})
        highlighted_components = {int(value) for value in self.server.highlights.get("components", [])}
        highlighted_features = {str(value) for value in self.server.highlights.get("features", [])}
        ica_erf = load_erf(self.server.stage_a_root, self.server.layer)
        sae_erf = load_sae_erf(self.server.sae_erf_csv, self.server.sae_coeff_dir, self.server.layer)
        q = (query.get("q", [""])[0] or "").lower().strip()
        status = query.get("status", ["all"])[0]
        side = query.get("side", ["all"])[0]
        source = query.get("source", ["all"])[0]
        interpretation_type = query.get("interpretation_type", ["all"])[0]
        semantic_tag_filter = query.get("semantic_tag", ["all"])[0]
        highlight_filter = query.get("highlight_filter", ["all"])[0]
        sort_by = query.get("sort_by", ["selected_voxels"])[0]
        sort_then = query.get("sort_then", ["erf"])[0]
        sort_dir = query.get("sort_dir", ["desc"])[0]
        limit = min(500, max(1, int(query.get("limit", ["100"])[0] or 100)))
        offset = max(0, int(query.get("offset", ["0"])[0] or 0))
        filtered = []
        for packet in packets:
            packet_id = packet["packet_id"]
            label = labels.get(packet_id, {})
            component = int(packet["component"])
            packet_source = str(packet.get("feature_source") or "ica")
            feature_key = make_feature_key(packet_source, component)
            is_highlighted = feature_key in highlighted_features or (
                packet_source == "ica" and component in highlighted_components
            )
            if source != "all" and packet_source != source:
                continue
            if side != "all" and packet.get("side") != side:
                continue
            if status == "approved" and not label.get("approved"):
                continue
            if status == "unapproved" and label.get("approved"):
                continue
            if interpretation_type != "all" and label.get("interpretation_type", "unclear") != interpretation_type:
                continue
            if highlight_filter == "highlighted" and not is_highlighted:
                continue
            if highlight_filter == "unhighlighted" and is_highlighted:
                continue
            label = dict(label)
            label["semantic_tags"] = normalize_semantic_tags(label.get("semantic_tags")) or broad_semantic_tags(packet, label)
            if semantic_tag_filter != "all" and semantic_tag_filter not in label["semantic_tags"]:
                continue
            search_score = semantic_search_score(q, packet, label) if q else 0.0
            if q and search_score <= 0.0:
                continue
            packet_copy = {key: value for key, value in packet.items() if not key.startswith("_")}
            packet_copy["label"] = label
            packet_copy["highlighted"] = is_highlighted
            packet_copy["feature_key"] = feature_key
            packet_copy["search_score"] = search_score
            if packet_source == "ica" and component in ica_erf:
                packet_copy["erf"] = ica_erf[component]
            elif packet_source == "sae" and component in sae_erf:
                packet_copy["erf"] = sae_erf[component]
            filtered.append(packet_copy)
        if q and sort_by != "search_score":
            sort_then = sort_by
            sort_by = "search_score"
        filtered = sort_packets(filtered, sort_by=sort_by, sort_then=sort_then, sort_dir=sort_dir)
        page = filtered[offset : offset + limit]
        return {
            "packets": page,
            "total_packets": len(packets),
            "shown_packets": len(filtered),
            "returned_packets": len(page),
            "offset": offset,
            "limit": limit,
            "approved_count": sum(1 for value in labels.values() if value.get("approved")),
            "highlighted_count": len(highlighted_components),
            "confidence_values": list(CONFIDENCE_VALUES),
            "interpretation_types": list(INTERPRETATION_TYPES),
            "semantic_tags": sorted(SEMANTIC_GROUP_TOKENS),
            "sources": sorted({str(packet.get("feature_source") or "ica") for packet in packets}),
            "sort_values": ["search_score", "selected_voxels", "activation_count", "max_score", "erf", "confidence", "component"],
        }

    def send_json(self, payload: Any) -> None:
        raw = json.dumps(payload, ensure_ascii=False).encode("utf-8")
        self.send_response(HTTPStatus.OK)
        self.send_header("content-type", "application/json; charset=utf-8")
        self.send_header("content-length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def send_html(self, html: str) -> None:
        raw = html.encode("utf-8")
        self.send_response(HTTPStatus.OK)
        self.send_header("content-type", "text/html; charset=utf-8")
        self.send_header("content-length", str(len(raw)))
        self.end_headers()
        self.wfile.write(raw)

    def send_ic_brainmap(self, query: dict[str, list[str]]) -> None:
        try:
            component = int(query.get("component", [""])[0])
            side = query.get("side", [""])[0]
            hemi = query.get("hemi", [""])[0]
            source = query.get("source", ["ica"])[0] or "ica"
            style = query.get("style", ["flatmap"])[0] or "flatmap"
            mask_reliable = query_bool(
                query,
                "mask_reliable",
                default=self.server.mask_reliable_vertices,
            )
            mask_significant = query_bool(
                query,
                "mask_significant",
                default=mask_reliable,
            )
        except ValueError:
            self.send_error(HTTPStatus.BAD_REQUEST, "Invalid feature brainmap query.")
            return
        if source not in {"ica", "sae"} or side not in {"positive", "negative"} or style not in {"flatmap", "litcoder"}:
            self.send_error(HTTPStatus.BAD_REQUEST, "Invalid feature brainmap source, side, or hemisphere.")
            return
        if style == "flatmap" and hemi not in HEMISPHERES:
            self.send_error(HTTPStatus.BAD_REQUEST, "Invalid feature brainmap hemisphere.")
            return
        try:
            raw = self.server.brainmap_context().render(
                source,
                component,
                side,
                hemi,
                style=style,
                mask_reliable_vertices=mask_reliable,
                mask_significant_vertices=mask_significant,
            )
        except Exception as exc:  # noqa: BLE001
            self.send_error(HTTPStatus.INTERNAL_SERVER_ERROR, f"Could not render feature brainmap: {exc}")
            return
        self.send_response(HTTPStatus.OK)
        self.send_header("content-type", "image/png")
        self.send_header("content-length", str(len(raw)))
        self.send_header("cache-control", "no-store")
        self.end_headers()
        self.wfile.write(raw)

    def send_erf_brainmap(self, query: dict[str, list[str]]) -> None:
        hemi = query.get("hemi", [""])[0]
        style = query.get("style", ["flatmap"])[0] or "flatmap"
        if style not in {"flatmap", "litcoder"}:
            self.send_error(HTTPStatus.BAD_REQUEST, "Invalid ERF brainmap style.")
            return
        if style == "flatmap" and hemi not in HEMISPHERES:
            self.send_error(HTTPStatus.BAD_REQUEST, "Invalid ERF brainmap hemisphere.")
            return
        try:
            raw = self.server.brainmap_context().render_erf(hemi, style=style)
        except Exception as exc:  # noqa: BLE001
            self.send_error(HTTPStatus.INTERNAL_SERVER_ERROR, f"Could not render ERF brainmap: {exc}")
            return
        self.send_response(HTTPStatus.OK)
        self.send_header("content-type", "image/png")
        self.send_header("content-length", str(len(raw)))
        self.send_header("cache-control", "no-store")
        self.end_headers()
        self.wfile.write(raw)

    def log_message(self, format: str, *args: Any) -> None:
        return


def normalize_choice(value: Any, allowed: tuple[str, ...]) -> str:
    text = str(value or "unclear")
    return text if text in allowed else "unclear"


def make_feature_key(source: str, component: int) -> str:
    return f"{source}:{int(component)}"


def semantic_search_score(query: str, packet: dict[str, Any], label: dict[str, Any]) -> float:
    query = query.lower().strip()
    if not query:
        return 0.0
    label_text = search_label_text(packet, label)
    evidence_text = str(packet.get("_search_evidence_text") or search_evidence_text(packet))
    all_text = f"{label_text} {evidence_text}".lower()
    if query in all_text:
        score = 30.0 if query in label_text.lower() else 18.0
    else:
        score = 0.0

    query_tokens = semantic_tokens(query)
    if not query_tokens:
        return score
    label_tokens = semantic_tokens(label_text)
    evidence_tokens = set(packet.get("_search_evidence_tokens") or semantic_tokens(evidence_text))
    all_tokens = label_tokens | evidence_tokens
    expanded_terms = expanded_query_terms(query_tokens)
    tags = set(label.get("semantic_tags") or broad_semantic_tags(packet, label))

    required = [token for token in query_tokens if token not in STOPWORDS]
    if required and not any(search_token_matches(token, label_tokens, evidence_tokens, tags) for token in required):
        return score

    for token in query_tokens:
        if token in tags:
            score += 25.0
        if token in label_tokens:
            score += 8.0
        elif token in evidence_tokens:
            score += 4.0
        else:
            prefix_match = token_prefix_match(token, all_tokens)
            if prefix_match:
                score += 2.0

    for token in expanded_terms - query_tokens:
        if token in label_tokens:
            score += 2.0

    if len(query_tokens) > 1:
        covered = sum(1 for token in query_tokens if search_token_matches(token, label_tokens, evidence_tokens, tags))
        score += 3.0 * covered / max(1, len(query_tokens))
    return score


def search_token_matches(token: str, label_tokens: set[str], evidence_tokens: set[str], tags: set[str]) -> bool:
    if token in tags or token in label_tokens or token in evidence_tokens:
        return True
    if token_prefix_match(token, label_tokens | evidence_tokens):
        return True
    return bool(SEMANTIC_EXPANSIONS.get(token, set()) & label_tokens)


def broad_semantic_tags(packet: dict[str, Any], label: dict[str, Any], *, max_tags: int = 3) -> list[str]:
    label_text = " ".join(str(label.get(key) or "") for key in ("working_label", "notes"))
    label_tokens = semantic_tokens(label_text)
    evidence_tokens = set(packet.get("_search_evidence_tokens") or semantic_tokens(search_evidence_text(packet)))
    scored = []
    for group, raw_tokens in SEMANTIC_GROUP_TOKENS.items():
        group_tokens = {normalize_search_token(token) for token in raw_tokens}
        group_tokens = {token for token in group_tokens if token and token not in STOPWORDS}
        label_hits = group_tokens & label_tokens
        evidence_hits = group_tokens & evidence_tokens
        score = 3.0 * len(label_hits) + 0.35 * min(len(evidence_hits), 8)
        if group in label_tokens:
            score += 8.0
        if score >= 2.5 and (label_hits or score >= 5.0):
            scored.append((score, group))
    scored.sort(key=lambda item: (-item[0], item[1]))
    return [group for _score, group in scored[: max(0, int(max_tags))]]


def normalize_semantic_tags(value: Any) -> list[str]:
    allowed = set(SEMANTIC_GROUP_TOKENS)
    if value is None:
        return []
    if isinstance(value, str):
        raw_values = re.split(r"[,;\s]+", value)
    elif isinstance(value, (list, tuple, set)):
        raw_values = list(value)
    else:
        raw_values = [value]
    tags = []
    for raw in raw_values:
        tag = normalize_search_token(str(raw))
        if tag in allowed and tag not in tags:
            tags.append(tag)
    return tags


def search_label_text(packet: dict[str, Any], label: dict[str, Any]) -> str:
    return " ".join(
        str(value or "")
        for value in (
            packet.get("packet_id"),
            packet.get("feature_label"),
            packet.get("feature_source"),
            packet.get("side"),
            packet.get("component"),
            label.get("working_label"),
            label.get("notes"),
            " ".join(label.get("semantic_tags") or []),
            label.get("confidence"),
            label.get("interpretation_type"),
        )
    )


def search_evidence_text(packet: dict[str, Any]) -> str:
    parts = []
    for key in ("strong_examples", "near_zero_examples", "opposite_examples"):
        for example in packet.get(key) or []:
            parts.extend(
                str(example.get(field) or "")
                for field in (
                    "story",
                    "target_text",
                    "snippet",
                    "sentence_text",
                    "left_text",
                    "right_text",
                    "context_text",
                )
            )
    erf = packet.get("erf") or {}
    parts.extend(str(value or "") for value in erf.values())
    return " ".join(parts)


def semantic_tokens(text: str) -> set[str]:
    tokens = set()
    for raw in re.findall(r"[a-z0-9']+", text.lower()):
        token = normalize_search_token(raw)
        if token and token not in STOPWORDS:
            tokens.add(token)
    return tokens


def expanded_query_terms(tokens: set[str]) -> set[str]:
    expanded = set(tokens)
    for token in tokens:
        expanded.update(SEMANTIC_EXPANSIONS.get(token, ()))
    return {normalize_search_token(token) for token in expanded if normalize_search_token(token)}


def normalize_search_token(token: str) -> str:
    token = token.lower().strip("'")
    if token.endswith("'s"):
        token = token[:-2]
    if len(token) > 6 and token.endswith("ing"):
        token = token[:-3]
    elif len(token) > 5 and token.endswith("ed"):
        token = token[:-2]
    elif len(token) > 5 and token.endswith("es"):
        token = token[:-2]
    elif len(token) > 4 and token.endswith("s") and not token.endswith(("ss", "us", "is")):
        token = token[:-1]
    return token


def token_prefix_match(query_token: str, tokens: set[str]) -> bool:
    if len(query_token) < 4:
        return False
    return any(token.startswith(query_token) or query_token.startswith(token) for token in tokens if len(token) >= 4)


def save_json_atomic(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with tempfile.NamedTemporaryFile("w", encoding="utf-8", dir=path.parent, delete=False) as handle:
        json.dump(payload, handle, indent=2, sort_keys=True)
        handle.write("\n")
        tmp = Path(handle.name)
    tmp.replace(path)


def query_bool(query: dict[str, list[str]], key: str, *, default: bool) -> bool:
    value = query.get(key, [None])[0]
    if value is None or value == "":
        return bool(default)
    return str(value).strip().lower() not in {"0", "false", "no", "off"}


def sort_packets(packets: list[dict[str, Any]], *, sort_by: str, sort_then: str, sort_dir: str) -> list[dict[str, Any]]:
    reverse = sort_dir != "asc"
    def packet_key(packet: dict[str, Any]) -> tuple[float, float, float, int, int]:
        return (
            sort_value(packet, sort_by),
            sort_value(packet, sort_then),
            sort_value(packet, "selected_voxels"),
            -int(packet.get("component") or 0),
            1 if packet.get("side") == "positive" else 0,
        )

    pinned = [packet for packet in packets if packet.get("highlighted")]
    unpinned = [packet for packet in packets if not packet.get("highlighted")]
    return sorted(pinned, key=packet_key, reverse=reverse) + sorted(unpinned, key=packet_key, reverse=reverse)


def sort_value(packet: dict[str, Any], key: str) -> float:
    if key == "selected_voxels":
        return float(packet.get("encoding_summary", {}).get("selected_voxels") or 0)
    if key == "activation_count":
        summary = packet.get("encoding_summary", {})
        return float(summary.get("activation_count") or summary.get("selected_rows") or 0)
    if key == "max_score":
        summary = packet.get("encoding_summary", {})
        return float(summary.get("max_coefficient") or summary.get("max_score") or 0)
    if key == "erf":
        erf = packet.get("erf") or {}
        try:
            return float(erf.get("mean_erf"))
        except (TypeError, ValueError):
            return -1.0
    if key == "confidence":
        confidence = str(packet.get("label", {}).get("confidence") or "unclear")
        return {"high": 3.0, "medium": 2.0, "low": 1.0, "unclear": 0.0}.get(confidence, 0.0)
    if key == "search_score":
        return float(packet.get("search_score") or 0.0)
    if key == "component":
        return float(packet.get("component") or 0)
    return 0.0


def png_bytes(image: Any) -> bytes:
    buffer = BytesIO()
    image.save(buffer, format="PNG")
    return buffer.getvalue()


def load_erf(stage_a_root: Path, layer: int) -> dict[int, dict[str, str]]:
    for path in (
        stage_a_root / "plots" / f"layer_{layer:02d}_component_erf.csv",
        stage_a_root / "erf" / f"layer_{layer:02d}_component_erf.csv",
    ):
        rows = load_erf_csv(path)
        if rows:
            return rows
    return {}


def load_sae_erf(path: Path | None, sae_coeff_dir: Path | None, layer: int) -> dict[int, dict[str, str]]:
    if path is None and sae_coeff_dir is not None:
        path = sae_coeff_dir.parent / "erf" / f"layer_{layer:02d}_sae_component_erf.csv"
    if path is None:
        return {}
    return load_erf_csv(path)


def load_erf_csv(path: Path) -> dict[int, dict[str, str]]:
    if not path.exists():
        return {}
    rows = {}
    with path.open(newline="", encoding="utf-8") as handle:
        for row in csv.DictReader(handle):
            rows[int(row["component"])] = dict(row)
    return rows


class ICBrainmapContext:
    def __init__(
        self,
        *,
        layer: int,
        encoding_dir: Path,
        ica_contribution_dir: Path | None,
        sae_coeff_dir: Path | None,
        noise_ceiling_npz: Path | None,
        ica_predictivity_npz: Path | None,
        sae_predictivity_npz: Path | None,
        ncsnr_threshold: float,
        mask_reliable_vertices: bool,
        stage_a_root: Path,
        roi_dir: Path,
        output_dir: Path,
        pycortex_db: Path,
        freesurfer_subjects: Path,
        fsaverage5: Path,
        flatmap_height: int,
    ) -> None:
        scripts_dir = Path(__file__).resolve().parent
        if str(scripts_dir) not in sys.path:
            sys.path.insert(0, str(scripts_dir))
        import numpy as np
        from PIL import Image, ImageDraw
        from plot_predictivity_brainmaps import (
            HEMI_HALF,
            HEMI_LABEL,
            composite_rgba_on_white,
            flatmap_roi_masks,
            load_subject_flatmap,
            load_subject_roi_masks,
            native_to_flatmap,
            project_fsaverage5_to_native,
        )
        from flatmap_plotting import (
            colorize_positive,
            colorize_scalar,
            crop_box_for_half,
            half_mask,
            load_font,
            render_colorbar,
            render_positive_colorbar,
            text_size,
        )
        from litcoder_style_plotting import litcoder_style_png_bytes
        from roi_overlay import draw_roi_legend, overlay_roi_borders

        self.np = np
        self.Image = Image
        self.ImageDraw = ImageDraw
        self.HEMI_HALF = HEMI_HALF
        self.HEMI_LABEL = HEMI_LABEL
        self.composite_rgba_on_white = composite_rgba_on_white
        self.native_to_flatmap = native_to_flatmap
        self.project_fsaverage5_to_native = project_fsaverage5_to_native
        self.colorize_positive = colorize_positive
        self.colorize_scalar = colorize_scalar
        self.crop_box_for_half = crop_box_for_half
        self.half_mask = half_mask
        self.load_font = load_font
        self.render_colorbar = render_colorbar
        self.render_positive_colorbar = render_positive_colorbar
        self.text_size = text_size
        self.draw_roi_legend = draw_roi_legend
        self.overlay_roi_borders = overlay_roi_borders
        self.litcoder_style_png_bytes = litcoder_style_png_bytes

        self.layer = int(layer)
        self.output_dir = output_dir
        self.fsaverage5 = fsaverage5
        self.flatmap = load_subject_flatmap(SUBJECT, pycortex_db, freesurfer_subjects, fsaverage5, flatmap_height)
        roi_masks = load_subject_roi_masks(roi_dir, SUBJECT)
        flat_roi_masks = flatmap_roi_masks(roi_masks, self.flatmap)
        self.roi_groups = self.requested_roi_groups(flat_roi_masks)
        self.encoding_dir = encoding_dir
        self.ica_contribution_dir = ica_contribution_dir
        self.ica_predictivity_npz = ica_predictivity_npz
        self.sae_predictivity_npz = sae_predictivity_npz
        self.component_side_values, self.vmax = self.load_component_side_values(encoding_dir / "selected_ic_records.csv")
        self.ica_contribution_ids, self.ica_contribution_values, self.ica_contribution_vmax, self.ica_contribution_manifest = (
            self.load_ica_contribution_maps(ica_contribution_dir, encoding_dir)
        )
        self.ica_contribution_signature = self.ica_contribution_file_signature()
        self.sae_coeff_dir = sae_coeff_dir
        self.sae_feature_ids, self.sae_coefficients, self.sae_vmax, self.sae_coefficient_scale = self.load_sae_coefficients(sae_coeff_dir)
        self.ncsnr_threshold = float(ncsnr_threshold)
        self.reliability_mask = self.load_ncsnr_mask(noise_ceiling_npz, self.ncsnr_threshold)
        self.ica_significant_mask = self.load_positive_predictivity_mask(ica_predictivity_npz)
        self.ica_predictivity_mtime = self.file_mtime(ica_predictivity_npz)
        self.sae_significant_mask = self.load_positive_predictivity_mask(sae_predictivity_npz)
        self.sae_predictivity_mtime = self.file_mtime(sae_predictivity_npz)
        self.mask_reliable_vertices = bool(mask_reliable_vertices)
        self.erf_values, self.erf_vmax, self.erf_voxel_count = self.load_top5_erf_values(
            encoding_dir / "selected_ic_records.csv",
            self.component_erf_csv(stage_a_root, self.layer),
        )
        self.flatmap_background_rgba = (226, 230, 235, 255)
        self.unselected_voxel_rgba = (0, 0, 0, 255)

    def file_mtime(self, path: Path | None) -> float:
        return path.stat().st_mtime if path is not None and path.exists() else 0.0

    def requested_roi_groups(self, flat_roi_masks: dict[str, Any]) -> list[dict[str, object]]:
        specs = [
            ("AC", "auditory cortex", (0, 220, 255), ("AC",)),
            ("Language", "temporal/parietal/frontal", (255, 150, 0), ("temporal_language", "parietal_language", "language_frontal")),
            ("High-level visual", "functional visual ROIs", (255, 35, 210), ("high_level_visual",)),
            ("Theory of mind", "approximate ToM network", (170, 105, 255), ("theory_of_mind",)),
            ("Multiple demand", "frontoparietal MD", (0, 180, 85), ("multiple_demand",)),
            ("Default mode", "Yeo default mode", (255, 60, 100), ("default_mode",)),
        ]
        groups = []
        for label, description, color, keys in specs:
            present = [key for key in keys if key in flat_roi_masks]
            if not present:
                continue
            mask = self.np.zeros(self.flatmap.mask.shape, dtype=bool)
            for key in present:
                mask |= flat_roi_masks[key]
            if not self.np.any(mask):
                continue
            groups.append(
                {
                    "label": label,
                    "description": description,
                    "color": color,
                    "rois": present,
                    "mask": mask,
                }
            )
        return groups

    def load_component_side_values(self, path: Path) -> tuple[dict[tuple[int, str], dict[int, float]], float]:
        if not path.exists():
            return {}, 1.0
        values: dict[tuple[int, str], dict[int, float]] = {}
        all_values = []
        with path.open(newline="", encoding="utf-8") as handle:
            for row in csv.DictReader(handle):
                coef = float(row["signed_coef_sum"])
                if coef == 0:
                    continue
                side = "positive" if coef > 0 else "negative"
                magnitude = abs(coef)
                component = int(row["ic"])
                vertex = int(row["voxel_index"])
                if not 0 <= vertex < N_VERTICES:
                    continue
                component_values = values.setdefault((component, side), {})
                component_values[vertex] = max(component_values.get(vertex, 0.0), magnitude)
                all_values.append(magnitude)
        if not all_values:
            return values, 1.0
        vmax = float(self.np.percentile(self.np.asarray(all_values, dtype=self.np.float32), 99.0))
        if not self.np.isfinite(vmax) or vmax <= 0:
            vmax = 1.0
        return values, vmax

    def ica_contribution_candidates(self, path: Path | None, encoding_dir: Path) -> list[Path]:
        candidates = []
        if path is not None:
            candidates.append(path)
        candidates.extend(
            [
                encoding_dir / "signed_unique_contribution_maps",
                encoding_dir.parent / "signed_unique_contribution_maps",
                BRAIN_ENCODING / "results/gemma_pile10k_ica_labels_layer13_residpost/signed_unique_contribution_maps",
                BRAIN_ENCODING / "results/gemma_ica_labels_UTS03_layer13_residpost/signed_unique_contribution_maps",
            ]
        )
        return candidates

    def ica_contribution_file_signature(self) -> tuple[tuple[str, float], ...]:
        out = []
        for directory in self.ica_contribution_candidates(self.ica_contribution_dir, self.encoding_dir):
            for name in (
                "component_ids.npy",
                "signed_unique_partial_correlation_thresholded.npy",
                "signed_unique_partial_correlation.npy",
                "manifest.json",
            ):
                path = directory / name
                if path.exists():
                    out.append((str(path), path.stat().st_mtime))
        return tuple(out)

    def reload_ica_contribution_maps_if_changed(self) -> None:
        signature = self.ica_contribution_file_signature()
        if signature != self.ica_contribution_signature:
            self.ica_contribution_ids, self.ica_contribution_values, self.ica_contribution_vmax, self.ica_contribution_manifest = (
                self.load_ica_contribution_maps(self.ica_contribution_dir, self.encoding_dir)
            )
            self.ica_contribution_signature = signature

    def reload_predictivity_masks_if_changed(self) -> None:
        ica_mtime = self.file_mtime(self.ica_predictivity_npz)
        if ica_mtime != self.ica_predictivity_mtime:
            self.ica_significant_mask = self.load_positive_predictivity_mask(self.ica_predictivity_npz)
            self.ica_predictivity_mtime = ica_mtime
        sae_mtime = self.file_mtime(self.sae_predictivity_npz)
        if sae_mtime != self.sae_predictivity_mtime:
            self.sae_significant_mask = self.load_positive_predictivity_mask(self.sae_predictivity_npz)
            self.sae_predictivity_mtime = sae_mtime

    def load_ica_contribution_maps(self, path: Path | None, encoding_dir: Path) -> tuple[Any, Any, float, dict[str, Any]]:
        candidates = self.ica_contribution_candidates(path, encoding_dir)
        chosen = None
        values_path = None
        for candidate in candidates:
            ids_candidate = candidate / "component_ids.npy"
            thresholded_candidate = candidate / "signed_unique_partial_correlation_thresholded.npy"
            raw_candidate = candidate / "signed_unique_partial_correlation.npy"
            value_candidate = thresholded_candidate if thresholded_candidate.exists() else raw_candidate
            if ids_candidate.exists() and value_candidate.exists():
                chosen = candidate
                values_path = value_candidate
                break
        if chosen is None or values_path is None:
            return None, None, 1.0, {}

        ids_path = chosen / "component_ids.npy"
        component_ids = self.np.load(ids_path, mmap_mode="r")
        values = self.np.load(values_path, mmap_mode="r")
        manifest = {}
        manifest_path = chosen / "manifest.json"
        if manifest_path.exists():
            try:
                manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
            except Exception:
                manifest = {}
        vmax = float(manifest.get("vmax_abs_p99_thresholded") or 1.0)
        if not self.np.isfinite(vmax) or vmax <= 0:
            nonzero = self.np.asarray(values[self.np.abs(values) > 0], dtype=self.np.float32)
            vmax = float(self.np.percentile(self.np.abs(nonzero), 99.0)) if nonzero.size else 1.0
        if not self.np.isfinite(vmax) or vmax <= 0:
            vmax = 1.0
        manifest["loaded_path"] = str(chosen)
        manifest["loaded_values"] = values_path.name
        return component_ids, values, vmax, manifest

    def load_sae_coefficients(self, path: Path | None) -> tuple[Any, Any, float, float]:
        if path is None:
            return None, None, 1.0, 1.0
        ids_path = path / "feature_ids.npy"
        coef_path = path / "coef_magnitude.npy"
        quantized_path = path / "coef_magnitude_uint8.npy"
        manifest_path = path / "manifest.json"
        if not ids_path.exists():
            return None, None, 1.0, 1.0
        use_quantized = not coef_path.exists() and quantized_path.exists()
        if not coef_path.exists() and not use_quantized:
            return None, None, 1.0, 1.0
        feature_ids = self.np.load(ids_path, mmap_mode="r")
        coefficients = self.np.load(quantized_path if use_quantized else coef_path, mmap_mode="r")
        vmax = 1.0
        scale = 1.0
        if manifest_path.exists():
            try:
                manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
                vmax = float(manifest.get("vmax") or manifest.get("scale_vmax") or vmax)
                if use_quantized:
                    scale_denominator = float(manifest.get("scale_denominator") or 255.0)
                    scale_vmax = float(manifest.get("scale_vmax") or vmax)
                    scale = scale_vmax / scale_denominator if scale_denominator > 0 else 1.0
            except Exception:
                vmax = 1.0
        if not self.np.isfinite(vmax) or vmax <= 0:
            vmax = 1.0
        return feature_ids, coefficients, vmax, scale

    def load_ncsnr_mask(self, path: Path | None, threshold: float) -> Any:
        if path is None or not path.exists():
            return None
        payload = self.np.load(path, allow_pickle=False)
        if "ncsnr" in payload.files:
            ncsnr = self.np.asarray(payload["ncsnr"], dtype=self.np.float32)
        elif "signal_power" in payload.files and "noise_power" in payload.files:
            signal = self.np.asarray(payload["signal_power"], dtype=self.np.float32)
            noise = self.np.asarray(payload["noise_power"], dtype=self.np.float32)
            ncsnr = self.np.sqrt(
                self.np.divide(
                    self.np.maximum(signal, 0.0),
                    noise,
                    out=self.np.zeros_like(signal, dtype=self.np.float32),
                    where=noise > 0,
                )
            )
        else:
            return None
        if ncsnr.shape != (N_VERTICES,):
            raise ValueError(f"{path} NCSNR array has shape {ncsnr.shape}, expected {(N_VERTICES,)}")
        mask = self.np.isfinite(ncsnr) & (ncsnr > float(threshold))
        if "coverage_mask" in payload.files:
            mask &= self.np.asarray(payload["coverage_mask"], dtype=bool)
        return mask

    def load_positive_predictivity_mask(self, path: Path | None) -> Any:
        if path is None or not path.exists():
            return None
        payload = self.np.load(path, allow_pickle=False)
        if "significant_mask" not in payload.files or "correlations" not in payload.files:
            return None
        significant = self.np.asarray(payload["significant_mask"], dtype=bool)
        correlations = self.np.asarray(payload["correlations"], dtype=self.np.float32)
        if significant.shape != (N_VERTICES,) or correlations.shape != (N_VERTICES,):
            raise ValueError(f"{path} predictivity arrays have shape {significant.shape}/{correlations.shape}, expected {(N_VERTICES,)}")
        return significant & self.np.isfinite(correlations) & (correlations > 0.0)

    def load_component_erfs(self, path: Path) -> dict[int, float]:
        if not path.exists():
            raise FileNotFoundError(f"Missing component ERF file: {path}")
        out: dict[int, float] = {}
        with path.open(newline="", encoding="utf-8") as handle:
            for row in csv.DictReader(handle):
                value = float(row["mean_erf"])
                if self.np.isfinite(value):
                    out[int(row["component"])] = value
        return out

    def component_erf_csv(self, stage_a_root: Path, layer: int) -> Path:
        candidates = [
            stage_a_root / "plots" / f"layer_{int(layer):02d}_component_erf.csv",
            stage_a_root / "erf" / f"layer_{int(layer):02d}_component_erf.csv",
        ]
        return next((path for path in candidates if path.exists()), candidates[0])

    def load_top5_erf_values(self, selected_path: Path, erf_path: Path) -> tuple[Any, float, int]:
        if not selected_path.exists() or not erf_path.exists():
            return self.np.zeros(N_VERTICES, dtype=self.np.float32), 11.0, 0
        component_erfs = self.load_component_erfs(erf_path)
        ranked_by_vertex: dict[int, list[tuple[float, float]]] = {}
        with selected_path.open(newline="", encoding="utf-8") as handle:
            for row in csv.DictReader(handle):
                vertex = int(row["voxel_index"])
                if not 0 <= vertex < N_VERTICES:
                    continue
                component = int(row["ic"])
                erf = component_erfs.get(component)
                if erf is None:
                    continue
                score = abs(float(row["signed_coef_sum"]))
                ranked_by_vertex.setdefault(vertex, []).append((score, erf))

        values = self.np.zeros(N_VERTICES, dtype=self.np.float32)
        voxel_count = 0
        for vertex, entries in ranked_by_vertex.items():
            top = sorted(entries, key=lambda item: item[0], reverse=True)[:5]
            if not top:
                continue
            values[vertex] = float(self.np.mean([erf for _score, erf in top]))
            voxel_count += 1
        return values, 11.0, voxel_count

    def render(
        self,
        source: str,
        component: int,
        side: str,
        hemi: str,
        *,
        style: str = "flatmap",
        mask_reliable_vertices: bool | None = None,
        mask_significant_vertices: bool | None = None,
    ) -> bytes:
        values = self.np.zeros(N_VERTICES, dtype=self.np.float32)
        metric_label = "ridge coefficient magnitude"
        scale_label = "ridge coefficient magnitude, shared p99 scale"
        signed_scale = False
        intrinsic_mask_labels = []
        if source == "sae":
            self.reload_predictivity_masks_if_changed()
            if side != "positive":
                raise ValueError("SAE brainmaps are only defined for the positive side.")
            if self.sae_feature_ids is None or self.sae_coefficients is None:
                self.sae_feature_ids, self.sae_coefficients, self.sae_vmax, self.sae_coefficient_scale = self.load_sae_coefficients(self.sae_coeff_dir)
            if self.sae_feature_ids is None or self.sae_coefficients is None:
                raise FileNotFoundError("SAE coefficient maps are not available yet.")
            matches = self.np.flatnonzero(self.sae_feature_ids == int(component))
            if matches.size == 0:
                raise KeyError(f"SAE feature {component} has no coefficient map.")
            values = self.np.asarray(self.sae_coefficients[int(matches[0])], dtype=self.np.float32) * float(self.sae_coefficient_scale)
            vmax = self.sae_vmax
        else:
            self.reload_ica_contribution_maps_if_changed()
            self.reload_predictivity_masks_if_changed()
            if self.ica_contribution_ids is not None and self.ica_contribution_values is not None:
                matches = self.np.flatnonzero(self.ica_contribution_ids == int(component))
                if matches.size == 0:
                    raise KeyError(f"ICA component {component} has no signed contribution map.")
                values = self.np.asarray(self.ica_contribution_values[int(matches[0])], dtype=self.np.float32)
                vmax = self.ica_contribution_vmax
                signed_scale = True
                metric_label = "signed unique partial r"
                scale_label = "signed IC-specific held-out contribution"
                if self.ica_contribution_manifest.get("loaded_values", "").endswith("_thresholded.npy"):
                    alpha = self.ica_contribution_manifest.get("permutation_alpha")
                    if alpha is not None:
                        intrinsic_mask_labels.append(f"IC permutation p<={float(alpha):g}")
                    else:
                        intrinsic_mask_labels.append("IC permutation null")
            else:
                component_values = self.component_side_values.get((int(component), side), {})
                for vertex, value in component_values.items():
                    values[vertex] = value
                vmax = self.vmax

        use_reliability_mask = self.mask_reliable_vertices if mask_reliable_vertices is None else bool(mask_reliable_vertices)
        use_significant_mask = use_reliability_mask if mask_significant_vertices is None else bool(mask_significant_vertices)
        mask_labels = list(intrinsic_mask_labels)
        if use_reliability_mask or use_significant_mask:
            values = self.np.asarray(values, dtype=self.np.float32).copy()
        if use_reliability_mask:
            if self.reliability_mask is not None:
                values[~self.reliability_mask] = 0.0
                mask_labels.append(f"NCSNR>{self.ncsnr_threshold:g}")
        if use_significant_mask:
            significant_mask = self.sae_significant_mask if source == "sae" else self.ica_significant_mask
            if significant_mask is not None:
                values[~significant_mask] = 0.0
                mask_labels.append("model-significant positive predictivity")
        n_selected = int(self.np.count_nonzero(values))
        max_value = float(self.np.max(self.np.abs(values))) if signed_scale and values.size else float(self.np.max(values)) if values.size else 0.0
        if style == "litcoder":
            return self.write_litcoder_image(
                values=values,
                source=source,
                component=int(component),
                side=side,
                n_selected=n_selected,
                max_value=max_value,
                vmax=vmax,
                metric_label=metric_label,
                mask_labels=mask_labels,
                signed_scale=signed_scale,
            )
        if style != "flatmap":
            raise ValueError(f"Unknown brainmap style: {style}")

        native_values = self.project_fsaverage5_to_native(values, self.flatmap)
        flat_values = self.native_to_flatmap(native_values, self.flatmap)
        return self.write_image(
            flat_values=flat_values,
            source=source,
            component=int(component),
            side=side,
            hemi=hemi,
            n_selected=n_selected,
            max_value=max_value,
            vmax=vmax,
            metric_label=metric_label,
            scale_label=scale_label,
            mask_labels=mask_labels,
            mask_requested=use_reliability_mask or use_significant_mask,
            signed_scale=signed_scale,
        )

    def render_erf(self, hemi: str, *, style: str = "flatmap") -> bytes:
        values = self.np.asarray(self.erf_values, dtype=self.np.float32).copy()
        if self.reliability_mask is not None:
            values[~self.reliability_mask] = 0.0
        if self.ica_significant_mask is not None:
            values[~self.ica_significant_mask] = 0.0
        if style == "litcoder":
            title = f"{SUBJECT} Gemma-2-2B L{self.layer} top-5 IC mean ERF; NCSNR>{self.ncsnr_threshold:g} + significant predictivity"
            return self.litcoder_style_png_bytes(
                values,
                title=title,
                fsaverage5=self.fsaverage5,
                positive=True,
                vmax=self.erf_vmax,
                cmap="viridis",
            )
        if style != "flatmap":
            raise ValueError(f"Unknown ERF brainmap style: {style}")
        native_values = self.project_fsaverage5_to_native(values, self.flatmap)
        flat_values = self.native_to_flatmap(native_values, self.flatmap)
        return self.write_erf_image(flat_values=flat_values, hemi=hemi, n_selected=int(self.np.count_nonzero(values)))

    def write_litcoder_image(
        self,
        *,
        values: Any,
        source: str,
        component: int,
        side: str,
        n_selected: int,
        max_value: float,
        vmax: float,
        metric_label: str,
        mask_labels: list[str],
        signed_scale: bool,
    ) -> bytes:
        feature_name = "SAE" if source == "sae" else "ICA"
        mask_label = " + ".join(mask_labels) if mask_labels else "unmasked"
        title = (
            f"{SUBJECT} Gemma-2-2B L{self.layer} {feature_name} {component} {side}; "
            f"{metric_label}; {mask_label}; {n_selected} nonzero; max {max_value:.3g}"
        )
        return self.litcoder_style_png_bytes(
            values,
            title=title,
            fsaverage5=self.fsaverage5,
            positive=not signed_scale,
            vmax=vmax,
            cmap="magma" if not signed_scale else None,
        )

    def write_image(
        self,
        *,
        flat_values: Any,
        source: str,
        component: int,
        side: str,
        hemi: str,
        n_selected: int,
        max_value: float,
        vmax: float,
        metric_label: str,
        scale_label: str,
        mask_labels: list[str],
        mask_requested: bool,
        signed_scale: bool,
    ) -> bytes:
        half = self.HEMI_HALF[hemi]
        panel_mask = self.half_mask(self.flatmap.mask, self.flatmap.split, half)
        crop = self.crop_box_for_half(self.flatmap.mask, self.flatmap.split, half, margin=0)
        crop_mask = panel_mask[crop]
        roi_groups = [group for group in self.roi_groups if self.np.any(self.np.asarray(group["mask"], dtype=bool) & panel_mask)]

        display_values = self.np.asarray(flat_values[crop], dtype=self.np.float32).copy()
        display_values[~self.np.isfinite(display_values)] = self.np.nan
        value_mask = crop_mask & self.np.isfinite(display_values)
        if signed_scale:
            value_mask &= self.np.abs(display_values) > 0.0
            panel = self.colorize_scalar(display_values, value_mask, vmax, cmap_name="coolwarm")
        else:
            value_mask &= display_values > 0.0
            panel = self.colorize_positive(display_values, value_mask, vmax, cmap_name="magma")
        panel = self.compose_flatmap_panel(panel, crop_mask)
        panel = self.overlay_roi_borders(panel, roi_groups, crop, crop_mask, width=1, smooth_iterations=1, dot_step=7).convert("RGB")

        pad = 18
        title_h = 44
        label_h = 34
        colorbar_h = 52
        title_font = self.load_font(24, bold=True)
        label_font = self.load_font(18, bold=True)
        small_font = self.load_font(15)
        roi_title_font = self.load_font(15, bold=True)

        probe = self.Image.new("RGB", (1, 1), "white")
        probe_draw = self.ImageDraw.Draw(probe)
        roi_legend_h = self.estimate_roi_legend_height(probe_draw, roi_groups, panel.width, small_font, roi_title_font)
        legend_h = colorbar_h + (roi_legend_h + 10 if roi_groups else 0)
        canvas_w = pad * 2 + panel.width
        canvas_h = pad * 2 + title_h + label_h + panel.height + legend_h

        canvas = self.Image.new("RGB", (canvas_w, canvas_h), "white")
        draw = self.ImageDraw.Draw(canvas)

        feature_name = "SAE" if source == "sae" else "ICA"
        title = f"{SUBJECT} {self.HEMI_LABEL[hemi]} {feature_name} {component} {side} map"
        draw.text((pad, pad), title, fill=(20, 20, 20), font=title_font)
        mask_label = " + ".join(mask_labels) if mask_labels else ("unmasked; mask unavailable" if mask_requested else "unmasked")
        max_prefix = "max |" if signed_scale else "max "
        max_suffix = "|" if signed_scale else ""
        label = f"Gemma-2-2B L{self.layer}; {mask_label}; {n_selected} nonzero voxels; {max_prefix}{metric_label}{max_suffix} {max_value:.3g}"
        label_w, _ = self.text_size(draw, label, label_font)
        y0 = pad + title_h
        draw.text((pad + max((panel.width - label_w) // 2, 0), y0), label, fill=(20, 20, 20), font=label_font)
        canvas.paste(panel, (pad, y0 + label_h))

        legend_y = y0 + label_h + panel.height + 14
        bar_w = min(720, canvas_w - pad * 2)
        bar = self.render_colorbar(bar_w, 18, cmap_name="coolwarm") if signed_scale else self.render_positive_colorbar(bar_w, 18, cmap_name="magma")
        canvas.paste(bar.convert("RGB"), (pad, legend_y))
        draw.rectangle((pad, legend_y, pad + bar_w - 1, legend_y + 17), outline=(40, 40, 40), width=1)
        left = f"-{vmax:.3g}" if signed_scale else "0"
        draw.text((pad, legend_y + 22), left, fill=(20, 20, 20), font=small_font)
        right = f"+{vmax:.3g}"
        right_w, _ = self.text_size(draw, right, small_font)
        draw.text((pad + bar_w - right_w, legend_y + 22), right, fill=(20, 20, 20), font=small_font)
        mid = scale_label
        mid_w, _ = self.text_size(draw, mid, small_font)
        draw.text((pad + max((bar_w - mid_w) // 2, 0), legend_y + 22), mid, fill=(20, 20, 20), font=small_font)
        if roi_groups:
            self.draw_roi_legend(draw, roi_groups, pad, legend_y + 54, canvas_w - pad * 2, small_font, roi_title_font)

        return png_bytes(canvas)

    def write_erf_image(self, *, flat_values: Any, hemi: str, n_selected: int) -> bytes:
        half = self.HEMI_HALF[hemi]
        panel_mask = self.half_mask(self.flatmap.mask, self.flatmap.split, half)
        crop = self.crop_box_for_half(self.flatmap.mask, self.flatmap.split, half, margin=0)
        crop_mask = panel_mask[crop]
        roi_groups = [group for group in self.roi_groups if self.np.any(self.np.asarray(group["mask"], dtype=bool) & panel_mask)]

        display_values = self.np.asarray(flat_values[crop], dtype=self.np.float32).copy()
        display_values[~self.np.isfinite(display_values)] = self.np.nan
        value_mask = crop_mask & self.np.isfinite(display_values) & (display_values > 0.0)
        panel = self.colorize_positive(display_values, value_mask, self.erf_vmax, cmap_name="viridis")
        panel = self.compose_flatmap_panel(panel, crop_mask)
        panel = self.overlay_roi_borders(panel, roi_groups, crop, crop_mask, width=1, smooth_iterations=1, dot_step=7).convert("RGB")

        pad = 18
        title_h = 44
        label_h = 34
        colorbar_h = 52
        title_font = self.load_font(24, bold=True)
        label_font = self.load_font(18, bold=True)
        small_font = self.load_font(15)
        roi_title_font = self.load_font(15, bold=True)

        probe = self.Image.new("RGB", (1, 1), "white")
        probe_draw = self.ImageDraw.Draw(probe)
        roi_legend_h = self.estimate_roi_legend_height(probe_draw, roi_groups, panel.width, small_font, roi_title_font)
        legend_h = colorbar_h + (roi_legend_h + 10 if roi_groups else 0)
        canvas_w = pad * 2 + panel.width
        canvas_h = pad * 2 + title_h + label_h + panel.height + legend_h

        canvas = self.Image.new("RGB", (canvas_w, canvas_h), "white")
        draw = self.ImageDraw.Draw(canvas)

        title = f"{SUBJECT} {self.HEMI_LABEL[hemi]} top-5 IC mean ERF"
        draw.text((pad, pad), title, fill=(20, 20, 20), font=title_font)
        label = f"Gemma-2-2B layer {self.layer}; NCSNR>{self.ncsnr_threshold:g} + significant predictivity; {n_selected} voxels"
        label_w, _ = self.text_size(draw, label, label_font)
        y0 = pad + title_h
        draw.text((pad + max((panel.width - label_w) // 2, 0), y0), label, fill=(20, 20, 20), font=label_font)
        canvas.paste(panel, (pad, y0 + label_h))

        legend_y = y0 + label_h + panel.height + 14
        bar_w = min(720, canvas_w - pad * 2)
        bar = self.render_positive_colorbar(bar_w, 18, cmap_name="viridis")
        canvas.paste(bar.convert("RGB"), (pad, legend_y))
        draw.rectangle((pad, legend_y, pad + bar_w - 1, legend_y + 17), outline=(40, 40, 40), width=1)
        draw.text((pad, legend_y + 22), "0", fill=(20, 20, 20), font=small_font)
        right = f"{self.erf_vmax:.0f}"
        right_w, _ = self.text_size(draw, right, small_font)
        draw.text((pad + bar_w - right_w, legend_y + 22), right, fill=(20, 20, 20), font=small_font)
        mid = "mean component ERF over top 5 selected ICs per voxel"
        mid_w, _ = self.text_size(draw, mid, small_font)
        draw.text((pad + max((bar_w - mid_w) // 2, 0), legend_y + 22), mid, fill=(20, 20, 20), font=small_font)
        if roi_groups:
            self.draw_roi_legend(draw, roi_groups, pad, legend_y + 54, canvas_w - pad * 2, small_font, roi_title_font)

        return png_bytes(canvas)

    def compose_flatmap_panel(self, layer: Any, crop_mask: Any) -> Any:
        base = self.Image.new("RGBA", layer.size, self.flatmap_background_rgba)
        rgba = self.np.asarray(base).copy()
        rgba[crop_mask] = self.unselected_voxel_rgba
        base = self.Image.fromarray(rgba, mode="RGBA")
        base.alpha_composite(layer.convert("RGBA"))
        return base

    def estimate_roi_legend_height(self, draw: Any, roi_groups: list[dict[str, object]], max_width: int, font: Any, title_font: Any) -> int:
        if not roi_groups:
            return 0
        title = "ROI borders:"
        title_w, _ = self.text_size(draw, title, title_font)
        cursor_x = title_w + 18
        cursor_y = 1
        line_h = 24
        for group in roi_groups:
            label = f"{group['label']}: {group['description']}"
            label_w, _ = self.text_size(draw, label, font)
            item_w = 34 + label_w + 22
            if cursor_x + item_w > max_width and cursor_x > title_w + 18:
                cursor_x = 0
                cursor_y += line_h
            cursor_x += item_w
        return cursor_y + line_h + 6


INDEX_HTML = r"""<!doctype html>
<html>
<head>
  <meta charset="utf-8">
  <title>Feature Label Review</title>
  <style>
    :root { color-scheme: light; --border:#c8cdd4; --muted:#5f6875; --bg:#f7f8fa; --accent:#244f8f; --highlight:#fff3bf; }
    body { margin:0; font:14px/1.4 system-ui, -apple-system, Segoe UI, sans-serif; color:#17202a; background:white; }
    header { min-height:48px; display:flex; align-items:center; gap:8px; flex-wrap:wrap; padding:8px 12px; border-bottom:1px solid var(--border); background:#fff; position:sticky; top:0; z-index:2; }
    header strong { white-space:nowrap; margin-right:4px; }
    header input[type=text], header select { height:30px; border:1px solid var(--border); border-radius:4px; padding:0 8px; background:white; max-width:190px; }
    header input[type=text] { width:180px; min-width:140px; flex:1 1 160px; }
    header select { width:auto; min-width:112px; }
    .toolbar-check { display:flex; align-items:center; gap:5px; margin:0; font-weight:600; white-space:nowrap; color:#26313d; }
    .toolbar-check input { width:auto; min-width:0; margin:0; }
    main { display:grid; grid-template-columns:minmax(300px, 360px) minmax(0, 1fr); min-height:calc(100vh - 65px); }
    #list { border-right:1px solid var(--border); background:var(--bg); overflow:auto; max-height:calc(100vh - 49px); }
    .row { padding:10px 12px; border-bottom:1px solid #e3e6ea; cursor:pointer; }
    .row.highlighted { background:var(--highlight); }
    .row.active { background:#e8f0fb; border-left:4px solid var(--accent); padding-left:8px; }
    .row.active.highlighted { background:#ffe8a3; }
    .row small { color:var(--muted); display:block; margin-top:3px; }
    #detail { padding:18px 22px 80px; overflow:auto; max-height:calc(100vh - 49px); }
    .meta { display:flex; gap:8px; flex-wrap:wrap; margin:8px 0 14px; }
    .pill { border:1px solid var(--border); border-radius:999px; padding:2px 8px; color:#26313d; background:#fff; }
    .panel { border:1px solid var(--border); border-radius:6px; padding:12px; margin:12px 0; background:white; }
    .example { border-top:1px solid #edf0f3; padding:10px 0; }
    .example:first-child { border-top:0; }
    mark { background:#ffec99; padding:1px 2px; border-radius:2px; }
    textarea, input[type=text], select { width:100%; box-sizing:border-box; border:1px solid var(--border); border-radius:4px; padding:8px; font:inherit; background:white; }
    textarea { min-height:82px; resize:vertical; }
    label { display:block; font-weight:600; margin:10px 0 4px; }
    button { border:1px solid #1f4d8a; background:#245a9d; color:white; border-radius:4px; padding:6px 10px; cursor:pointer; white-space:nowrap; }
    button.secondary { background:white; color:#1f4d8a; }
    .grid { display:grid; grid-template-columns:1fr 1fr; gap:12px; }
    .brainmaps { display:grid; grid-template-columns:repeat(auto-fit, minmax(320px, 1fr)); gap:12px; }
    .brainmap img { width:100%; border:1px solid #dfe3e8; border-radius:4px; background:white; display:block; }
    .brainmap h4 { margin:0 0 6px; font-size:13px; color:#26313d; }
    .brainmap-wide { grid-column:1 / -1; }
    .brainmap-wide img { max-height:760px; object-fit:contain; }
    .muted { color:var(--muted); }
    pre { white-space:pre-wrap; background:#f3f5f7; border:1px solid #dfe3e8; padding:8px; border-radius:4px; }
    @media (max-width: 920px) {
      main { grid-template-columns:1fr; }
      #list, #detail { max-height:none; }
      #list { border-right:0; border-bottom:1px solid var(--border); }
      .brainmaps { grid-template-columns:1fr; }
    }
  </style>
</head>
<body>
<header>
  <strong>Feature Label Review</strong>
  <input id="q" placeholder="Semantic search labels/examples">
  <select id="source_filter"><option value="all">All sources</option></select>
  <select id="status"><option value="all">All</option><option value="unapproved">Unapproved</option><option value="approved">Approved</option></select>
  <select id="side"><option value="all">Both signs</option><option value="positive">Positive</option><option value="negative">Negative</option></select>
  <select id="type_filter"><option value="all">All types</option></select>
  <select id="semantic_filter"><option value="all">All semantic tags</option></select>
  <select id="highlight_filter"><option value="all">All pins</option><option value="highlighted">Pinned</option><option value="unhighlighted">Unpinned</option></select>
    <select id="sort_by"><option value="selected_voxels">Sort selected voxels</option><option value="search_score">Sort relevance</option><option value="activation_count">Sort activations</option><option value="max_score">Sort max score</option><option value="erf">Sort ERF</option><option value="confidence">Sort confidence</option><option value="component">Sort feature</option></select>
  <select id="sort_then"><option value="erf">Then ERF</option><option value="search_score">Then relevance</option><option value="selected_voxels">Then voxels</option><option value="activation_count">Then activations</option><option value="max_score">Then max score</option><option value="confidence">Then confidence</option><option value="component">Then feature</option></select>
  <select id="sort_dir"><option value="desc">Desc</option><option value="asc">Asc</option></select>
  <label class="toolbar-check" title="Mask maps to vertices with NCSNR above 0.2 and global model-level significant positive held-out predictivity"><input id="mask_reliable" type="checkbox" checked> NCSNR + sig mask</label>
  <button class="secondary" onclick="prevPage()">Prev</button>
  <button class="secondary" onclick="nextPage()">Next</button>
  <button class="secondary" onclick="loadState()">Refresh</button>
  <span id="counts" class="muted"></span>
</header>
<main>
  <section id="list"></section>
  <section id="detail"><p class="muted">Select a feature.</p></section>
</main>
<script>
let state = null;
let current = null;
let offset = 0;
const limit = 100;
const $ = id => document.getElementById(id);
async function loadState() {
  const params = new URLSearchParams({
    q:$('q').value,
    source:$('source_filter').value,
    status:$('status').value,
    side:$('side').value,
    interpretation_type:$('type_filter').value,
    semantic_tag:$('semantic_filter').value,
    highlight_filter:$('highlight_filter').value,
    sort_by:$('sort_by').value,
    sort_then:$('sort_then').value,
    sort_dir:$('sort_dir').value,
    offset,
    limit
  });
  state = await (await fetch('/api/state?' + params)).json();
  populateTypeFilter();
  populateSemanticFilter();
  populateSourceFilter();
  const start = state.shown_packets ? state.offset + 1 : 0;
  const end = state.offset + state.returned_packets;
  $('counts').textContent = `${start}-${end}/${state.shown_packets} shown, ${state.total_packets} total, ${state.approved_count} approved, ${state.highlighted_count} pinned`;
  renderList();
  if (current) {
    const found = state.packets.find(p => p.packet_id === current.packet_id);
    if (found) selectPacket(found);
  }
}
function renderList() {
  $('list').innerHTML = state.packets.map(p => {
    const featureLabel = p.feature_label || (p.feature_source === 'sae' ? 'SAE feature' : 'IC');
    const label = p.label?.working_label || '(unlabeled)';
    const ok = p.label?.approved ? 'approved' : 'draft';
    const conf = p.label?.confidence || 'unclear';
    const type = p.label?.interpretation_type || 'unclear';
    const erf = p.erf?.mean_erf ?? 'pending';
    const classes = ['row', current?.packet_id===p.packet_id ? 'active' : '', p.highlighted ? 'highlighted' : ''].filter(Boolean).join(' ');
    const marker = p.highlighted ? ' [H]' : '';
    return `<div class="${classes}" onclick="selectById('${p.packet_id}')">
      <strong>${escapeHtml(featureLabel)} ${p.component}${marker}</strong> ${p.side}<br>
      ${escapeHtml(label)}
      <small>${escapeHtml(p.feature_source || 'ica')} · ${ok} · ${escapeHtml(conf)} · ${escapeHtml(type)} · ${tagsText(p)}${searchText(p)}ERF ${escapeHtml(erf)} · ${featureVoxelText(p)} · activations ${p.encoding_summary?.activation_count ?? p.encoding_summary?.selected_rows ?? 0}</small>
    </div>`;
  }).join('');
}
function populateTypeFilter() {
  const select = $('type_filter');
  if (select.dataset.ready === '1') return;
  const currentValue = select.value || 'all';
  select.innerHTML = '<option value="all">All types</option>' + state.interpretation_types.map(v => `<option value="${escapeAttr(v)}">${escapeHtml(v)}</option>`).join('');
  select.value = currentValue;
  select.dataset.ready = '1';
}
function populateSourceFilter() {
  const select = $('source_filter');
  if (select.dataset.ready === '1') return;
  const currentValue = select.value || 'all';
  select.innerHTML = '<option value="all">All sources</option>' + state.sources.map(v => `<option value="${escapeAttr(v)}">${escapeHtml(v)}</option>`).join('');
  select.value = currentValue;
  select.dataset.ready = '1';
}
function populateSemanticFilter() {
  const select = $('semantic_filter');
  if (select.dataset.ready === '1') return;
  const currentValue = select.value || 'all';
  select.innerHTML = '<option value="all">All semantic tags</option>' + state.semantic_tags.map(v => `<option value="${escapeAttr(v)}">${escapeHtml(v)}</option>`).join('');
  select.value = currentValue;
  select.dataset.ready = '1';
}
function selectById(id) {
  const packet = state.packets.find(p => p.packet_id === id);
  if (packet) selectPacket(packet);
}
function selectPacket(packet) {
  current = packet;
  renderList();
  const label = packet.label || {};
  const featureLabel = packet.feature_label || (packet.feature_source === 'sae' ? 'SAE feature' : 'IC');
  $('detail').innerHTML = `
    <h2>${escapeHtml(featureLabel)} ${packet.component} · ${packet.side}</h2>
    <div class="meta">
      <span class="pill">${escapeHtml(packet.feature_source || 'ica')}</span>
      <span class="pill">Layer ${packet.layer}</span>
      <span class="pill">${featureVoxelText(packet)}</span>
      <span class="pill">Activations ${packet.encoding_summary?.activation_count ?? packet.encoding_summary?.selected_rows ?? 0}</span>
      <span class="pill">Max score ${formatMaybe(packet.encoding_summary?.max_coefficient ?? packet.encoding_summary?.max_score)}</span>
      <span class="pill">Rows ${packet.encoding_summary?.selected_rows ?? 0}</span>
      <span class="pill">ERF ${packet.erf?.mean_erf ?? 'pending'}</span>
      ${(label.semantic_tags || []).map(tag => `<span class="pill">${escapeHtml(tag)}</span>`).join('')}
      ${packet.highlighted ? '<span class="pill">Pinned</span>' : ''}
    </div>
    <button class="secondary" onclick="toggleHighlight()">${packet.highlighted ? 'Unpin feature' : 'Pin feature'}</button>
    ${packet.erf ? `<pre>${escapeHtml(JSON.stringify(packet.erf, null, 2))}</pre>` : '<p class="muted">ERF pending or not applicable for this feature source.</p>'}
    ${['ica', 'sae'].includes(packet.feature_source || 'ica') ? brainmapPanel(packet) : ''}
    <div class="panel">
      <div class="grid">
        <div><label>Working Label</label><input id="working_label" type="text" value="${escapeAttr(label.working_label || '')}"></div>
        <div><label>Confidence</label>${selectHtml('confidence', state.confidence_values, label.confidence || 'unclear')}</div>
      </div>
      <label>Interpretation Type</label>${selectHtml('interpretation_type', state.interpretation_types, label.interpretation_type || 'unclear')}
      <label>Broad Semantic Tags</label><input id="semantic_tags" type="text" value="${escapeAttr((label.semantic_tags || []).join(', '))}" list="semantic_tag_options">
      <datalist id="semantic_tag_options">${state.semantic_tags.map(v => `<option value="${escapeAttr(v)}"></option>`).join('')}</datalist>
      <label>Notes / Evidence Check</label><textarea id="notes">${escapeHtml(label.notes || '')}</textarea>
      <label><input id="approved" type="checkbox" ${label.approved ? 'checked' : ''}> Approved</label>
      <button onclick="saveLabel()">Save Label</button>
    </div>
    ${examplesPanel('Strong same-sign examples', packet.strong_examples)}
    ${examplesPanel('Near-zero examples', packet.near_zero_examples)}
    ${examplesPanel('Opposite-sign examples', packet.opposite_examples)}
  `;
}
function brainmapPanel(packet) {
  const source = packet.feature_source || 'ica';
  const mask = $('mask_reliable')?.checked === false ? '0' : '1';
  const base = `/api/feature-brainmap?source=${encodeURIComponent(source)}&component=${encodeURIComponent(packet.component)}&side=${encodeURIComponent(packet.side)}&mask_reliable=${mask}&mask_significant=${mask}&v=signed_unique_ic_v1`;
  const label = source === 'sae' ? 'SAE Coefficient' : 'ICA Signed Unique Contribution';
  const maskLabel = mask === '1' ? 'NCSNR > 0.2 + significant predictivity' : 'unmasked';
  return `<div class="panel">
    <h3>UTS03 ${escapeHtml(label)} Maps <span class="muted">(${maskLabel})</span></h3>
    <div class="brainmaps">
      <div class="brainmap"><h4>Flatmap LH</h4><img src="${base}&style=flatmap&hemi=lh" loading="lazy" alt="UTS03 left hemisphere ${escapeAttr(source)} ${escapeAttr(packet.component)} flatmap"></div>
      <div class="brainmap"><h4>Flatmap RH</h4><img src="${base}&style=flatmap&hemi=rh" loading="lazy" alt="UTS03 right hemisphere ${escapeAttr(source)} ${escapeAttr(packet.component)} flatmap"></div>
      <div class="brainmap brainmap-wide"><h4>LitCoder Surface</h4><img src="${base}&style=litcoder&hemi=both" loading="lazy" alt="UTS03 ${escapeAttr(source)} ${escapeAttr(packet.component)} LitCoder style surface map"></div>
    </div>
  </div>`;
}
function examplesPanel(title, examples) {
  return `<div class="panel"><h3>${title}</h3>${(examples || []).map(ex => `
    <div class="example">
      <div class="muted">#${ex.rank} · ${ex.story} token ${ex.token_index} · score ${Number(ex.score).toFixed(4)} · target ${escapeHtml(JSON.stringify(ex.target_text))}</div>
      <div>${exampleText(ex)}</div>
    </div>`).join('')}</div>`;
}
function featureVoxelText(packet) {
  return (packet.feature_source || 'ica') === 'sae'
    ? 'Dense coefficient map'
    : `Selected voxels ${packet.encoding_summary?.selected_voxels ?? 0}`;
}
function searchText(packet) {
  return Number(packet.search_score || 0) > 0 ? `relevance ${Number(packet.search_score).toFixed(1)} · ` : '';
}
function tagsText(packet) {
  const tags = packet.label?.semantic_tags || [];
  return tags.length ? `tags ${tags.map(escapeHtml).join(', ')} · ` : '';
}
function exampleText(ex) {
  if (ex.left_text || ex.right_text || ex.target_text) {
    return `${escapeHtml(ex.left_text)}<mark>${escapeHtml(ex.target_text)}</mark>${escapeHtml(ex.right_text)}`;
  }
  if (ex.sentence_text) return highlightTarget(ex.sentence_text, ex.target_text);
  if (ex.context_text) return highlightTarget(ex.context_text, ex.target_text);
  return '';
}
function highlightTarget(text, target) {
  const source = String(text ?? '');
  const needle = String(target ?? '').trim();
  if (!source || !needle) return escapeHtml(source);
  const sourceLower = source.toLowerCase();
  const needleLower = needle.toLowerCase();
  let idx = sourceLower.lastIndexOf(needleLower);
  if (idx < 0 && needleLower.startsWith(' ')) idx = sourceLower.lastIndexOf(needleLower.trimStart());
  if (idx < 0) return escapeHtml(source);
  const matched = source.slice(idx, idx + needle.length);
  return `${escapeHtml(source.slice(0, idx))}<mark>${escapeHtml(matched)}</mark>${escapeHtml(source.slice(idx + needle.length))}`;
}
async function saveLabel() {
  const payload = {
    packet_id: current.packet_id,
    component: current.component,
    feature_source: current.feature_source || 'ica',
    side: current.side,
    working_label: $('working_label').value,
    confidence: $('confidence').value,
    interpretation_type: $('interpretation_type').value,
    semantic_tags: $('semantic_tags').value.split(',').map(s => s.trim()).filter(Boolean),
    notes: $('notes').value,
    approved: $('approved').checked
  };
  await fetch('/api/label', {method:'POST', headers:{'content-type':'application/json'}, body:JSON.stringify(payload)});
  await loadState();
}
async function toggleHighlight() {
  if (!current) return;
  await fetch('/api/highlight', {
    method:'POST',
    headers:{'content-type':'application/json'},
    body:JSON.stringify({component: current.component, feature_source: current.feature_source || 'ica', highlighted: !current.highlighted})
  });
  await loadState();
}
function resetAndLoad() { offset = 0; loadState(); }
function prevPage() { offset = Math.max(0, offset - limit); loadState(); }
function nextPage() {
  if (!state) return;
  if (offset + limit < state.shown_packets) offset += limit;
  loadState();
}
function selectHtml(id, values, selected) {
  return `<select id="${id}">${values.map(v => `<option value="${escapeAttr(v)}" ${v===selected?'selected':''}>${escapeHtml(v)}</option>`).join('')}</select>`;
}
function escapeHtml(s) { return String(s ?? '').replace(/[&<>"']/g, c => ({'&':'&amp;','<':'&lt;','>':'&gt;','"':'&quot;',"'":'&#39;'}[c])); }
function escapeAttr(s) { return escapeHtml(s).replace(/`/g, '&#96;'); }
function formatMaybe(v) {
  const n = Number(v);
  return Number.isFinite(n) ? n.toPrecision(4) : 'n/a';
}
$('q').addEventListener('input', () => resetAndLoad());
$('source_filter').addEventListener('change', () => resetAndLoad());
$('status').addEventListener('change', () => resetAndLoad());
$('side').addEventListener('change', () => resetAndLoad());
$('type_filter').addEventListener('change', () => resetAndLoad());
$('semantic_filter').addEventListener('change', () => resetAndLoad());
$('highlight_filter').addEventListener('change', () => resetAndLoad());
$('sort_by').addEventListener('change', () => resetAndLoad());
$('sort_then').addEventListener('change', () => resetAndLoad());
$('sort_dir').addEventListener('change', () => resetAndLoad());
$('mask_reliable').addEventListener('change', () => { if (current) selectPacket(current); });
loadState();
</script>
</body>
</html>
"""


if __name__ == "__main__":
    raise SystemExit(main())