visualization.py

"""Visualization utilities for Minerva multimodal contact fingerprints.

The public plotting path accepts named fingerprint results directly:

    fp = model.get_fingerprints(seq, tokenizer)
    plot_fingerprints(fp)

For notebooks that compose their own panels, use ``render_fingerprints(fp)`` to
produce an RGB image. Legacy RGB helpers remain available for compatibility.
"""

from typing import Dict, List, Optional, Union

import numpy as np


# =============================================================================
# Color scheme
# =============================================================================
COLORS = {
    "base_pairing": "#F25560",  # Coral   - RFAM / RNA base pairing
    "protein":      "#2ED5C5",  # Teal    - PDB protein contacts
    "repeat":       "#4F46E5",  # Deep indigo - DNA repeats
    "other":        "#D1D5DB",  # Light gray - unclassified (usually not rendered)
}

# Channels rendered in the overlay, in argmax-stack order.
OVERLAY_CHANNELS = ["base_pairing", "repeat", "protein"]

# RNA base-pairing wins wherever its normalized value exceeds this, so faint
# base-pairing stays visible against strong protein backgrounds.
DEFAULT_OVERLAY_OVERRIDES = {"base_pairing": 0.5}

# Raw Frobenius-norm contact strengths roughly span 0..10 after centering / APC;
# this vmax matches the UG27 / DRT2 fingerprint convention.
DEFAULT_FP_VMAX = 10.0

# Publication palette used by the Minerva UG27 figure scripts.
PUBLICATION_DATASET_COLORS = {
    "rna": (218 / 255, 56 / 255, 50 / 255),       # #DA3832
    "pdb": (125 / 255, 180 / 255, 224 / 255),     # #7DB4E0
    "repeat": (95 / 255, 48 / 255, 140 / 255),    # #5F308C
}
PUBLICATION_JACOBIAN_COLOR = (224 / 255, 159 / 255, 14 / 255)  # #E09F0E
PUBLICATION_OVERLAY_COLORS = {
    "protein": PUBLICATION_DATASET_COLORS["pdb"],
    "base_pairing": PUBLICATION_DATASET_COLORS["rna"],
    "repeat": PUBLICATION_DATASET_COLORS["repeat"],
    "jacobian": PUBLICATION_JACOBIAN_COLOR,
    "other": PUBLICATION_JACOBIAN_COLOR,
}
PUBLICATION_LINE_WIDTH = 0.8


def _hex_to_rgb(c):
    if isinstance(c, str):
        c = c.lstrip("#")
        return np.array([int(c[i:i + 2], 16) / 255 for i in (0, 2, 4)])
    return np.asarray(c, dtype=float)


def _finite_channel(x, *, nan=0.0, posinf=1.0, neginf=0.0):
    """Float32 array with non-finite values replaced for stable rendering."""
    return np.nan_to_num(np.asarray(x, dtype=np.float32), nan=nan, posinf=posinf, neginf=neginf)


def _finite_rgb(rgb):
    """RGB image in [0, 1], with invalid pixels rendered as white."""
    return np.clip(_finite_channel(rgb, nan=1.0, posinf=1.0, neginf=0.0), 0.0, 1.0)


def contact_rgb_overlay(
    channels: Dict[str, np.ndarray],
    channel_order: Optional[List[str]] = None,
    colors: Optional[Dict[str, object]] = None,
    vmin: Union[float, Dict[str, float]] = 0.0,
    vmax: Union[float, Dict[str, float]] = 1.0,
    overrides: Optional[Dict[str, float]] = DEFAULT_OVERLAY_OVERRIDES,
) -> np.ndarray:
    """Winner-take-all RGB overlay of contact channels over white.

    1. Each channel is independently min-max rescaled to [0, 1] via
       ``clip((value - vmin) / (vmax - vmin), 0, 1)``.
    2. Per pixel, argmax picks the dominant channel (no overlap blending).
    3. The painted color is ``1 - (1 - color) * normalized`` -- a subtract-from-
       white blend, so 0 stays white and 1 is fully saturated.
    4. ``overrides`` force a channel to win wherever its normalized value exceeds
       the threshold (default: base-pairing at 0.5).

    Args:
        channels: dict channel name -> (H, W) array.
        channel_order: argmax stack order. Defaults to ``channels`` insertion order.
        colors: channel name -> hex/(R,G,B). Defaults to module ``COLORS``.
        vmin, vmax: per-channel normalization bounds (float or {channel: bound}).
        overrides: {channel: threshold} forced wins, or None for pure argmax.

    Returns:
        (H, W, 3) float array in [0, 1] for ``imshow``.
    """
    if not channels:
        raise ValueError("channels is empty")

    order = channel_order if channel_order is not None else list(channels.keys())
    order = [ch for ch in order if ch in channels]
    if not order:
        raise ValueError("no channels in `channel_order` matched `channels`")

    palette = colors if colors is not None else COLORS

    def _resolve(arg, ch, default):
        if isinstance(arg, dict):
            return float(arg.get(ch, default))
        return float(arg) if arg is not None else default

    first = channels[order[0]]
    h, w = first.shape

    stack = np.empty((len(order), h, w), dtype=np.float32)
    for i, ch in enumerate(order):
        lo = _resolve(vmin, ch, 0.0)
        hi = _resolve(vmax, ch, 1.0)
        denom = max(hi - lo, 1e-12)
        values = _finite_channel(channels[ch], nan=lo, posinf=hi, neginf=lo)
        stack[i] = np.clip((values - lo) / denom, 0.0, 1.0)

    max_idx = np.argmax(stack, axis=0)
    max_val = np.max(stack, axis=0)

    if overrides:
        for ch, thr in overrides.items():
            if ch not in order:
                continue
            ch_idx = order.index(ch)
            mask = stack[ch_idx] >= thr
            max_idx[mask] = ch_idx
            max_val[mask] = stack[ch_idx][mask]

    palette_rgb = np.array([_hex_to_rgb(palette[ch]) for ch in order], dtype=np.float32)
    target_color = palette_rgb[max_idx]                    # (H, W, 3)
    rgb = 1.0 - (1.0 - target_color) * max_val[..., np.newaxis]
    return rgb


def _to_numpy_array(x):
    if hasattr(x, "detach"):
        x = x.detach().cpu().numpy()
    return np.asarray(x)


def _overlay_channel_name(name: str) -> str:
    aliases = {
        "basepairing": "base_pairing",
        "base_pair": "base_pairing",
        "bp": "base_pairing",
    }
    key = str(name).lower()
    return aliases.get(key, key)


def _default_fingerprint_channel_names(count: int) -> List[str]:
    if count == 3:
        return ["basepairing", "repeat", "other"]
    if count == 4:
        return ["basepairing", "repeat", "protein", "other"]
    if count == 5:
        return ["bp_forward", "bp_reverse", "repeat", "protein", "other"]
    return [f"channel_{i}" for i in range(count)]


def _fingerprint_channels(fingerprints, channel_names: Optional[List[str]] = None):
    if hasattr(fingerprints, "channels"):
        channels = dict(fingerprints.channels)
        names = channel_names or list(getattr(fingerprints, "channel_names", channels.keys()))
    elif isinstance(fingerprints, dict):
        channels = dict(fingerprints)
        names = channel_names or list(channels.keys())
    elif isinstance(fingerprints, tuple) and len(fingerprints) == 2:
        arr, names = fingerprints
        names = list(channel_names or names)
        arr = _to_numpy_array(arr)
        channels = {name: arr[i] for i, name in enumerate(names)}
    else:
        arr = _to_numpy_array(fingerprints)
        if arr.ndim != 3:
            raise ValueError(
                "fingerprints must be a FingerprintResult, dict, or (C, L, L) array"
            )
        names = list(channel_names or _default_fingerprint_channel_names(arr.shape[0]))
        if len(names) != arr.shape[0]:
            raise ValueError(
                "channel_names length must match fingerprints channel dimension: "
                f"{len(names)} != {arr.shape[0]}"
            )
        channels = {name: arr[i] for i, name in enumerate(names)}

    normalized = {}
    for name, values in channels.items():
        normalized[_overlay_channel_name(name)] = _to_numpy_array(values)

    if "base_pairing" not in normalized:
        bp_parts = [normalized[ch] for ch in ("bp_forward", "bp_reverse") if ch in normalized]
        if bp_parts:
            normalized["base_pairing"] = np.sum(np.stack(bp_parts, axis=0), axis=0)
    return normalized


def render_fingerprints(
    fingerprints,
    channel_names: Optional[List[str]] = None,
    *,
    style: str = "default",
    include_other: Optional[bool] = None,
    colors: Optional[Dict[str, object]] = None,
    vmin: Union[float, Dict[str, float]] = 0.0,
    vmax: Union[float, Dict[str, float]] = DEFAULT_FP_VMAX,
    overrides: Optional[Dict[str, float]] = DEFAULT_OVERLAY_OVERRIDES,
) -> np.ndarray:
    """Render named fingerprint channels as a winner-take-all RGB overlay.

    Args:
        fingerprints: ``FingerprintResult``, ``{channel: map}``, or a channel-first
            ``(C, L, L)`` tensor/array.
        channel_names: Required only when passing a raw array with non-default
            channel order.
        style: ``"default"`` or ``"publication"``. Publication style uses the
            Minerva figure palette and includes the ``other`` channel by default.
        include_other: Whether to include the unclassified channel in the overlay.
            Defaults to True only for ``style="publication"``.
        colors, vmin, vmax, overrides: Forwarded to :func:`contact_rgb_overlay`.

    Returns:
        ``(L, L, 3)`` float RGB image in ``[0, 1]``.
    """
    if style not in ("default", "publication"):
        raise ValueError("style must be 'default' or 'publication'")
    if include_other is None:
        include_other = style == "publication"
    if style == "publication" and colors is None:
        colors = PUBLICATION_OVERLAY_COLORS

    by_name = _fingerprint_channels(fingerprints, channel_names=channel_names)
    order = [ch for ch in OVERLAY_CHANNELS if ch in by_name]
    if include_other and "other" in by_name:
        order.append("other")
    if not order:
        raise ValueError(
            f"no renderable fingerprint channels; got {list(by_name)}, "
            f"expected some of {OVERLAY_CHANNELS}"
        )
    return contact_rgb_overlay(
        {ch: by_name[ch] for ch in order},
        channel_order=order,
        colors=colors,
        vmin=vmin,
        vmax=vmax,
        overrides=overrides,
    )


def jacobian_fingerprint_rgb(
    jac,
    tokens: List[str],
    *,
    bp_threshold: Optional[float] = None,
    repeat_threshold: Optional[float] = None,
    protein_threshold: Optional[float] = None,
    aa_start: int = 4,
    jac_aa_order: Optional[List[str]] = None,
    colors: Optional[Dict[str, object]] = None,
    vmin: Union[float, Dict[str, float]] = 0.0,
    vmax: Union[float, Dict[str, float]] = DEFAULT_FP_VMAX,
    overrides: Optional[Dict[str, float]] = DEFAULT_OVERLAY_OVERRIDES,
    include_other: bool = False,
):
    """One-shot Jacobian -> multimodal fingerprint -> RGB overlay.

    Wraps :func:`minerva.jacobian.fingerprint_jacobian_multimodality` (classifies
    each position-pair into base-pairing / repeat / protein / other via cosine
    similarity to the reference fingerprints) and :func:`contact_rgb_overlay`. The
    unclassified "other" channel is dropped from the overlay so only the three
    modality colors appear over white.

    Args:
        jac: Jacobian (L, A, L, A), numpy or torch, full alphabet channels
            (i.e. ``fast=False`` mode of ``get_categorical_jacobian``).
        tokens: L tokens aligned to Jacobian positions.
        bp_threshold, repeat_threshold, protein_threshold: optional similarity
            cutoffs; None -> Minerva defaults.
        aa_start, jac_aa_order: forwarded to the classifier.
        colors, vmin, vmax, overrides: forwarded to :func:`contact_rgb_overlay`.
        include_other: include the classifier's unassigned "other" channel.
            This is the gold Jacobian background used in the Minerva UG27
            publication fingerprint plots.

    Returns:
        (rgb, fingerprints, channel_names): rgb is (L, L, 3) in [0, 1];
        fingerprints is (C, L, L); channel_names lists the classifier channels.
    """
    try:
        from .jacobian import fingerprint_jacobian_multimodality
    except ImportError:  # HF snapshot imported as top-level visualization.py
        from jacobian import fingerprint_jacobian_multimodality

    fp_kwargs = {"split_bp": False, "aa_start": aa_start, "jac_aa_order": jac_aa_order}
    if bp_threshold is not None:
        fp_kwargs["bp_threshold"] = bp_threshold
    if repeat_threshold is not None:
        fp_kwargs["repeat_threshold"] = repeat_threshold
    if protein_threshold is not None:
        fp_kwargs["protein_threshold"] = protein_threshold

    _, fingerprints, channel_names = fingerprint_jacobian_multimodality(
        jac, tokens, **fp_kwargs,
    )
    # Classifier returns 'basepairing'; the palette key is 'base_pairing'.
    aliases = {"basepairing": "base_pairing"}
    norm = [aliases.get(n, n) for n in channel_names]
    by_name = {n: fingerprints[i] for i, n in enumerate(norm)}

    # Render only the three colored modality channels by default. Publication
    # fingerprint plots add "other" as a gold Jacobian-strength background.
    order = [ch for ch in OVERLAY_CHANNELS if ch in by_name]
    if include_other and "other" in by_name:
        order.append("other")
    rgb = contact_rgb_overlay(
        {ch: by_name[ch] for ch in order},
        channel_order=order,
        colors=colors,
        vmin=vmin,
        vmax=vmax,
        overrides=overrides,
    )
    return rgb, fingerprints, channel_names


def head_contacts_rgb(
    contacts: Dict[str, np.ndarray],
    tokens: Optional[List[str]] = None,
    colors: Optional[Dict[str, object]] = None,
    vmin: Union[float, Dict[str, float]] = 0.0,
    vmax: Union[float, Dict[str, float]] = 1.0,
    overrides: Optional[Dict[str, float]] = DEFAULT_OVERLAY_OVERRIDES,
) -> np.ndarray:
    """RGB overlay from regression-head contact maps.

    Args:
        contacts: dict channel name (``base_pairing`` / ``repeat`` / ``protein``)
            -> (L, L) contact map. Head-derived maps are sigmoid probabilities in
            [0, 1], so ``vmax=1.0`` is the natural scale.
        tokens: optional length-L token list. If given, applies token-aware
            masking (matching the Jacobian fingerprint convention): the protein
            channel is kept only on amino-acid/amino-acid pairs, and the
            base-pairing / repeat channels only on nucleotide/nucleotide pairs.
            This removes spurious protein signal on RNA/intergenic regions.
        colors, vmin, vmax, overrides: forwarded to :func:`contact_rgb_overlay`.

    Returns:
        (L, L, 3) float RGB array in [0, 1].
    """
    order = [ch for ch in OVERLAY_CHANNELS if ch in contacts]
    if not order:
        raise ValueError(
            f"no overlay channels in contacts; got keys {list(contacts)}, "
            f"expected some of {OVERLAY_CHANNELS}"
        )
    chan = {ch: _finite_channel(contacts[ch]) for ch in order}
    if tokens is not None and len(tokens) == chan[order[0]].shape[0]:
        tt = np.array(token_types(tokens))
        aa_ij = (tt == "protein")[:, None] & (tt == "protein")[None, :]
        nuc_ij = (tt == "nucleotide")[:, None] & (tt == "nucleotide")[None, :]
        for ch in order:
            mask = aa_ij if ch == "protein" else nuc_ij
            chan[ch] = np.where(mask, chan[ch], 0.0).astype(np.float32, copy=False)
    return contact_rgb_overlay(
        chan, channel_order=order, colors=colors, vmin=vmin, vmax=vmax,
        overrides=overrides,
    )


def _auto_contrast_contact_channels(
    channels: Dict[str, np.ndarray],
    *,
    percentile: float = 99.0,
    min_dynamic: float = 1e-8,
):
    """Scale structured low-probability contact maps without coloring constants."""
    scaled = {}
    vmax = {}
    stats = {}
    for ch, arr in channels.items():
        values = _finite_channel(arr)
        finite = values[np.isfinite(values)]
        if finite.size == 0:
            scaled[ch] = np.zeros_like(values, dtype=np.float32)
            vmax[ch] = 1.0
            stats[ch] = {"kept": False, "reason": "non-finite"}
            continue
        lo = float(np.min(finite))
        hi = float(np.max(finite))
        dynamic = hi - lo
        if dynamic <= min_dynamic:
            scaled[ch] = np.zeros_like(values, dtype=np.float32)
            vmax[ch] = 1.0
            stats[ch] = {"kept": False, "reason": "constant", "min": lo, "max": hi}
            continue
        positive = finite[finite > 0]
        ref = positive if positive.size else finite
        p = float(np.percentile(ref, percentile))
        if not np.isfinite(p) or p <= 0:
            p = hi if hi > 0 else 1.0
        scaled[ch] = values
        vmax[ch] = max(p, min_dynamic)
        stats[ch] = {"kept": True, "min": lo, "max": hi, f"p{percentile:g}": p}
    return scaled, vmax, stats


def _mask_contact_channels_by_tokens(
    contacts: Dict[str, np.ndarray],
    tokens: Optional[List[str]] = None,
) -> Dict[str, np.ndarray]:
    """Apply the same token-aware channel mask used by head contact overlays."""
    order = [ch for ch in OVERLAY_CHANNELS if ch in contacts]
    chan = {ch: _finite_channel(contacts[ch]) for ch in order}
    if not order:
        return chan
    if tokens is not None and len(tokens) == chan[order[0]].shape[0]:
        tt = np.array(token_types(tokens))
        aa_ij = (tt == "protein")[:, None] & (tt == "protein")[None, :]
        nuc_ij = (tt == "nucleotide")[:, None] & (tt == "nucleotide")[None, :]
        for ch in order:
            mask = aa_ij if ch == "protein" else nuc_ij
            chan[ch] = np.where(mask, chan[ch], 0.0).astype(np.float32, copy=False)
    return chan


def publication_head_contacts_rgb(
    contacts: Dict[str, np.ndarray],
    tokens: Optional[List[str]] = None,
    auto_contrast: bool = False,
    contrast_percentile: float = 99.0,
    return_stats: bool = False,
    **kwargs,
) -> np.ndarray:
    """Regression-head overlay using the Minerva publication palette.

    By default this renders calibrated 0..1 contact probabilities, matching the
    Minerva publication contact-overlay convention. Set ``auto_contrast=True``
    only for exploratory viewing of very low-probability structure.
    """
    kwargs.setdefault("colors", PUBLICATION_OVERLAY_COLORS)
    rendered_contacts = _mask_contact_channels_by_tokens(contacts, tokens=tokens)
    stats = None
    if auto_contrast:
        rendered_contacts, auto_vmax, stats = _auto_contrast_contact_channels(
            rendered_contacts, percentile=contrast_percentile)
        kwargs.setdefault("vmax", auto_vmax)
    rgb = head_contacts_rgb(rendered_contacts, tokens=None, **kwargs)
    if return_stats:
        return rgb, stats
    return rgb


def publication_jacobian_fingerprint_rgb(jac, tokens: List[str], **kwargs):
    """Jacobian fingerprint overlay matching the Minerva UG27 figure style."""
    kwargs.setdefault("colors", PUBLICATION_OVERLAY_COLORS)
    kwargs.setdefault("include_other", True)
    return jacobian_fingerprint_rgb(jac, tokens, **kwargs)


def legend_handles(channels: Optional[List[str]] = None, colors: Optional[Dict] = None):
    """matplotlib Patch handles for a base-pairing / repeat / protein legend."""
    from matplotlib.patches import Patch

    channels = channels if channels is not None else OVERLAY_CHANNELS
    palette = colors if colors is not None else COLORS
    labels = {"base_pairing": "Base pairing (RNA)", "repeat": "Repeat (DNA)",
              "protein": "Protein"}
    return [Patch(facecolor=palette[ch], edgecolor="none", label=labels.get(ch, ch))
            for ch in channels]


# =============================================================================
# Locus annotation track (protein / nucleotide / special per token)
# =============================================================================
TRACK_COLORS = {
    "protein": COLORS["protein"],       # teal  -- CDS / gene (amino acids)
    "nucleotide": COLORS["repeat"],     # indigo -- intergenic / RNA (nucleotides)
    "special": "#C7CBD1",               # gray  -- orientation / special tokens
}


def token_types(tokens: List[str]) -> List[str]:
    """Classify each token as 'protein' | 'nucleotide' | 'special'.

    Amino acids are upper-case, nucleotides lower-case (a/c/g/t), and
    orientation/special tokens contain '<' or '>'.
    """
    out = []
    for t in tokens:
        if not t or "<" in t or ">" in t:
            out.append("special")
        elif t[:1].islower():
            out.append("nucleotide")
        elif t[:1].isupper():
            out.append("protein")
        else:
            out.append("special")
    return out


def _track_rgb(tokens: List[str]) -> np.ndarray:
    """(L, 3) RGB strip coloring each token by type."""
    return np.array([_hex_to_rgb(TRACK_COLORS[t]) for t in token_types(tokens)],
                    dtype=np.float32)


def plot_locus(
    rgb: np.ndarray,
    tokens: Optional[List[str]] = None,
    title: str = "Minerva locus",
    figsize=(9, 9),
    show_legend: bool = True,
    save: Optional[str] = None,
    dpi: int = 200,
):
    """Static contact overlay with an optional token-type track (top + left).

    ``tokens`` (len L, aligned to the map) draws a protein/nucleotide/special
    strip so gene vs intergenic structure is visible. ``save`` writes the figure
    (e.g. a ``.pdf``) at ``dpi``. Returns the matplotlib Figure.
    """
    import matplotlib.pyplot as plt

    L = rgb.shape[0]
    has_track = tokens is not None and len(tokens) == L
    if not has_track:
        fig, ax = plt.subplots(figsize=figsize)
        ax.imshow(_finite_rgb(rgb), aspect="equal", interpolation="nearest")
        ax.set_title(title, fontsize=11)
        ax.set_xlabel("Position"); ax.set_ylabel("Position")
        if show_legend:
            ax.legend(handles=legend_handles(), loc="upper right", fontsize=9,
                      framealpha=0.9)
    else:
        strip = _track_rgb(tokens)
        fig = plt.figure(figsize=figsize)
        gs = fig.add_gridspec(2, 2, height_ratios=[0.035, 1], width_ratios=[0.035, 1],
                              hspace=0.015, wspace=0.015)
        ax_top = fig.add_subplot(gs[0, 1])
        ax_top.imshow(strip[None, :, :], aspect="auto", interpolation="nearest")
        ax_top.set_xticks([]); ax_top.set_yticks([])
        ax_top.set_title(title, fontsize=11)
        ax_left = fig.add_subplot(gs[1, 0])
        ax_left.imshow(strip[:, None, :], aspect="auto", interpolation="nearest")
        ax_left.set_xticks([]); ax_left.set_yticks([])
        ax = fig.add_subplot(gs[1, 1])
        ax.imshow(_finite_rgb(rgb), aspect="equal", interpolation="nearest")
        ax.set_xticks([]); ax.set_yticks([])
        if show_legend:
            ax.legend(handles=legend_handles(), loc="upper right", fontsize=8,
                      framealpha=0.9)
    if save:
        fig.savefig(save, dpi=dpi, bbox_inches="tight")
    return fig


def interactive_overlay(rgb: np.ndarray, title: str = "Minerva locus"):
    """Zoom/pan/hover view of the RGB overlay via Plotly. Returns a plotly Figure.

    Use in Colab/Jupyter: ``interactive_overlay(rgb).show()``, or
    ``fig.write_html("locus.html")`` for a standalone interactive file.
    """
    import plotly.express as px

    img = (_finite_rgb(rgb) * 255).astype(np.uint8)
    fig = px.imshow(img, title=title)
    fig.update_layout(dragmode="pan", margin=dict(l=10, r=10, t=40, b=10),
                      height=760, width=800)
    fig.update_xaxes(title="Position", constrain="domain")
    fig.update_yaxes(title="Position", scaleanchor="x")
    return fig


def plot_fingerprint_overlay(
    rgb: np.ndarray,
    title: str = "Minerva multimodal fingerprint",
    ax=None,
    extent=None,
    show_legend: bool = True,
    figsize=(9, 9),
    legend_channels: Optional[List[str]] = None,
    colors: Optional[Dict[str, object]] = None,
):
    """imshow an RGB overlay with a channel legend. Returns the Axes."""
    import matplotlib.pyplot as plt

    if ax is None:
        _, ax = plt.subplots(figsize=figsize)
    ax.imshow(_finite_rgb(rgb), aspect="equal", interpolation="nearest",
              extent=extent, origin="upper")
    ax.set_title(title, fontsize=12)
    ax.set_xlabel("Position")
    ax.set_ylabel("Position")
    if show_legend:
        ax.legend(handles=legend_handles(legend_channels, colors=colors),
                  loc="upper right", framealpha=0.9, fontsize=9)
    return ax


def plot_fingerprints(
    fingerprints,
    channel_names: Optional[List[str]] = None,
    *,
    title: str = "Minerva multimodal fingerprint",
    style: str = "default",
    include_other: Optional[bool] = None,
    ax=None,
    extent=None,
    show_legend: bool = True,
    figsize=(9, 9),
    return_rgb: bool = False,
    **render_kwargs,
):
    """Plot a ``FingerprintResult`` or channel-first fingerprint tensor.

    Returns the matplotlib Axes by default. Set ``return_rgb=True`` to get
    ``(ax, rgb)`` for notebook panels that reuse the rendered image.
    """
    rgb = render_fingerprints(
        fingerprints,
        channel_names=channel_names,
        style=style,
        include_other=include_other,
        **render_kwargs,
    )
    colors = render_kwargs.get("colors")
    if style == "publication" and colors is None:
        colors = PUBLICATION_OVERLAY_COLORS
    legend_channels = [ch for ch in OVERLAY_CHANNELS if ch in _fingerprint_channels(fingerprints, channel_names)]
    ax = plot_fingerprint_overlay(
        rgb,
        title=title,
        ax=ax,
        extent=extent,
        show_legend=show_legend,
        figsize=figsize,
        legend_channels=legend_channels,
        colors=colors,
    )
    if return_rgb:
        return ax, rgb
    return ax


def _contact_map_2d(values, batch_index: int = 0) -> np.ndarray:
    arr = _to_numpy_array(values)
    if arr.ndim == 3:
        arr = arr[batch_index]
    if arr.ndim != 2:
        raise ValueError(f"expected a 2D contact map or batched 3D map, got shape {arr.shape}")
    return arr


def render_interactions(
    interactions,
    tokens: Optional[List[str]] = None,
    *,
    style: str = "default",
    batch_index: int = 0,
    colors: Optional[Dict[str, object]] = None,
    vmin: Union[float, Dict[str, float]] = 0.0,
    vmax: Union[float, Dict[str, float]] = 1.0,
    overrides: Optional[Dict[str, float]] = DEFAULT_OVERLAY_OVERRIDES,
) -> np.ndarray:
    """Render model interaction-head outputs as an RGB overlay."""
    if hasattr(interactions, "interactions") and interactions.interactions is not None:
        interactions = interactions.interactions
    if interactions is None:
        raise ValueError("interactions is None; call model(..., output_interactions=True)")
    if style not in ("default", "publication"):
        raise ValueError("style must be 'default' or 'publication'")

    contacts = {
        _overlay_channel_name(name): _contact_map_2d(values, batch_index=batch_index)
        for name, values in dict(interactions).items()
    }
    if style == "publication":
        return publication_head_contacts_rgb(
            contacts,
            tokens=tokens,
            colors=colors or PUBLICATION_OVERLAY_COLORS,
            vmin=vmin,
            vmax=vmax,
            overrides=overrides,
        )
    return head_contacts_rgb(
        contacts,
        tokens=tokens,
        colors=colors,
        vmin=vmin,
        vmax=vmax,
        overrides=overrides,
    )


def plot_interactions(
    interactions,
    tokens: Optional[List[str]] = None,
    *,
    title: str = "Minerva interactions",
    style: str = "default",
    batch_index: int = 0,
    ax=None,
    extent=None,
    show_legend: bool = True,
    figsize=(9, 9),
    return_rgb: bool = False,
    **render_kwargs,
):
    """Plot model interaction-head outputs from ``outputs.interactions``."""
    rgb = render_interactions(
        interactions,
        tokens=tokens,
        style=style,
        batch_index=batch_index,
        **render_kwargs,
    )
    colors = render_kwargs.get("colors")
    if style == "publication" and colors is None:
        colors = PUBLICATION_OVERLAY_COLORS
    ax = plot_fingerprint_overlay(
        rgb,
        title=title,
        ax=ax,
        extent=extent,
        show_legend=show_legend,
        figsize=figsize,
        colors=colors,
    )
    if return_rgb:
        return ax, rgb
    return ax

# =============================================================================
# Dot plots and publication-style single panels
# =============================================================================

def setup_publication_style():
    """Matplotlib rcParams matching the Minerva publication figures."""
    import matplotlib.pyplot as plt

    plt.rcParams["font.family"] = "sans-serif"
    plt.rcParams["font.sans-serif"] = [
        "Helvetica Neue", "Arial", "Helvetica", "DejaVu Sans"
    ]
    plt.rcParams["font.size"] = 12
    plt.rcParams["axes.labelsize"] = 14
    plt.rcParams["axes.titlesize"] = 14
    plt.rcParams["axes.linewidth"] = 1.0
    plt.rcParams["xtick.labelsize"] = 11
    plt.rcParams["ytick.labelsize"] = 11
    plt.rcParams["legend.fontsize"] = 9
    plt.rcParams["hatch.linewidth"] = 0.5
    plt.rcParams["svg.fonttype"] = "none"
    plt.rcParams["pdf.fonttype"] = 42
    plt.rcParams["xtick.direction"] = "out"
    plt.rcParams["ytick.direction"] = "out"
    plt.rcParams["xtick.major.size"] = 5
    plt.rcParams["ytick.major.size"] = 5
    plt.rcParams["xtick.major.width"] = 0.8
    plt.rcParams["ytick.major.width"] = 0.8
    plt.rcParams["xtick.minor.size"] = 3
    plt.rcParams["ytick.minor.size"] = 3
    plt.rcParams["xtick.bottom"] = True
    plt.rcParams["ytick.left"] = True


def compute_dotplot_fwd_rc(seq, window: int = 6, threshold: Optional[int] = None):
    """Forward and reverse-complement self-comparison dot plot kernel.

    Both axes use forward-sequence coordinates. Forward exact k-mer matches mark
    repeat diagonals; reverse-complement matches mark potential stem/base-pairing
    anti-diagonals.
    """
    if threshold is None:
        threshold = window
    seq = seq.upper()
    n = len(seq)
    if n < window:
        empty = np.array([], dtype=int)
        return empty, empty, empty, empty, n
    arr = np.frombuffer(seq.encode(), dtype=np.uint8)

    comp_table = np.zeros(256, dtype=np.uint8)
    for a, b in [(ord("A"), ord("T")), (ord("T"), ord("A")),
                 (ord("C"), ord("G")), (ord("G"), ord("C"))]:
        comp_table[a] = b
    arr_comp = comp_table[arr]

    fwd_r, fwd_c = [], []
    for offset in range(1, n - window + 1):
        a = arr[:n - offset]
        b = arr[offset:]
        if len(a) < window:
            continue
        matches = (a == b).astype(np.int32)
        cs = np.cumsum(matches)
        ws = np.empty(len(matches) - window + 1, dtype=np.int32)
        ws[0] = cs[window - 1]
        ws[1:] = cs[window:] - cs[:len(matches) - window]
        hits = np.where(ws >= threshold)[0]
        for h in hits:
            for k in range(window):
                fwd_r.append(h + k)
                fwd_c.append(h + offset + k)
                fwd_r.append(h + offset + k)
                fwd_c.append(h + k)

    arr_comp_rev = arr_comp[::-1].copy()
    rc_r, rc_c = [], []
    for offset in range(-(n - window), n - window + 1):
        if offset >= 0:
            a = arr[:n - offset]
            b = arr_comp_rev[offset:offset + len(a)]
        else:
            a = arr[-offset:]
            b = arr_comp_rev[:len(a)]
        if len(a) < window:
            continue
        matches = (a == b).astype(np.int32)
        cs = np.cumsum(matches)
        ws = np.empty(len(matches) - window + 1, dtype=np.int32)
        ws[0] = cs[window - 1]
        ws[1:] = cs[window:] - cs[:len(matches) - window]
        hits = np.where(ws >= threshold)[0]
        for h in hits:
            for k in range(window):
                if offset >= 0:
                    ri = h + k
                    ci = n - 1 - (h + offset + k)
                else:
                    ri = h - offset + k
                    ci = n - 1 - (h + k)
                if 0 <= ri < n and 0 <= ci < n:
                    rc_r.append(ri)
                    rc_c.append(ci)

    return (np.array(fwd_r, dtype=int), np.array(fwd_c, dtype=int),
            np.array(rc_r, dtype=int), np.array(rc_c, dtype=int), n)


def dotplot_rgb_from_tokens(
    tokens: List[str],
    word_size: int = 6,
    threshold: Optional[int] = None,
    colors: Optional[Dict[str, object]] = None,
):
    """Build a forward/RC dot plot RGB image for a mixed-token window.

    Non-nucleotide tokens remain white. Forward matches are repeat purple;
    reverse-complement matches are RNA red under the publication palette.
    """
    palette = colors if colors is not None else PUBLICATION_DATASET_COLORS
    dotplot_img = np.ones((len(tokens), len(tokens), 3), dtype=np.float32)
    nuc_subseq = ""
    nuc_sub_to_tok = {}
    for i, tok in enumerate(tokens):
        if len(tok) == 1 and tok.upper() in "ACGTN":
            nuc_sub_to_tok[len(nuc_subseq)] = i
            nuc_subseq += tok.upper()

    if len(nuc_subseq) < word_size:
        return dotplot_img, {"forward": 0, "revcomp": 0,
                             "nucleotide_tokens": len(nuc_subseq)}

    fwd_r, fwd_c, rc_r, rc_c, _ = compute_dotplot_fwd_rc(
        nuc_subseq, window=word_size, threshold=threshold)
    clr_fwd = np.array(palette["repeat"], dtype=np.float32)
    clr_rc = np.array(palette["rna"], dtype=np.float32)

    for r_nuc, c_nuc in zip(fwd_r, fwd_c):
        r_tok = nuc_sub_to_tok.get(int(r_nuc))
        c_tok = nuc_sub_to_tok.get(int(c_nuc))
        if r_tok is not None and c_tok is not None:
            dotplot_img[r_tok, c_tok] = clr_fwd

    for r_nuc, c_nuc in zip(rc_r, rc_c):
        r_tok = nuc_sub_to_tok.get(int(r_nuc))
        c_tok = nuc_sub_to_tok.get(int(c_nuc))
        if r_tok is not None and c_tok is not None:
            dotplot_img[r_tok, c_tok] = clr_rc

    return dotplot_img, {
        "forward": int(len(fwd_r)),
        "revcomp": int(len(rc_r)),
        "nucleotide_tokens": int(len(nuc_subseq)),
    }


def _publication_handles(kind: str = "heads"):
    from matplotlib.patches import Patch

    if kind == "fingerprint":
        items = [
            (PUBLICATION_JACOBIAN_COLOR, "Jacobian"),
            (PUBLICATION_DATASET_COLORS["rna"], "RNA"),
            (PUBLICATION_DATASET_COLORS["repeat"], "Repeats"),
        ]
    elif kind == "dotplot":
        items = [
            (PUBLICATION_DATASET_COLORS["repeat"], "Forward"),
            (PUBLICATION_DATASET_COLORS["rna"], "Rev. comp."),
        ]
    else:
        items = [
            (PUBLICATION_DATASET_COLORS["pdb"], "PDB"),
            (PUBLICATION_DATASET_COLORS["rna"], "RNA"),
            (PUBLICATION_DATASET_COLORS["repeat"], "Repeats"),
        ]
    return [Patch(facecolor=color, edgecolor="black", linewidth=0.4, label=label)
            for color, label in items]


def render_publication_panel(ax, rgb, label: str, show_yticks: bool = True,
                             genome_offset: int = 0):
    """Render one square publication-style image panel."""
    L = rgb.shape[0]
    ax.imshow(_finite_rgb(rgb), aspect="equal", interpolation="nearest")
    tick_step = max(1, L // 4)
    ticks = np.arange(0, L, tick_step)
    tick_labels = [f"{int(t + genome_offset):,}" for t in ticks]
    ax.set_xticks(ticks)
    ax.set_xticklabels(tick_labels, fontsize=6, rotation=0, ha="center")
    ax.set_yticks(ticks)
    ax.set_yticklabels(tick_labels if show_yticks else [], fontsize=6)
    ax.set_title(label, fontsize=11, pad=4)
    for spine in ax.spines.values():
        spine.set_visible(True)
        spine.set_linewidth(PUBLICATION_LINE_WIDTH)
    ax.tick_params(axis="both", which="major", direction="out",
                   length=3, width=PUBLICATION_LINE_WIDTH, labelsize=6,
                   bottom=True, left=True, top=False, right=False)


def plot_dotplot(
    tokens: List[str],
    title: str = "Dot plot",
    *,
    word_size: int = 6,
    threshold: Optional[int] = None,
    ax=None,
    figsize=(5, 5),
    save: Optional[str] = None,
    dpi: int = 600,
):
    """Plot a reusable forward/RC nucleotide dot plot for any token window."""
    setup_publication_style()
    import matplotlib.pyplot as plt

    img, stats = dotplot_rgb_from_tokens(tokens, word_size=word_size,
                                         threshold=threshold)
    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.figure
    render_publication_panel(ax, img, title, show_yticks=True, genome_offset=0)
    leg = ax.legend(handles=_publication_handles("dotplot"), loc="upper right",
                    fontsize=8, frameon=True, framealpha=0.9,
                    edgecolor="black", fancybox=False)
    leg.get_frame().set_linewidth(0.6)
    if save:
        fig.savefig(save, dpi=dpi, bbox_inches="tight", facecolor="white")
    return fig, ax, stats


def plot_publication_locus(
    rgb: np.ndarray,
    title: str = "Minerva locus",
    *,
    genome_offset: int = 0,
    overlay_kind: str = "heads",
    show_legend: bool = True,
    ax=None,
    figsize=(5, 5),
    save: Optional[str] = None,
    dpi: int = 600,
):
    """Single-panel publication-style renderer for arbitrary loci.

    This intentionally does not add UG27-specific dot-plot, triptych, or inset
    panels. It just applies the Minerva figure styling to one RGB overlay.
    """
    setup_publication_style()
    import matplotlib.pyplot as plt

    if ax is None:
        fig, ax = plt.subplots(figsize=figsize)
    else:
        fig = ax.figure
    render_publication_panel(ax, rgb, title, show_yticks=True,
                             genome_offset=genome_offset)
    if show_legend:
        leg = ax.legend(handles=_publication_handles(overlay_kind), loc="upper right",
                        fontsize=8, frameon=True, framealpha=0.9,
                        edgecolor="black", fancybox=False)
        leg.get_frame().set_linewidth(0.6)
    if save:
        fig.savefig(save, dpi=dpi, bbox_inches="tight", facecolor="white")
    return fig, ax


# =============================================================================
# Interactive Bokeh contact viewer (zoom / pan / hover / per-channel sliders)
# =============================================================================
def _rgb01_to_hex(rgb) -> str:
    return "#%02x%02x%02x" % tuple(int(round(max(0.0, min(1.0, c)) * 255)) for c in rgb[:3])


# Publication channel -> color, matching the static overlay / legend.
_BOKEH_CHANNEL_COLOR = {
    "base_pairing": PUBLICATION_DATASET_COLORS["rna"],
    "repeat":       PUBLICATION_DATASET_COLORS["repeat"],
    "protein":      PUBLICATION_DATASET_COLORS["pdb"],
}
_BOKEH_CHANNEL_LABEL = {
    "base_pairing": "RNA (base pairing)",
    "repeat":       "Repeats",
    "protein":      "PDB (protein)",
}


def bokeh_contact_viewer(
    channels: Dict[str, np.ndarray],
    tokens: Optional[List[str]] = None,
    *,
    title: str = "Minerva contacts",
    genome_offset: int = 0,
    prefilter: float = 0.05,
    init_threshold: float = 0.3,
    vmax: Union[float, Dict[str, float]] = 1.0,
    overrides: Optional[Dict[str, float]] = DEFAULT_OVERLAY_OVERRIDES,
    size: int = 760,
    track_px: int = 46,
):
    """Interactive Bokeh contact map: wheel-zoom/pan for per-pixel inspection,
    per-position token-type tracks (top + left), hover tooltips, and one
    real-time threshold slider per channel. Keeps the publication colors.

    Winner-take-all per pixel: each position-pair is assigned to a single
    channel via argmax of the vmax-normalized values (with the base-pairing
    ``overrides`` win), exactly like the static ``contact_rgb_overlay``. The
    per-channel sliders then threshold that channel's won pixels (on the
    normalized 0..1 scale).

    Args:
        channels: dict with any of ``base_pairing`` / ``repeat`` / ``protein``
            -> (L, L) contact-probability arrays.
        tokens: length-L token list; draws the protein/nucleotide/special
            per-position tracks and enables token hover. Optional.
        genome_offset: added to position labels in hover.
        prefilter: contacts below this are dropped before sending to the browser.
        init_threshold: initial slider value per channel.
        vmax: per-channel (or scalar) value mapped to full color saturation.
        size: main panel size in px. track_px: annotation track thickness.

    Returns:
        A Bokeh layout. In a notebook: ``from bokeh.io import output_notebook,
        show; output_notebook(); show(layout)``. To a file: ``save_bokeh_html``.
    """
    from bokeh.plotting import figure
    from bokeh.models import (ColumnDataSource, CustomJS, Slider, HoverTool,
                              Legend, LegendItem, Range1d, WheelZoomTool, Div)
    from bokeh.layouts import gridplot, column, row

    order = [c for c in OVERLAY_CHANNELS if c in channels]
    if not order:
        raise ValueError("channels must include one of base_pairing/repeat/protein")
    # Draw (z-order) bottom -> top: protein background, then repeat, then base
    # pairing on top -- matches the publication overlay's channel priority so
    # dense protein contacts don't bury the sparser repeat / base-pairing signal.
    draw_order = [c for c in ("protein", "repeat", "base_pairing") if c in channels]
    L = int(np.asarray(channels[order[0]]).shape[0])
    if isinstance(vmax, (int, float)):
        vmax = {c: float(vmax) for c in order}

    # Winner-take-all assignment, identical to contact_rgb_overlay: normalize
    # each channel by its vmax, argmax per pixel, then let `overrides` force a
    # channel to win where its normalized value clears the threshold.
    stack = np.stack([np.clip(_finite_channel(channels[c]) / max(vmax[c], 1e-12), 0.0, 1.0)
                      for c in order], axis=0)                     # (C, L, L)
    max_idx = np.argmax(stack, axis=0)
    max_val = np.max(stack, axis=0)
    for ch, thr in (overrides or {}).items():
        if ch in order:
            ci = order.index(ch)
            m = stack[ci] >= thr
            max_idx[m] = ci
            max_val[m] = stack[ci][m]

    def _points(channel, color, _vm=None):
        ci = order.index(channel)
        rows, cols = np.where((max_idx == ci) & (max_val > prefilter))
        vals = max_val[rows, cols]                                 # normalized 0..1
        base = np.asarray(color, dtype=float)
        blended = 1.0 - (1.0 - base)[None, :] * vals[:, None]      # subtract-from-white
        colors = [_rgb01_to_hex(c) for c in blended]
        return dict(i=cols.astype(int).tolist(), j=rows.astype(int).tolist(),
                    val=[round(float(v), 4) for v in vals], color=colors)

    x_range = Range1d(start=0, end=L, bounds=(0, L))
    y_range = Range1d(start=L, end=0, bounds=(0, L))
    p = figure(title=None, x_range=x_range, y_range=y_range,
               width=size, height=size, tools="pan,box_zoom,reset,save",
               toolbar_location="right", output_backend="webgl",
               x_axis_location="below", y_axis_location="left")
    wheel = WheelZoomTool(dimensions="both")
    p.add_tools(wheel); p.toolbar.active_scroll = wheel
    p.grid.visible = False
    p.line([0, L], [0, L], line_color="#BBBBBB", line_width=1, line_dash="dashed")
    p.xaxis.axis_label = "Position"; p.yaxis.axis_label = "Position"

    renderers, full_src, shown_src = {}, {}, {}
    for c in draw_order:   # draw bottom -> top (protein ... base_pairing)
        full = ColumnDataSource(_points(c, _BOKEH_CHANNEL_COLOR[c]))
        vals = np.asarray(full.data["val"])
        keep = vals >= init_threshold if len(vals) else np.zeros(0, bool)
        init = {k: list(np.asarray(v, dtype=object)[keep]) for k, v in full.data.items()}
        shown = ColumnDataSource(init)
        renderers[c] = p.rect(x="i", y="j", width=1.0, height=1.0, source=shown,
                              fill_color="color", line_color=None, fill_alpha=0.95)
        full_src[c], shown_src[c] = full, shown

    sliders = []
    for c in order:        # controls / legend in publication order
        slider = Slider(start=round(prefilter, 3), end=1.0, value=init_threshold, step=0.01,
                        title=_BOKEH_CHANNEL_LABEL[c], width=size // 2,
                        bar_color=_rgb01_to_hex(_BOKEH_CHANNEL_COLOR[c]))
        slider.js_on_change("value", CustomJS(args=dict(full=full_src[c], shown=shown_src[c]), code="""
            const t = cb_obj.value, f = full.data;
            const out = {i:[], j:[], val:[], color:[]};
            for (let k = 0; k < f.val.length; k++) {
                if (f.val[k] >= t) { out.i.push(f.i[k]); out.j.push(f.j[k]);
                    out.val.push(f.val[k]); out.color.push(f.color[k]); }
            }
            shown.data = out; shown.change.emit();
        """))
        sliders.append(slider)

    tips = [("position", "@i, @j"), ("value", "@val")]
    p.add_tools(HoverTool(renderers=list(renderers.values()), tooltips=tips, point_policy="follow_mouse"))

    # Legend OUTSIDE the plot box (to the right).
    legend = Legend(items=[LegendItem(label=_BOKEH_CHANNEL_LABEL[c], renderers=[renderers[c]]) for c in order],
                    location="top", border_line_color=None)
    p.add_layout(legend, "right")

    # Per-position token-type tracks (top + left), sharing the map's ranges.
    top_track = left_track = None
    if tokens is not None and len(tokens) == L:
        types = token_types(tokens)
        tcol = [_rgb01_to_hex(_hex_to_rgb(TRACK_COLORS[t])) for t in types]
        idx = list(range(L))
        top_src = ColumnDataSource(dict(left=idx, right=[i + 1 for i in idx],
                                        top=[1] * L, bottom=[0] * L, color=tcol,
                                        tok=list(tokens), typ=types))
        top_track = figure(x_range=x_range, y_range=Range1d(0, 1), width=size, height=track_px,
                           tools="", toolbar_location=None, output_backend="webgl")
        top_track.quad(left="left", right="right", top="top", bottom="bottom",
                       source=top_src, fill_color="color", line_color=None)
        top_track.add_tools(HoverTool(tooltips=[("pos", "@left"), ("token", "@tok"), ("type", "@typ")]))
        top_track.grid.visible = False; top_track.yaxis.visible = False
        top_track.xaxis.visible = False; top_track.title = title

        left_src = ColumnDataSource(dict(bottom=idx, top=[i + 1 for i in idx],
                                         left=[0] * L, right=[1] * L, color=tcol,
                                         tok=list(tokens), typ=types))
        left_track = figure(x_range=Range1d(0, 1), y_range=y_range, width=track_px, height=size,
                            tools="", toolbar_location=None, output_backend="webgl")
        left_track.quad(left="left", right="right", top="top", bottom="bottom",
                        source=left_src, fill_color="color", line_color=None)
        left_track.grid.visible = False; left_track.xaxis.visible = False; left_track.yaxis.visible = False

    if top_track is not None:
        grid = gridplot([[None, top_track], [left_track, p]], toolbar_location="right",
                        merge_tools=False)
    else:
        grid = gridplot([[p]], toolbar_location="right", merge_tools=False)
    header = Div(text=f"<b>{title}</b>", styles={"font-size": "14px", "margin": "2px 0"})
    return column(header, row(*sliders), grid)


def save_bokeh_html(layout, path: str, title: str = "Minerva contacts") -> str:
    """Save a Bokeh layout to a standalone interactive HTML file. Returns path."""
    from bokeh.io import output_file, save, reset_output
    output_file(path, title=title)
    save(layout)
    reset_output()
    return path