src/genome_annotation/plotting.py

"""Figures of the annotation notebooks (matplotlib, one call per figure)."""

import itertools
from collections.abc import Mapping, Sequence
from typing import Any

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from matplotlib.colors import to_rgb
from matplotlib.figure import Figure
from matplotlib.patches import Rectangle

from . import decoding
from .analysis import cds_span
from .benchmark import TestSpecies
from .evaluation import Gene

# Paper hues; the gold is darkened to sit in the readable lightness band. `correct`
# and `invented` are deliberate neutrals so the three error classes carry the colour.
CLASS_COLORS = {
    "correct": "#b3b3b3",
    "boundary": "#d99a00",
    "missed": "#b428a0",
    "fragmented": "#ff554d",
    "invented": "#4b5563",
}
EXON_COLORS = {
    "exact": CLASS_COLORS["correct"],
    "shifted_one": CLASS_COLORS["boundary"],
    "shifted_both": CLASS_COLORS["fragmented"],
    "missed": CLASS_COLORS["missed"],
}
NTV3_COLOR = "#3973fc"
OTHER_COLOR = "#808080"
INK, MUTED = "#1f2328", "#6b7280"
BAR_HEIGHT = 0.6


def _style(ax: Axes) -> None:
    for side in ("top", "right"):
        ax.spines[side].set_visible(False)
    for side in ("left", "bottom"):
        ax.spines[side].set_color(MUTED)
    ax.tick_params(colors=MUTED, labelcolor=INK)
    ax.grid(axis="x", color="#e5e7eb", linewidth=0.8)
    ax.set_axisbelow(True)


def _text_color(background: str) -> str:
    r, g, b = to_rgb(background)
    return "white" if 0.2126 * r + 0.7152 * g + 0.0722 * b < 0.45 else INK


def _stacked_rows(
    ax: Axes,
    rows: Sequence[str],
    shares: Sequence[Mapping[str, float]],
    colors: Mapping[str, str],
) -> None:
    """Horizontal 100% bars, one per row, labelled inside segments wide enough."""
    for y, share in enumerate(shares):
        left = 0.0
        for klass, color in colors.items():
            width = share.get(klass, 0.0)
            ax.barh(
                y,
                width,
                left=left,
                height=BAR_HEIGHT,
                color=color,
                edgecolor="white",
                linewidth=2,
            )
            if width >= 6:
                ax.text(
                    left + width / 2,
                    y,
                    f"{width:.0f}%",
                    ha="center",
                    va="center",
                    fontsize=8,
                    color=_text_color(color),
                )
            left += width
    ax.set_yticks(range(len(rows)), rows)
    ax.invert_yaxis()
    ax.set_xlim(0, 100)
    ax.set_xlabel("share (%)")
    handles = [Rectangle((0, 0), 1, 1, color=c) for c in colors.values()]
    labels = [k.replace("_", " ") for k in colors]
    ax.legend(
        handles,
        labels,
        ncol=len(colors),
        frameon=False,
        loc="lower left",
        bbox_to_anchor=(0, 1.0),
        fontsize=8,
    )
    _style(ax)


def _percent(counts: Mapping[str, int], keys: Sequence[str]) -> dict[str, float]:
    total = sum(counts[k] for k in keys)
    return {k: 100.0 * counts[k] / total if total else 0.0 for k in keys}


def plot_scores(
    results: Mapping[str, Mapping[str, Mapping[str, float]]],
    species: TestSpecies | None = None,
) -> Figure:
    """Gene and exon F1 of each annotation, with the published NTv3 score as a tick.

    Args:
        results: Annotation label to ``evaluation.evaluate`` output.
        species: Benchmark species, to mark the paper's NTv3 numbers.
    """
    fig, axes = plt.subplots(1, 2, figsize=(9, 0.6 + 0.5 * len(results)), sharey=True)
    labels = list(results)
    for ax, (level, title) in zip(
        axes, (("Gene", "Gene F1 (%)"), ("CDS", "Exon (CDS) F1 (%)")), strict=True
    ):
        values = [100 * results[label][level]["F1"] for label in labels]
        colors = [NTV3_COLOR if "NTv3" in label else OTHER_COLOR for label in labels]
        ax.barh(labels, values, height=BAR_HEIGHT, color=colors)
        for y, value in enumerate(values):
            ax.text(value + 1, y, f"{value:.2f}", va="center", fontsize=9, color=INK)
        if species is not None:
            paper = species.gene_f1 if level == "Gene" else species.cds_f1
            ax.axvline(paper, color=INK, linestyle=":", linewidth=1, label="paper")
            ax.legend(
                frameon=False, fontsize=8, loc="lower right", bbox_to_anchor=(1, 1)
            )
        ax.set_xlim(0, 110)
        ax.set_title(title, loc="left", fontsize=10, color=INK)
        _style(ax)
    axes[0].invert_yaxis()
    fig.tight_layout()
    return fig


def plot_locus_composition(summaries: Mapping[str, Mapping[str, float]]) -> Figure:
    """Where every reference gene ends up, and the share of invented predictions."""
    fig, (ax, side) = plt.subplots(
        1,
        2,
        figsize=(10, 1.2 + 0.6 * len(summaries)),
        gridspec_kw={"width_ratios": (5, 1)},
        sharey=True,
    )
    classes = ("correct", "boundary", "missed", "fragmented")
    shares = [
        _percent({k: int(s[k]) for k in classes}, classes) for s in summaries.values()
    ]
    _stacked_rows(ax, list(summaries), shares, {k: CLASS_COLORS[k] for k in classes})
    ax.set_xlabel("reference genes (%)")
    for y, summary in enumerate(summaries.values()):
        share = 100 * summary["invented"] / summary["n_pred"]
        side.barh(y, share, height=BAR_HEIGHT, color=CLASS_COLORS["invented"])
        side.text(share, y, f" {share:.1f}%", va="center", fontsize=8, color=INK)
    side.set_xlim(0, max(10, 1.4 * side.get_xlim()[1]))
    side.set_xlabel("invented\n(% of predictions)")
    _style(side)
    fig.tight_layout()
    return fig


def plot_recovery_by_size(recovery: Mapping[str, Any]) -> Figure:
    """Locus outcome by CDS span, and exact recovery by exon count."""
    fig, (left, right) = plt.subplots(1, 2, figsize=(11, 3.6))
    span_bins = [_short_span(label) for label in recovery["span_bins"]]
    totals = recovery["span_total"]
    outcome = recovery["span_outcome"]
    shares = [
        _percent({k: outcome[k][i] for k in outcome}, list(outcome))
        for i in range(len(span_bins))
    ]
    rows = [f"{b} (n={n:,})" for b, n in zip(span_bins, totals, strict=True)]
    _stacked_rows(left, rows, shares, {k: CLASS_COLORS[k] for k in outcome})
    left.set_title("by CDS span", loc="left", fontsize=10, color=INK, pad=22)

    bins = recovery["exon_count_bins"]
    total = recovery["exon_total"]
    correct = recovery["exon_correct"]
    rate = [100 * c / t if t else 0.0 for c, t in zip(correct, total, strict=True)]
    right.bar(bins, rate, color=NTV3_COLOR, width=0.6)
    for x, (value, n) in enumerate(zip(rate, total, strict=True)):
        right.text(x, value + 1, f"{value:.0f}%\nn={n:,}", ha="center", fontsize=8)
    right.set_ylim(0, 110)
    right.set_xlabel("coding exons per gene")
    right.set_ylabel("genes recovered exactly (%)")
    right.set_title("by exon count", loc="left", fontsize=10, color=INK)
    _style(right)
    right.grid(axis="x", visible=False)
    right.grid(axis="y", color="#e5e7eb", linewidth=0.8)
    fig.tight_layout()
    return fig


def _short_span(label: str) -> str:
    lo, hi = (int(x) for x in label.split("-"))
    fmt = lambda v: f"{v / 1e6:.0f} Mb" if v >= 1e6 else f"{v / 1e3:.0f} kb"  # noqa: E731
    return (
        f">{fmt(lo)}" if hi >= 10**12 else f"{fmt(lo)}-{fmt(hi)}".replace("0 kb-", "0-")
    )


def plot_exon_taxonomy(exons: Mapping[str, Any]) -> Figure:
    """Reference exons: exact, one or both boundaries wrong, or missed."""
    by_position = exons["by_position"]
    by_length = exons["by_length"]
    classes = list(EXON_COLORS)
    fig, (top, bottom) = plt.subplots(2, 1, figsize=(9, 6.6))
    positions = [
        p
        for p in ("first", "internal", "last", "single")
        if sum(by_position[p].values())
    ]
    totals = exons["totals"]
    _stacked_rows(
        top,
        [f"all exons (n={sum(totals.values()):,})"]
        + [f"{p} (n={sum(by_position[p].values()):,})" for p in positions],
        [_percent(totals, classes)]
        + [_percent(by_position[p], classes) for p in positions],
        EXON_COLORS,
    )
    top.set_title("by position in the gene", loc="left", fontsize=10, color=INK, pad=22)
    lengths = exons["exon_length_bins"]
    _stacked_rows(
        bottom,
        [
            f"{b.replace('+', '+ ')} bp (n={sum(c.values()):,})"
            for b, c in zip(lengths, by_length, strict=True)
        ],
        [_percent(c, classes) for c in by_length],
        EXON_COLORS,
    )
    bottom.set_title("by exon length", loc="left", fontsize=10, color=INK, pad=22)
    fig.tight_layout()
    return fig


def plot_failure_anatomy(anatomy: Mapping[str, Any]) -> Figure:
    """How many exons are wrong in each failed gene, and which one when it is one."""
    grid = anatomy["wrong_by_gene_size"]
    sizes = anatomy["exon_count_bins"]
    buckets = anatomy["wrong_exon_buckets"]
    fig, (left, right) = plt.subplots(1, 2, figsize=(11, 3.6))
    ramp = ["#c3d5fc", "#7ea4fc", "#3973fc", "#1d56d6", "#004697", "#002a5c"]
    shares = [
        {
            b: 100 * c / sum(row) if sum(row) else 0.0
            for b, c in zip(buckets, row, strict=True)
        }
        for row in grid
    ]
    rows = [
        f"{s} exon{'' if s == '1' else 's'} (n={sum(r):,})"
        for s, r in zip(sizes, grid, strict=True)
    ]
    _stacked_rows(left, rows, shares, dict(zip(buckets, ramp, strict=False)))
    left.set_title(
        "wrong exons per failed gene", loc="left", fontsize=10, color=INK, pad=22
    )

    observed = anatomy["single_wrong_observed"]
    expected = anatomy["single_wrong_expected"]
    positions = list(observed)
    x = range(len(positions))
    right.bar(
        [i - 0.2 for i in x],
        [observed[p] for p in positions],
        width=0.4,
        edgecolor="white",
        linewidth=2,
        color=NTV3_COLOR,
        label="observed",
    )
    right.bar(
        [i + 0.2 for i in x],
        [expected[p] for p in positions],
        width=0.4,
        edgecolor="white",
        linewidth=2,
        color=OTHER_COLOR,
        label="if uniform",
    )
    right.set_xticks(list(x), positions)
    right.set_ylabel("failed genes")
    right.set_title(
        "position of the single wrong exon", loc="left", fontsize=10, color=INK
    )
    right.legend(frameon=False, fontsize=8)
    _style(right)
    right.grid(axis="x", visible=False)
    fig.tight_layout()
    return fig


def plot_independence(check: Mapping[str, Any]) -> Figure:
    """Observed gene recovery vs. what independent exon errors would give."""
    bins = check["exon_count_bins"]
    n = check["n_genes"]
    observed = [
        100 * o / t if t else 0 for o, t in zip(check["observed"], n, strict=True)
    ]
    predicted = [
        100 * p / t if t else 0 for p, t in zip(check["predicted"], n, strict=True)
    ]
    fig, ax = plt.subplots(figsize=(6, 3.4))
    ax.plot(bins, observed, marker="o", color=NTV3_COLOR, linewidth=2, label="observed")
    ax.plot(
        bins,
        predicted,
        marker="o",
        color=OTHER_COLOR,
        linewidth=2,
        linestyle="--",
        label="independent exons",
    )
    ax.set_ylim(0, 100)
    ax.set_xlabel("coding exons per gene")
    ax.set_ylabel("genes recovered exactly (%)")
    ax.legend(frameon=False, fontsize=8)
    _style(ax)
    ax.grid(axis="y", color="#e5e7eb", linewidth=0.8)
    fig.tight_layout()
    return fig


def plot_boundary_offsets(offsets: Mapping[str, Any]) -> Figure:
    """How far a wrong exon boundary is from the reference one."""
    bins = offsets["offset_bins"]
    counts = offsets["offsets"]
    total = sum(counts)
    fig, ax = plt.subplots(figsize=(7, 3.2))
    shares = [100 * c / total if total else 0 for c in counts]
    ax.bar([f"{b} bp" for b in bins], shares, color=NTV3_COLOR, width=0.6)
    for x, value in enumerate(shares):
        ax.text(x, value + 0.5, f"{value:.0f}%", ha="center", fontsize=8)
    ax.set_ylabel("shifted boundaries (%)")
    ax.set_xlabel("distance to the reference boundary")
    _style(ax)
    ax.grid(axis="x", visible=False)
    fig.tight_layout()
    return fig


def plot_structure(chains: Mapping[str, Mapping[str, float]]) -> Figure:
    """Exact gene F1 vs. intron-chain F1 (outer ends forgiven), by gene type."""
    groups = [
        ("multi_exon", "multi-exon"),
        ("single_exon", "single-exon"),
        ("all", "all"),
    ]
    fig, ax = plt.subplots(figsize=(6.5, 3.2))
    for i, (key, _) in enumerate(groups):
        for offset, (metric, color, label) in zip(
            (-0.2, 0.2),
            (
                ("exact_f1", NTV3_COLOR, "exact structure"),
                ("chain_f1", "#004697", "intron chain"),
            ),
            strict=True,
        ):
            value = chains[key][metric]
            ax.bar(
                i + offset,
                value,
                width=0.4,
                edgecolor="white",
                linewidth=2,
                color=color,
                label=label if i == 0 else None,
            )
            ax.text(i + offset, value + 1, f"{value:.1f}", ha="center", fontsize=8)
    ax.set_xticks(range(len(groups)), [label for _, label in groups])
    ax.set_ylim(0, 110)
    ax.set_ylabel("gene F1 (%)")
    ax.legend(frameon=False, fontsize=8, ncol=2, loc="upper left")
    _style(ax)
    ax.grid(axis="x", visible=False)
    fig.tight_layout()
    return fig


def plot_proteins(counts: Mapping[str, Any]) -> Figure:
    """Identity of each reference protein's best predicted counterpart."""
    histogram = counts["identity_histogram"]
    total = sum(histogram.values())
    fig, ax = plt.subplots(figsize=(6.5, 3.2))
    labels = list(histogram)
    shares = [100 * histogram[k] / total if total else 0 for k in labels]
    ax.bar([f"{k}%" for k in labels], shares, color=NTV3_COLOR, width=0.6)
    for x, value in enumerate(shares):
        ax.text(x, value + 1, f"{value:.0f}%", ha="center", fontsize=8)
    ax.set_ylabel("reference proteins (%)")
    ax.set_xlabel("identity to the best overlapping prediction")
    _style(ax)
    ax.grid(axis="x", visible=False)
    fig.tight_layout()
    return fig


def plot_failed_gene_repair(
    exons: Mapping[str, Any], anatomy: Mapping[str, Any], summary: Mapping[str, Any]
) -> Figure:
    """Wrong exons per failed gene, and the genes a fix of at most k exons would add."""
    buckets = exons["wrong_exon_buckets"]
    per_gene = exons["wrong_exons_per_failed_gene"]
    fig, (left, right) = plt.subplots(1, 2, figsize=(11, 3.4))
    failed = sum(per_gene)
    shares = [100 * n / failed if failed else 0 for n in per_gene]
    left.bar(buckets, shares, color=NTV3_COLOR, width=0.6)
    for x, value in enumerate(shares):
        left.text(x, value + 1, f"{value:.0f}%", ha="center", fontsize=8)
    left.set_xlabel("wrong exons in the gene")
    left.set_ylabel("failed genes (%)")
    left.set_title("how wrong a failed gene is", loc="left", fontsize=10, color=INK)
    n_loci = summary["n_loci"]
    recall = 100 * summary["correct"] / n_loci
    gained = [100 * c / n_loci for c in anatomy["cumulative_convertible"]]
    right.bar(
        buckets,
        [recall] * len(buckets),
        color=OTHER_COLOR,
        width=0.6,
        label="recovered",
    )
    right.bar(
        buckets,
        gained,
        bottom=recall,
        color=NTV3_COLOR,
        width=0.6,
        label="after the fix",
    )
    for x, value in enumerate(gained):
        right.text(
            x, recall + value + 1, f"{recall + value:.0f}%", ha="center", fontsize=8
        )
    right.set_ylim(0, 105)
    right.set_xlabel("fix at most this many exons per failed gene")
    right.set_ylabel("reference genes exact (%)")
    right.set_title(
        "what fixing a few exons would buy", loc="left", fontsize=10, color=INK
    )
    right.legend(frameon=False, fontsize=8, loc="lower right")
    for ax in (left, right):
        _style(ax)
        ax.grid(axis="x", visible=False)
    fig.tight_layout()
    return fig


def plot_structure_and_merges(
    chains: Mapping[str, Mapping[str, float]], merges: Mapping[str, Any]
) -> Figure:
    """Exact vs. intron-chain gene F1, and how often one prediction merges loci."""
    fig, (left, right) = plt.subplots(1, 2, figsize=(10, 3.2))
    values = [chains["all"]["exact_f1"], chains["all"]["chain_f1"]]
    left.bar(
        ["exact structure", "intron chain"],
        values,
        color=[NTV3_COLOR, "#004697"],
        width=0.5,
    )
    for x, value in enumerate(values):
        left.text(x, value + 1, f"{value:.1f}", ha="center", fontsize=9)
    left.set_ylim(0, 105)
    left.set_ylabel("gene F1 (%)")
    shares = [
        100 * merges["absorbed_loci"] / merges["n_ref"],
        100 * merges["merging_predictions"] / merges["n_pred"],
    ]
    right.barh(
        ["reference genes\nabsorbed in a merge", "predictions that\nmerge 2+ genes"],
        shares,
        color=CLASS_COLORS["fragmented"],
        height=0.5,
    )
    for y, value in enumerate(shares):
        right.text(value, y, f" {value:.1f}%", va="center", fontsize=9)
    right.set_xlim(0, max(5, 1.3 * max(shares)))
    right.set_xlabel("%")
    right.invert_yaxis()
    for ax in (left, right):
        _style(ax)
    left.grid(axis="x", visible=False)
    fig.tight_layout()
    return fig


def plot_protein_recovery(counts: Mapping[str, Any]) -> Figure:
    """Reference proteins recovered at each identity, and well-formed predicted ORFs."""
    fig, (left, right) = plt.subplots(
        1, 2, figsize=(10, 3.2), gridspec_kw={"width_ratios": (3, 1)}
    )
    n_ref = counts["n_reference"]
    thresholds = list(counts["recovered"])
    shares = [100 * counts["recovered"][t] / n_ref for t in thresholds]
    left.bar([f">= {t}%" for t in thresholds], shares, color=NTV3_COLOR, width=0.6)
    for x, value in enumerate(shares):
        left.text(x, value + 1, f"{value:.1f}%", ha="center", fontsize=8)
    left.set_ylim(0, 105)
    left.set_xlabel("identity to the best overlapping prediction")
    left.set_ylabel("reference proteins (%)")
    well = 100 * counts["well_formed"] / counts["n_predicted"]
    right.bar(["well-formed\nORFs"], [well], color=NTV3_COLOR, width=0.5)
    right.text(0, well + 1, f"{well:.2f}%", ha="center", fontsize=8)
    right.set_ylim(0, 105)
    right.set_ylabel("predicted proteins (%)")
    for ax in (left, right):
        _style(ax)
        ax.grid(axis="x", visible=False)
    fig.tight_layout()
    return fig


def plot_locus(
    reference: Gene,
    predictions: Sequence[Gene],
    title: str = "",
    pad_fraction: float = 0.04,
) -> Figure:
    """Gene models of one reference locus and the predictions overlapping it."""
    genes = [reference, *predictions]
    spans = [cds_span(g) for g in genes if g.cds_keys]
    start, end = min(s for s, _ in spans), max(e for _, e in spans)
    pad = max(1, int((end - start) * pad_fraction))
    fig, ax = plt.subplots(figsize=(10, 0.8 + 0.5 * len(genes)))
    ref_boundaries = {k[:4] for k in reference.cds_keys}
    for y, gene in enumerate(genes):
        lo, hi = cds_span(gene)
        ax.plot([lo, hi], [y, y], color=MUTED, linewidth=1)
        for key in gene.cds_keys:
            good = y == 0 or key[:4] in ref_boundaries
            color = (
                (OTHER_COLOR if y == 0 else NTV3_COLOR)
                if good
                else CLASS_COLORS["fragmented"]
            )
            ax.add_patch(
                Rectangle((key[1], y - 0.25), key[2] - key[1] + 1, 0.5, color=color)
            )
    ax.set_yticks(range(len(genes)), ["reference", *[g.gene_id for g in predictions]])
    ax.invert_yaxis()
    ax.set_xlim(start - pad, end + pad)
    ax.xaxis.set_major_formatter(lambda x, _: f"{x:,.0f}")
    strand = "forward" if reference.strand == "+" else "reverse"
    ax.set_xlabel(f"{reference.chrom} ({strand} strand), bp")
    ax.set_title(title or reference.gene_id, loc="left", fontsize=10, color=INK)
    _style(ax)
    fig.tight_layout()
    return fig


# --- Decoding walkthrough (03b) ---------------------------------------------------

_INTRON = np.array([c.startswith("intron") for c in decoding.CLASSES])
_EXONIC = ~_INTRON & (np.arange(decoding.NUM_CLASSES) != 0)
GROUP_COLORS = {"intergenic": "#e5e7eb", "exon": NTV3_COLOR, "intron": "#c3d5fc"}


def plot_transitions() -> Figure:
    """The HMM's allowed transitions and their probabilities (log scale)."""
    probs = np.exp(decoding.LOG_TRANSITIONS)
    fig, ax = plt.subplots(figsize=(7, 6))
    shown = np.where(probs > 0, np.log10(np.maximum(probs, 1e-12)), np.nan)
    image = ax.imshow(shown, cmap="Blues", vmin=-5, vmax=0)
    ax.set_xticks(
        range(decoding.NUM_CLASSES), decoding.CLASSES, rotation=90, fontsize=8
    )
    ax.set_yticks(range(decoding.NUM_CLASSES), decoding.CLASSES, fontsize=8)
    ax.set_xlabel("to")
    ax.set_ylabel("from")
    fig.colorbar(image, ax=ax, label="log10 probability (blank: forbidden)")
    fig.tight_layout()
    return fig


def _feature_runs(
    states: np.ndarray,
) -> tuple[list[tuple[int, int]], list[tuple[int, int]]]:
    """(exon runs, intron runs) of a state path, as half-open index ranges."""
    kind = np.where(_INTRON[states], 2, np.where(_EXONIC[states], 1, 0))
    edges = np.flatnonzero(np.diff(np.concatenate([[-1], kind, [-1]])))
    runs = list(itertools.pairwise(edges))
    return (
        [(a, b) for a, b in runs if kind[a] == 1],
        [(a, b) for a, b in runs if kind[a] == 2],
    )


def plot_tracks(
    tracks: Mapping[str, np.ndarray],
    region_start: int,
    zoom: tuple[int, int],
    reference: Sequence[Gene] = (),
    title: str = "",
    mark_errors: bool = False,
) -> Figure:
    """Gene structure of each state path over ``zoom`` (contig bp), reference on top.

    Args:
        tracks: Label to one strand's state path over the region, genomic order.
        region_start: Contig bp of the paths' first position.
        zoom: Contig bp range to draw.
        reference: Reference genes of that strand to draw on the first row.
        title: Axes title.
        mark_errors: Draw predicted exons whose boundaries match no reference exon
            in red.
    """
    ref_exons = {(k[1], k[2]) for gene in reference for k in gene.cds_keys}
    rows = ["reference", *tracks] if reference else list(tracks)
    fig, ax = plt.subplots(figsize=(11, 0.9 + 0.45 * len(rows)))
    lo, hi = zoom
    y = 0
    if reference:
        for gene in reference:
            s, e = cds_span(gene)
            if e < lo or s > hi:
                continue
            ax.plot([s, e], [y, y], color=MUTED, linewidth=1)
            for key in gene.cds_keys:
                ax.add_patch(
                    Rectangle(
                        (key[1], y - 0.25), key[2] - key[1] + 1, 0.5, color=OTHER_COLOR
                    )
                )
        y += 1
    for states in tracks.values():
        window = states[lo - region_start : hi - region_start]
        exons, introns = _feature_runs(window)
        for a, b in introns:
            ax.plot([lo + a, lo + b], [y, y], color=MUTED, linewidth=1)
        for a, b in exons:
            wrong = mark_errors and (lo + a + 1, lo + b) not in ref_exons
            color = CLASS_COLORS["fragmented"] if wrong else NTV3_COLOR
            ax.add_patch(Rectangle((lo + a, y - 0.25), b - a, 0.5, color=color))
        y += 1
    ax.set_yticks(range(len(rows)), rows)
    ax.set_ylim(len(rows) - 0.5, -0.5)
    ax.set_xlim(lo, hi)
    ax.xaxis.set_major_formatter(lambda x, _: f"{x:,.0f}")
    ax.set_title(title, loc="left", fontsize=10, color=INK)
    _style(ax)
    fig.tight_layout()
    return fig


def plot_class_probabilities(
    log_probs: np.ndarray, region_start: int, zoom: tuple[int, int], title: str = ""
) -> Figure:
    """Probability of intergenic / exon / intron along ``zoom``, stacked to 1."""
    lo, hi = zoom
    probs = np.exp(log_probs[lo - region_start : hi - region_start])
    groups = {
        "intergenic": probs[:, 0],
        "exon": probs[:, _EXONIC].sum(1),
        "intron": probs[:, _INTRON].sum(1),
    }
    fig, ax = plt.subplots(figsize=(11, 2.4))
    x = np.arange(lo, hi)
    ax.stackplot(
        x, *groups.values(), colors=list(GROUP_COLORS.values()), labels=list(groups)
    )
    ax.set_xlim(lo, hi)
    ax.set_ylim(0, 1)
    ax.set_ylabel("probability")
    ax.xaxis.set_major_formatter(lambda v, _: f"{v:,.0f}")
    ax.legend(
        frameon=False, fontsize=8, ncol=3, loc="lower left", bbox_to_anchor=(0, 1)
    )
    ax.set_title(title, loc="right", fontsize=10, color=INK)
    _style(ax)
    fig.tight_layout()
    return fig


def plot_edge_drift(
    distance_kb: np.ndarray, drift: np.ndarray, crop_kb: float
) -> Figure:
    """How far predictions move from their full-context values near a window edge.

    Args:
        distance_kb: Bin centres, kb from the window edge.
        drift: Mean total-variation distance to the full-context prediction.
        crop_kb: Context margin the pipeline discards, drawn as a line.
    """
    fig, ax = plt.subplots(figsize=(7.5, 3.2))
    ax.axvspan(
        distance_kb.min() / 2, crop_kb, color="#eef2ff", label="discarded margin"
    )
    ax.plot(distance_kb, drift, color=NTV3_COLOR, linewidth=2, marker="o", markersize=5)
    ax.set_xscale("log")
    ax.set_xlim(distance_kb.min() / 2, distance_kb.max() * 1.5)
    ticks = [1, 4, 16, 64, round(crop_kb)]
    ax.set_xticks(ticks, [f"{t:g}" for t in ticks])
    ax.minorticks_off()
    ax.set_ylim(0, None)
    ax.set_xlabel("distance to the window edge (kb)")
    ax.set_ylabel("change in predicted\nprobabilities (TV)")
    ax.legend(frameon=False, fontsize=8)
    _style(ax)
    ax.grid(axis="y", color="#e5e7eb", linewidth=0.8)
    fig.tight_layout()
    return fig