sras-viewer/sras_compute.py

#!/usr/bin/env python3
"""Image computation and angle alignment for .sras scans.

Depends only on numpy/scipy/pyfftw and sras_format, so a multiprocessing
child can import it without loading Qt or matplotlib — which matters
because Python 3.14 on macOS spawns rather than forks.
"""

import atexit
import json
import os
import threading
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pyfftw
import scipy.fft as scipy_fft
import scipy.ndimage as scipy_ndimage

from sras_format import (CH1_IDX, CH3_IDX, CH4_IDX, MAX_PAD_FACTOR, SrasFile,
                         adc_to_mv)

# ---------------------------------------------------------------------------
# FFT worker pool and per-thread pyFFTW plans
# ---------------------------------------------------------------------------

_FFT_BLOCK_MAX = 512   # ceiling on waveforms/task: the knee measured on a
                       # 16-core machine at natural resolution (n_len == spf).
                       # Smaller blocks serialise on GIL-held numpy dispatch,
                       # larger ones lose cache residency and task granularity.
_FFT_BLOCK_MIN = 32    # floor, so task granularity never collapses at
                       # extreme pad factors — at the cost of exceeding
                       # _FFT_PLAN_BYTES_BUDGET there (see _fft_block_for).
_FFT_PLAN_BYTES_BUDGET = int(
    os.environ.get("SRAS_FFT_PLAN_BUDGET_MB", 16)) * 1024 * 1024
    # Per-thread ceiling on one cached pyFFTW plan's resident input+output
    # buffers (see _fft_block_for). This is PERMANENT memory — the plan
    # cache is never evicted — so the worst case across the whole pool for
    # one distinct (spf, n_len) is _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET.
    # Deliberately separate from _TOTAL_BYTES_BUDGET/SRAS_MEM_BUDGET_MB,
    # which bounds transient concurrently-live chunk buffers, freed at the
    # end of each chunk — this one bounds a permanent per-thread tax that
    # compounds with worker count, not chunk concurrency.

_pool_lock = threading.Lock()
_pool: ThreadPoolExecutor | None = None


def _fft_pool() -> ThreadPoolExecutor:
    """The persistent process-wide pool for FFT block tasks."""
    global _pool
    with _pool_lock:
        if _pool is None:
            _pool = ThreadPoolExecutor(max_workers=_MAX_WORKERS,
                                       thread_name_prefix="sras-fft")
            atexit.register(_pool.shutdown, wait=False, cancel_futures=True)
        return _pool


_WISDOM_PATH = Path.home() / ".cache" / "sras-viewer" / "fftw_wisdom"
_wisdom_lock = threading.Lock()
_wisdom_loaded = False
_fftw_local = threading.local()


def _load_wisdom_once():
    """Import saved FFTW wisdom so FFTW_MEASURE planning is a one-time cost
    per machine. Purely an optimisation: failures are ignored."""
    global _wisdom_loaded
    with _wisdom_lock:
        if _wisdom_loaded:
            return
        _wisdom_loaded = True
        try:
            pyfftw.import_wisdom(_WISDOM_PATH.read_bytes().split(b"\x00\n"))
        except Exception:
            pass


def _save_wisdom():
    with _wisdom_lock:
        try:
            _WISDOM_PATH.parent.mkdir(parents=True, exist_ok=True)
            _WISDOM_PATH.write_bytes(b"\x00\n".join(pyfftw.export_wisdom()))
        except Exception:
            pass


def _fft_block_for(spf: int, n_len: int) -> int:
    """Waveforms per FFT task/plan at transform length n_len. Bounds a
    cached pyFFTW plan's resident input+output buffers (permanent, one per
    pool thread, never evicted) to _FFT_PLAN_BYTES_BUDGET regardless of pad
    factor. Reproduces _FFT_BLOCK_MAX exactly at natural resolution
    (n_len == spf) — see docs/design.md for the worked numbers at high pad."""
    bytes_per_wf = 4 * spf + 8 * (n_len // 2 + 1)
    block = _FFT_PLAN_BYTES_BUDGET // max(1, bytes_per_wf)
    return int(min(_FFT_BLOCK_MAX, max(_FFT_BLOCK_MIN, block)))


def _block_rfft(waves: np.ndarray, n: int) -> np.ndarray:
    """rfft of a (B, spf) float32 block via a cached per-thread FFTW plan.

    Plans have a fixed (block, spf) input shape, block = _fft_block_for(spf,
    n), so each worker thread plans once per (block size, transform length)
    pair; a remainder block runs through the same plan with its tail rows
    ignored. Single-threaded (threads=1): outer parallelism comes from the
    pool, so the transform itself must not spin up threads. The returned
    array is the plan's output buffer — consume it before the next call on
    the same thread.
    """
    n_wf, spf = waves.shape
    block = _fft_block_for(spf, n)
    plans = getattr(_fftw_local, "plans", None)
    if plans is None:
        plans = _fftw_local.plans = {}
    key = (block, spf, n)
    plan = plans.get(key)
    if plan is None:
        _load_wisdom_once()
        buf = pyfftw.empty_aligned((block, spf), dtype="float32")
        # No overwrite_input: FFTW must not scribble on input_array, whose
        # zero-padded tail (columns spf..n) is zeroed exactly once here.
        plan = pyfftw.builders.rfft(buf, n=n, axis=-1, threads=1,
                                    planner_effort="FFTW_MEASURE")
        plan.input_array[:] = 0.0
        plans[key] = plan
        _save_wisdom()
    inp = plan.input_array
    inp[:n_wf, :spf] = waves
    return plan()[:n_wf]


# ---------------------------------------------------------------------------
# Peak search: full transform, block-parallel
# ---------------------------------------------------------------------------

def _peak_bins(waves: np.ndarray, n_len: int, min_bin: int = 1) -> np.ndarray:
    """Peak bin per waveform via the full transform. *min_bin* zeroes every
    bin below it (always >= 1, so true DC stays suppressed) before the
    argmax, so the peak search never returns a bin below the caller's
    frequency floor. Called in _fft_block_for(spf, n_len)-sized blocks,
    fanned out across _fft_pool() — see docs/design.md."""
    S = _block_rfft(waves, n_len)
    power = S.real ** 2
    power += S.imag ** 2
    power[:, :min_bin] = 0.0
    return np.argmax(power, axis=1)


# ---------------------------------------------------------------------------
# Chunking / parallel budget
# ---------------------------------------------------------------------------
#
# Row chunks are budgeted so all concurrently-live working buffers fit in
# memory; the worker count is derived from the budget, not vice versa.
# Rationale and the measured 1024 MB default: docs/design.md ("Memory budget
# and row chunking").
_TOTAL_BYTES_BUDGET = int(os.environ.get("SRAS_MEM_BUDGET_MB", 1024)) * 1024 * 1024
_CHUNK_ROWS_MAX = 32           # cap for small scans (original behavior)
_MAX_WORKERS = int(os.environ.get("SRAS_MAX_WORKERS", 0)) or (os.cpu_count() or 4)

def memory_budget_bytes() -> int:
    """The module-wide ceiling on concurrently-live working buffers, for
    callers outside this module that read the same data (the aligned exporter
    sizes its source-band reader against it)."""
    return _TOTAL_BYTES_BUDGET


def _chunk_rows_for(n_frames: int, samples_per_frame: int,
                    budget: int = _TOTAL_BYTES_BUDGET) -> int:
    bytes_per_row = max(1, n_frames * samples_per_frame * 4)   # float32
    return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))


def _plan_fft_rows(n_frames: int, samples_per_frame: int, budget: int) -> int:
    """Rows per outer chunk for the block-FFT path. Only the float32
    waveform read buffer scales with the chunk; spectrum memory is bounded
    independently, by _fft_block_for, to _MAX_WORKERS * _FFT_PLAN_BYTES_BUDGET
    regardless of chunk size or pad factor (see docs/design.md), so this
    formula budgets the read buffer alone, with 2x slack, and lets the block
    fan-out saturate the pool."""
    bytes_per_row = max(1, n_frames * samples_per_frame * 4 * 2)
    return int(max(1, min(_CHUNK_ROWS_MAX, budget // bytes_per_row)))


def _plan_chunks(n_rows: int, n_frames: int, samples_per_frame: int,
                 live_multiplier: int = 1,
                 max_workers: int | None = None,
                 budget: int | None = None) -> tuple[int, int]:
    """(chunk_rows, n_workers) such that all concurrent chunks together fit
    the budget: n_workers x live_multiplier x chunk_rows x n_frames x
    samples x 4 bytes <= budget.

    The worker count is chosen *first* and the chunk sized to it. Sizing the
    chunk first is the trap: _chunk_rows_for spends whatever budget it is
    given, so a single chunk would always consume the lot and leave room for
    exactly one worker — no concurrency, on precisely the large scans that
    need it most.

    A caller that is itself running several of these concurrently must pass
    *both* max_workers=1 and its share of the budget. Capping the workers
    alone is not enough: the chunk would still be sized against the whole
    budget, and N concurrent callers would each allocate all of it.
    """
    total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
    cap = max_workers if max_workers is not None else _MAX_WORKERS
    bytes_per_row = max(1, live_multiplier * n_frames * samples_per_frame * 4)

    # A chunk is at least one row, so that alone caps how many can be live.
    n_workers = int(max(1, min(cap, n_rows, total // bytes_per_row)))
    chunk_rows = _chunk_rows_for(n_frames, samples_per_frame,
                                 max(1, total // (n_workers * live_multiplier)))
    # With chunk_rows known, more workers than chunks buys nothing.
    n_workers = int(max(1, min(n_workers, -(-n_rows // chunk_rows))))
    return chunk_rows, n_workers


def _map_row_chunks(n_rows: int, chunk_rows: int, n_workers: int, fn,
                    should_stop=None):
    """Apply fn(r0, r1) over row chunks, in parallel when it pays.

    Chunks write to disjoint output slices, so no locking is needed. numpy
    ufuncs, scipy's pocketfft, and memmap page faults all release the GIL, so
    threads give real parallelism here — and on a very large file they also
    keep many more page-fault requests in flight, which is what the I/O path
    wants.

    *should_stop* is polled per chunk so a cancelled job abandons work within
    one chunk rather than one whole angle — on a large scan an angle is
    ~40 s, which is far too long to block application shutdown.
    """
    bounds = [(r0, min(r0 + chunk_rows, n_rows))
              for r0 in range(0, n_rows, chunk_rows)]

    def run(b):
        if should_stop is not None and should_stop():
            return
        fn(*b)

    if n_workers <= 1 or len(bounds) == 1:
        for b in bounds:
            run(b)
        return
    with ThreadPoolExecutor(max_workers=min(n_workers, len(bounds))) as pool:
        list(pool.map(run, bounds))


# ---------------------------------------------------------------------------
# Image computation
# ---------------------------------------------------------------------------

def plan_angle_level(sras: SrasFile,
                     live_multiplier: int = 1) -> tuple[int, int]:
    """(n_angle_workers, per_angle_budget) for a caller that parallelises
    over angles instead of over rows.

    Callers that parallelise over angles must not also let each angle
    parallelise over rows: the two levels multiply, in threads and in memory.
    Each angle is therefore given max_workers=1 and the returned budget share,
    which together keep total live buffers within _TOTAL_BYTES_BUDGET.
    """
    biggest = max(range(sras.n_angles), key=lambda a: int(sras.n_frames[a]))
    _, n_workers = _plan_chunks(int(sras.n_rows[biggest]),
                                int(sras.n_frames[biggest]),
                                sras.samples_per_frame,
                                live_multiplier=live_multiplier)
    n_workers = max(1, min(n_workers, sras.n_angles))
    return n_workers, max(1, _TOTAL_BYTES_BUDGET // n_workers)


def compute_dc_image(sras: SrasFile, angle_idx: int, ch_idx: int,
                     max_workers: int | None = None,
                     budget: int | None = None,
                     should_stop=None) -> np.ndarray:
    """Mean of each waveform → (n_rows, n_frames) float32, in ADC counts.

    Pass max_workers=1 when the caller is already parallelising over angles.
    If *should_stop* ever returns True the result is incomplete — the caller
    is expected to be abandoning it.
    """
    n_rows, n_frames = sras.image_shape(angle_idx)
    data = sras.data[angle_idx]
    chunk_rows, n_workers = _plan_chunks(n_rows, n_frames, sras.samples_per_frame,
                                         max_workers=max_workers, budget=budget)
    img = np.empty((n_rows, n_frames), dtype=np.float32)

    def chunk(r0: int, r1: int):
        img[r0:r1] = data[r0:r1, ch_idx, :, :].astype(np.float32).mean(axis=-1)

    _map_row_chunks(n_rows, chunk_rows, n_workers, chunk, should_stop=should_stop)
    return img


def dc_image_mv(sras: SrasFile, angle_idx: int, ch_idx: int,
                max_workers: int | None = None, budget: int | None = None,
                should_stop=None) -> np.ndarray:
    """DC image for (angle, channel) in mV, preferring a stored v5/v7 cache."""
    cached = sras.cached_dc_mv(angle_idx, ch_idx)
    if cached is not None:
        return cached
    return adc_to_mv(
        compute_dc_image(sras, angle_idx, ch_idx, max_workers=max_workers,
                         budget=budget, should_stop=should_stop),
        *sras.cal(ch_idx))



# ---------------------------------------------------------------------------
# Row-averaged FFT: same-row, distance-weighted CH1 waveform smoothing
#
# Averages a pixel's CH1 waveform with its same-row neighbors *before* the
# FFT peak search, to improve SNR on noisy scans. Never crosses rows: pixel
# pitch is strongly anisotropic and varies by scan, but the X pitch within
# one row is a single file-wide constant (SrasFile.pixel_x_mm), so a
# Gaussian in pixel-index distance along a row and one in true physical mm
# distance are the same function up to that constant scale factor — the
# kernel itself needs no pitch, only the GUI's physical-width hint label
# does. Rationale and the masked-average/background-subtraction proofs:
# docs/design.md ("Row-averaged FFT").
# ---------------------------------------------------------------------------

_ROW_AVG_SIGMA_FRAC = 0.5   # sigma = n * this; edge weight (distance n) is
                            # exp(-1/(2*frac**2)) ~= 0.135 of the center tap


def _row_average_weights(row_avg_n: int) -> np.ndarray:
    """(2n+1,) float32 Gaussian weights for distance-weighted row averaging,
    symmetric around the center tap. Not pre-normalized to sum to 1 —
    _row_average_waveforms renormalizes per output pixel by the actual sum
    of included, valid, in-window neighbor weights, not a fixed total."""
    n = max(0, int(row_avg_n))
    if n == 0:
        return np.ones(1, dtype=np.float32)
    d = np.arange(-n, n + 1, dtype=np.float64)
    sigma = n * _ROW_AVG_SIGMA_FRAC
    return np.exp(-0.5 * (d / sigma) ** 2).astype(np.float32)


def _row_average_waveforms(masked_waves: np.ndarray, valid: np.ndarray,
                           weights: np.ndarray) -> np.ndarray:
    """Distance-weighted mean of each row position's CH1 waveform with its
    same-row neighbors, counting only neighbors where *valid* is True.

    masked_waves: (n_frames, spf) float32 — CH1 samples at valid[f]
        positions; must already be 0.0 elsewhere (the caller must never
        have read the raw memmap at an invalid position).
    valid: (n_frames,) bool.
    weights: (2n+1,) float32 from _row_average_weights.

    Returns (n_frames, spf) float32, meaningful only where valid is True
    (the caller compacts by that same mask right after, matching
    compute_rf_image's existing masked-compaction contract).

    Two 1-D correlations along the frame axis — the numerator against the
    raw masked-zero waveforms, the denominator against the validity mask
    itself — so a masked neighbor contributes *zero weight* rather than a
    zero-amplitude sample at full weight, and a window truncated at a row's
    edge renormalizes correctly with no separate edge case: mode="constant",
    cval=0.0 pads both convolutions with zero beyond the row's own ends.
    """
    num = scipy_ndimage.correlate1d(masked_waves, weights, axis=0,
                                    mode="constant", cval=0.0)
    den = scipy_ndimage.correlate1d(valid.astype(np.float32), weights, axis=0,
                                    mode="constant", cval=0.0)
    den_safe = np.where(valid, den, np.float32(1.0))
    return num / den_safe[:, None]


def pad_factor_for(sras: SrasFile, n_fft: int | None) -> int:
    """The integer zero-padding factor an *n_fft* request represents, or 0
    if it represents none — i.e. it is not a whole multiple of this file's
    samples_per_frame, so no stored image (which only ever records an
    integer factor) can answer it.

    0 rather than None so callers can compare it straight against
    ``sras.precomputed_pad_factor``, which is never 0.
    """
    if n_fft is None:
        return 1
    spf = sras.samples_per_frame
    if spf <= 0 or n_fft % spf:
        return 0
    return max(1, n_fft // spf)


def cache_mismatch_reasons(sras: SrasFile, *, n_fft: int | None,
                           apply_bg_sub: bool, row_avg_n: int) -> list[str]:
    """Why this file's stored FFT images can't answer a request, as human-
    readable phrases; empty means they can.

    Padding, background subtraction and row-averaging are all baked into
    the stored numbers, so a request that differs in any of them has to go
    back through a real FFT. This is the single accept rule: cached_rf_image
    serves a stored image iff this returns nothing, and callers that want to
    *explain* the miss (rather than silently recompute) format these same
    strings, so the two can't drift apart.
    """
    reasons = []
    pad = pad_factor_for(sras, n_fft)
    if pad != sras.precomputed_pad_factor:
        want = f"pad {pad}x" if pad else "a ragged n_fft"
        reasons.append(f"stored at pad {sras.precomputed_pad_factor}x, "
                       f"requested at {want}")
    if sras.precomputed_bg_sub != (apply_bg_sub and sras.background is not None):
        reasons.append("stored with background subtraction "
                       f"{'on' if sras.precomputed_bg_sub else 'off'}")
    if sras.precomputed_row_avg_n != row_avg_n:
        have = (f"row-averaged (n={sras.precomputed_row_avg_n})"
                if sras.precomputed_row_avg_n else "raw per-pixel")
        want = (f"row-averaged (n={row_avg_n})"
                if row_avg_n else "raw per-pixel")
        reasons.append(f"stored {have}, requested {want}")
    return reasons


def cached_rf_image(sras: SrasFile, angle_idx: int,
                    dc_threshold_mv: float | None,
                    apply_bg_sub: bool = True,
                    n_fft: int | None = None,
                    dc4_mv: np.ndarray | None = None,
                    row_avg_n: int = 0,
                    allow_dc_recompute: bool = True) -> np.ndarray | None:
    """The precomputed-cache fast path for compute_rf_image: a ready-to-
    display peak-frequency image if this file already has one matching
    every setting the caller cares about, else None (caller must run a
    real FFT).

    A stored image stands in only when *all* hold: this angle actually has
    a stored image (v5 PREC or v7 CACH); the requested zero-padding matches
    what the store was computed at (see below); the stored bg-sub flag
    matches what the caller wants; and sras.precomputed_row_avg_n ==
    row_avg_n exactly (0 == "raw") — this last check is what stops a raw
    request from ever being silently served a row-averaged image, or vice
    versa, or a request at one window size being served a cache stored at a
    different one.

    Padding is checked the same way, and for the same reason: a padded FFT
    interpolates between the natural bins, so it resolves genuinely
    different peak frequencies. *n_fft* of None means natural resolution,
    i.e. pad 1; anything else must be an exact whole multiple of
    samples_per_frame equal to sras.precomputed_pad_factor. A ragged n_fft
    that is not such a multiple can never match a stored image, since the
    store only ever records an integer pad factor.

    If *allow_dc_recompute* is False and no DC4 image is already cached or
    supplied via *dc4_mv*, applying the mask would mean reading a whole
    channel on the caller's behalf; this returns None instead so a caller
    that wants to stay off the I/O path (e.g. a GUI thread) can choose to
    fall through to a real compute rather than block.
    """
    cached_freq = sras.precomputed_freq_mhz[angle_idx]
    if cached_freq is None or cache_mismatch_reasons(
            sras, n_fft=n_fft, apply_bg_sub=apply_bg_sub, row_avg_n=row_avg_n):
        return None
    dc4_img = None
    if dc_threshold_mv is not None:
        # DC4 mask, in priority order: already-cached DC block, caller-
        # supplied image, or a fresh (cheap — no FFT) recompute. Resolved
        # before the copy below so the allow_dc_recompute bail-out doesn't
        # allocate a full image it is about to throw away.
        dc4_img = sras.cached_dc_mv(angle_idx, CH4_IDX)
        if dc4_img is None:
            if dc4_mv is not None:
                dc4_img = dc4_mv
            elif allow_dc_recompute:
                dc4_img = adc_to_mv(compute_dc_image(sras, angle_idx, CH4_IDX),
                                    *sras.cal(CH4_IDX))
            else:
                return None
    freq_img = cached_freq.copy()
    if dc4_img is not None:
        freq_img[dc4_img < dc_threshold_mv] = 0.0
    return freq_img


def compute_rf_image(sras: SrasFile, angle_idx: int,
                     dc_threshold_mv: float | None,
                     apply_bg_sub: bool = True,
                     n_fft: int | None = None,
                     dc4_mv: np.ndarray | None = None,
                     max_workers: int | None = None,
                     budget: int | None = None,
                     should_stop=None,
                     row_avg_n: int = 0,
                     min_freq_mhz: float = 0.0) -> np.ndarray:
    """FFT of each CH1 waveform; pixel = peak frequency in MHz.

    Pixels where CH4_dc < dc_threshold_mv are set to 0 — and the FFT is
    never run for them, since that's the expensive part. The masked-out
    CH1 samples are also never *read*: the boolean mask is applied to the
    raw memmap slice before any dtype conversion, so numpy only pages in
    the bytes for pixels that pass the threshold (an untouched memmap page
    is never read from disk).

    *dc_threshold_mv* of None means "mask nothing", and skips reading and
    averaging CH4 entirely — worth a third of the I/O when caching a whole
    file, where the mask is applied later at display time.

    If the DC4 image for this angle is already known (e.g. from the
    DC-channel precompute cache), pass it as *dc4_mv* (mV, shape
    (n_rows, n_frames)) to reuse it instead of re-reading CH4 here.

    The peak search runs the full transform (_peak_bins) in
    _fft_block_for(spf, n_len)-sized blocks, fanned out across the
    pyFFTW-plan-caching thread pool — see docs/design.md for how the block
    size is bounded so this stays memory-safe at high pad factors.

    *row_avg_n* > 0 averages each pixel's CH1 waveform with its up-to-n
    same-row neighbors (distance-weighted, valid-neighbors-only per the
    same dc_threshold_mv mask) before the FFT runs — see
    _row_average_waveforms. 0 (default) is the raw, unaveraged behavior.

    *min_freq_mhz* excludes every bin below it (true DC, bin 0, is always
    excluded regardless) from the peak search, on top of — not instead of —
    the dc_threshold_mv mask. Without it, a pixel that passes the DC-bias
    threshold but carries only weak real signal can still resolve to a
    near-zero frequency: the un-subtracted background's DC-leakage skirt
    then has more power than the genuine (but weak) signal peak. Raising
    the floor above that skirt forces the search to report the strongest
    peak that is plausibly real signal instead. 0.0 (default) disables it
    (only true DC is excluded, the long-standing behavior).

    Fast path: if the file has a precomputed peak-frequency image for this
    angle (v5 PREC or v7 CACH) matching every one of the caller's settings
    — including row_avg_n exactly — the stored image is used directly, no
    FFT is run. See cached_rf_image. *min_freq_mhz* plays no part in that
    match — a stored image was baked without any floor, so it is returned
    as-is; the floor only ever affects a real (re)compute.
    """
    n_rows, n_frames = sras.image_shape(angle_idx)
    data = sras.data[angle_idx]

    fast = cached_rf_image(sras, angle_idx, dc_threshold_mv, apply_bg_sub=apply_bg_sub,
                           n_fft=n_fft, dc4_mv=dc4_mv, row_avg_n=row_avg_n)
    if fast is not None:
        return fast

    # ---- Chunked FFT path --------------------------------------------------
    spf = sras.samples_per_frame
    n_len = n_fft if n_fft is not None else spf
    block = _fft_block_for(spf, n_len)
    freq32 = sras.freq_axis_mhz(n_fft).astype(np.float32)
    # First fine bin at or above the floor; searchsorted is exact here since
    # freq32 is the very axis the floor is expressed against.
    min_bin = max(1, int(np.searchsorted(freq32, min_freq_mhz)))
    img = np.zeros((n_rows, n_frames), dtype=np.float32)
    background = sras.background if (apply_bg_sub and sras.background is not None) else None
    cal4 = sras.cal(CH4_IDX)
    row_avg_weights = _row_average_weights(row_avg_n) if row_avg_n > 0 else None

    total = _TOTAL_BYTES_BUDGET if budget is None else max(1, budget)
    if row_avg_n > 0:
        # One extra same-sized transient buffer (the pre-averaging full-row
        # scratch array) is live per in-flight row; halve the budget so
        # _plan_fft_rows accounts for it rather than relying on its existing
        # 2x slack to happen to cover it.
        total = max(1, total // 2)
    cap = max_workers if max_workers is not None else _MAX_WORKERS
    chunk_rows = _plan_fft_rows(n_frames, spf, total)
    pool = _fft_pool() if cap > 1 else None

    def process(r0: int, r1: int):
        if dc_threshold_mv is None:
            valid = None
        else:
            if dc4_mv is not None:
                dc4_chunk = dc4_mv[r0:r1]
            else:
                dc4_chunk = adc_to_mv(
                    data[r0:r1, CH4_IDX, :, :].astype(np.float32).mean(axis=-1), *cal4)
            valid = dc4_chunk >= dc_threshold_mv   # True = run the FFT
            if not valid.any():
                return

        counts = (valid.sum(axis=1) if valid is not None
                  else np.full(r1 - r0, n_frames, dtype=np.int64))
        offs = np.concatenate(([0], np.cumsum(counts)))
        n_wf = int(offs[-1])
        waves = np.empty((n_wf, spf), dtype=np.float32)

        def read_row(i: int):
            if should_stop is not None and should_stop():
                return
            raw = data[r0 + i, CH1_IDX]
            dst = waves[offs[i]:offs[i + 1]]
            row_valid = valid[i] if valid is not None else None
            if row_avg_weights is not None:
                v = row_valid if row_valid is not None else np.ones(n_frames, dtype=bool)
                # A plain cast + broadcast multiply, not a boolean-indexed
                # scatter into a zeroed buffer: numpy's fancy indexing holds
                # the GIL for its whole duration (measured zero speedup
                # across threads, and net negative once threads outnumber
                # physical cores), while a cast and a multiply are ordinary
                # ufuncs that release it — this is what lets read_row actually
                # parallelize across the pool instead of serializing on the
                # scatter/gather. Trade-off: every sample in the row is read,
                # valid or not (row averaging needs broad neighbor context
                # regardless, unlike the plain path below, which still skips
                # masked-out pixels entirely).
                full = raw.astype(np.float32) * v[:, None]
                avg = _row_average_waveforms(full, v, row_avg_weights)
                dst[:] = avg[v] if row_valid is not None else avg
            else:
                # Index the raw memmap slice with the boolean mask *before*
                # converting dtype — this is a lazy view until touched, so
                # only the selected elements are actually read from disk;
                # masked-out pixels' pages are never paged in at all.
                dst[:] = raw[row_valid] if row_valid is not None else raw
            if background is not None:
                dst -= background             # background is 1-D (spf,)

        if pool is None:
            for i in range(r1 - r0):
                read_row(i)
        else:
            list(pool.map(read_row, range(r1 - r0)))

        out = np.empty(n_wf, dtype=np.float32)

        def fft_block(b0: int):
            if should_stop is not None and should_stop():
                return
            b1 = min(b0 + block, n_wf)
            out[b0:b1] = freq32[_peak_bins(waves[b0:b1], n_len, min_bin)]

        if pool is None:
            for b0 in range(0, n_wf, block):
                fft_block(b0)
        else:
            list(pool.map(fft_block, range(0, n_wf, block)))

        # Boolean scatter is row-major, matching the read pass's
        # concatenation order.
        if valid is not None:
            img[r0:r1][valid] = out
        else:
            img[r0:r1] = out.reshape(r1 - r0, n_frames)

    for r0 in range(0, n_rows, chunk_rows):
        if should_stop is not None and should_stop():
            break
        process(r0, min(r0 + chunk_rows, n_rows))
    return img


# ---------------------------------------------------------------------------
# Batch cache (Convert menu) — process-pool entry point
# ---------------------------------------------------------------------------

def cache_file(path: str, mode: str, apply_bg_sub: bool,
               max_workers: int = 0,
               pad_factor: int = 1,
               dc_threshold_mv: float | None = None,
               row_avg_n: int = 0) -> str:
    """Compute and store DC or FFT images for every angle of one file,
    converting v6 → v7 in place. Returns "" on success or an error message.

    Module-level and picklable so it can run in a ProcessPoolExecutor. The
    worker cap is passed explicitly because module globals do not survive a
    spawn.

    mode is "dc", "fft" (raw per-pixel FFT, unmasked — masking applied at
    display time), or "fft_rowavg" (same-row, distance-weighted CH1
    averaging before the FFT — see compute_rf_image's row_avg_n).
    fft_rowavg needs *dc_threshold_mv* up front, unlike plain "fft":
    neighbor validity is baked into the stored numbers, so it can't be
    deferred to display time the way plain masking can.

    *pad_factor* is the zero-padding factor to resolve the peaks at: 1 (the
    default) is natural resolution, n_fft == samples_per_frame. It is
    recorded in the SFFT block so a reader knows which views the stored
    numbers answer for — and it has to be the pad the viewer is *actually*
    using, since a pad-1 cache is dead weight to a padded view and vice
    versa (cached_rf_image refuses the mismatch rather than showing peaks
    resolved at the wrong resolution).
    """
    global _MAX_WORKERS
    try:
        if mode not in ("dc", "fft", "fft_rowavg"):
            return f"unknown cache mode {mode!r} (expected 'dc', 'fft', or 'fft_rowavg')"
        # write_v7_cache enforces this bound too, but only once every angle's
        # FFT has already been computed. Checking up front is the difference
        # between a bad argument costing nothing and costing the whole run.
        if not (1 <= pad_factor <= MAX_PAD_FACTOR):
            return f"pad_factor must be 1-{MAX_PAD_FACTOR}, got {pad_factor}"
        if max_workers:
            _MAX_WORKERS = max_workers

        sras = SrasFile(path)
        if sras.version not in (6, 7):
            return (f"unsupported version {sras.version} — only v6/v7 files "
                    "can be batch-cached")

        n = sras.n_angles
        n_fft = sras.samples_per_frame * pad_factor if pad_factor > 1 else None
        n_workers, angle_budget = plan_angle_level(sras)
        if mode == "dc":
            dc3 = _parallel_map(
                lambda a: adc_to_mv(
                    compute_dc_image(sras, a, CH3_IDX, max_workers=1,
                                     budget=angle_budget), *sras.cal(CH3_IDX)),
                range(n), n_workers)
            dc4 = _parallel_map(
                lambda a: adc_to_mv(
                    compute_dc_image(sras, a, CH4_IDX, max_workers=1,
                                     budget=angle_budget), *sras.cal(CH4_IDX)),
                range(n), n_workers)
            sras.write_v7_cache(new_dc3_mv=dc3, new_dc4_mv=dc4)
        elif mode == "fft":
            effective_bg = apply_bg_sub and sras.background is not None
            # dc_threshold_mv=None: store unmasked images and mask at display
            # time (same convention as v5's PREC block). Skipping the mask
            # also skips reading CH4 entirely. The FFT path parallelises
            # internally over blocks, so angles run one at a time with the
            # full budget.
            freq = [compute_rf_image(sras, a, dc_threshold_mv=None,
                                     apply_bg_sub=effective_bg, n_fft=n_fft)
                    for a in range(n)]
            # new_row_avg_n=0 explicitly: these images are raw per-pixel FFTs,
            # and carrying forward a row_avg_n left by an earlier fft_rowavg
            # write would label them as something they are not.
            sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
                                new_row_avg_n=0, new_pad_factor=pad_factor)
        else:   # "fft_rowavg"
            if row_avg_n <= 0:
                return "row_avg_n must be a positive neighbor half-width for fft_rowavg mode"
            if dc_threshold_mv is None:
                return ("fft_rowavg mode requires a DC threshold "
                        "(neighbor validity depends on it)")
            effective_bg = apply_bg_sub and sras.background is not None
            freq = [compute_rf_image(sras, a, dc_threshold_mv=dc_threshold_mv,
                                     apply_bg_sub=effective_bg, n_fft=n_fft,
                                     row_avg_n=row_avg_n)
                    for a in range(n)]
            sras.write_v7_cache(new_freq_mhz=freq, new_bg_sub=effective_bg,
                                new_row_avg_n=row_avg_n,
                                new_pad_factor=pad_factor)
        return ""
    except Exception as exc:
        return str(exc)


# ---------------------------------------------------------------------------
# Angle alignment (Fusion menu)
#
# Rigid transforms only (rotation + translation, never scale), computed in mm
# on two frames: each angle's "local mm" (origin at its own array center) and
# the reference angle's local mm. Angle 0 is the sole coordinate authority;
# every other angle is placed purely by image content. Why, and the full
# frame/affine conventions: docs/design.md ("Angle alignment coordinate
# frames").
# ---------------------------------------------------------------------------

@dataclass
class AngleTransform:
    rotation_deg: float
    shift_mm: tuple[float, float]      # (dx_mm, dy_mm) in ref mm
    matrix: np.ndarray                 # (2,2): canvas (row,col) -> this angle's raw (row,col)
    offset: np.ndarray                 # (2,)
    score: float = 1.0                 # registration NCC (1.0 = reference/manual)
    source: str = ""                   # image that won registration: "signal"/"mask"


@dataclass
class AlignmentResult:
    ref_angle_idx: int
    dc_threshold_mv: float
    canvas_shape: tuple[int, int]          # (n_rows, n_cols)
    canvas_dx_mm: float
    canvas_dy_mm: float
    canvas_origin_mm: tuple[float, float]  # mm at canvas pixel index (0, 0)
    per_angle: dict[int, AngleTransform]


@dataclass
class ManualAngleParams:
    """One angle's manual-alignment state, independent of any canvas.

    rotation_deg/shift_mm are exactly AngleTransform's non-derived fields — the
    rigid map from this angle's local mm to ref mm, the pair a canvas-bound
    AngleTransform's matrix/offset get built from once a canvas is decided
    (build_manual_alignment). Defaults to identity: a fresh angle with no prior
    alignment is shown centered on the reference with no rotation, which is the
    same "fully unaligned" state Clear Alignment resets back to.
    """
    rotation_deg: float = 0.0
    shift_mm: tuple[float, float] = (0.0, 0.0)


# Side length of the grid the rotation refinement runs on. The whole cost of
# registration scales with it: ~325 ms per candidate rotation at 640, roughly
# quadrupling per doubling, against ~200 MB of live masked-correlation buffers
# (see _registration_workers). 640 puts a full-size scan at ~0.1 mm/px, which
# resolves rotation on a 20 mm sample to well under a tenth of a degree.
_DEFAULT_FINE_DIM = 640


@dataclass
class RigidFit:
    """What register_angle_to_reference found for one angle."""
    rotation_deg: float
    shift_mm: tuple[float, float]
    score: float                       # zero-mean NCC over the valid overlap
    source: str                        # "signal", "mask", or "reference"


def _skimage_phase_cross_correlation():
    """Import skimage.registration lazily and cache it.

    Deliberately not a module-level import: this module is imported by every
    multiprocessing child (see the module docstring), skimage costs ~0.6 s to
    import, and no child ever registers anything — registration runs in GUI-
    process threads.
    """
    global _pcc
    try:
        return _pcc
    except NameError:
        from skimage.registration import phase_cross_correlation as _fn
        _pcc = _fn
        return _pcc


# ---- Geometry: local mm, ref mm, and the one affine builder ---------------

def pixel_pitch_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
    """(dx, dy) mm/pixel for one angle: dx is the file-wide constant
    pixel_x_mm; dy is this angle's own row spacing (assumed uniform, the same
    assumption _redraw_image makes when it builds the display extent). dy keeps
    its sign, so +row always means the same physical direction as +y."""
    y = sras.y_positions_mm(angle_idx)
    return sras.pixel_x_mm, (float(y[1] - y[0]) if len(y) > 1 else 1.0)


def _center_idx(sras: SrasFile, angle_idx: int) -> np.ndarray:
    """(row, col) index of this angle's array center — the origin of its local
    mm frame. Purely geometric: it depends on the array shape and nothing
    else, which is what keeps local mm free of stage position."""
    n_rows, n_frames = sras.image_shape(angle_idx)
    return np.array([(n_rows - 1) / 2.0, (n_frames - 1) / 2.0])


def ref_center_mm(sras: SrasFile, ref_angle_idx: int) -> np.ndarray:
    """Stage mm of the reference angle's array center: ref mm + this == stage
    mm.

    The single bridge between ref mm and stage mm, and the only place in the
    whole alignment path where any angle's stage position is read at all —
    which is why it takes the *reference* index by name rather than an
    arbitrary angle.
    """
    x = sras.x_axis_mm(ref_angle_idx)
    y = sras.y_positions_mm(ref_angle_idx)
    return np.array([(x[0] + x[-1]) / 2.0, (y[0] + y[-1]) / 2.0], dtype=np.float64)


def _local_half_extent_mm(sras: SrasFile, angle_idx: int) -> tuple[float, float]:
    """(half width, half height) in mm from this angle's array center to the
    center of its outermost pixel."""
    n_rows, n_frames = sras.image_shape(angle_idx)
    dx, dy = pixel_pitch_mm(sras, angle_idx)
    return (n_frames - 1) / 2.0 * abs(dx), (n_rows - 1) / 2.0 * abs(dy)


def _rotation_matrix(theta_deg: float) -> np.ndarray:
    t = np.radians(theta_deg)
    c, s = np.cos(t), np.sin(t)
    return np.array([[c, -s], [s, c]])   # CCW rotation acting on (x, y)


def nominal_delta_deg(sras: SrasFile, angle_idx: int, ref_idx: int) -> float:
    """The rotation-stage's own reported angle change between two angles.

    Used only to *seed* the rotation search, never as the answer: the stage's
    sign convention relative to this module's math-positive (CCW, x toward y)
    convention in scan mm is not knowable from the file, so
    register_angle_to_reference scores both +this and -this and lets the image
    content decide (see _rotation_candidates).
    """
    return float(sras.angles_deg[angle_idx] - sras.angles_deg[ref_idx])


def _footprint_corners_ref_mm(sras: SrasFile, angle_idx: int,
                              rotation_deg: float,
                              shift_mm: tuple[float, float]) -> np.ndarray:
    """This angle's 4 footprint corners mapped into ref mm by its rigid
    transform, shape (4, 2). Zero-padding cost, not content: it is the scan
    window's corners, which is what the shared canvas has to cover."""
    hw, hh = _local_half_extent_mm(sras, angle_idx)
    corners = np.array([[sx * hw, sy * hh] for sx in (-1.0, 1.0) for sy in (-1.0, 1.0)])
    R = _rotation_matrix(rotation_deg)
    return corners @ R.T + np.asarray(shift_mm, dtype=np.float64)


def _affine_out_to_src(*, out_pitch_mm: tuple[float, float],
                       out_origin_ref_mm, src_dx_mm: float, src_dy_mm: float,
                       src_center_idx, src_center_off_mm=(0.0, 0.0),
                       rotation_deg: float = 0.0,
                       shift_mm: tuple[float, float] = (0.0, 0.0)
                       ) -> tuple[np.ndarray, np.ndarray]:
    """matrix, offset s.t. src_index = matrix @ [row_out, col_out] + offset,
    matching scipy.ndimage.affine_transform's output->input convention.

    The single affine builder behind every resampling in this module — the
    registration grid, the manual-alignment preview and the final canvas all
    differ only in their arguments.

    Pipeline (mm unless noted):
      [X;Y]   = A_out @ [row_out;col_out] + out_origin_ref_mm   # out idx -> ref mm
      [lx;ly] = R(rotation)^T @ ([X;Y] - shift)                 # ref mm -> this angle's local mm
      [row;col] = D @ ([lx;ly] - src_center_off) + src_center_idx

    *src_center_off_mm* is the local-mm position of the source array's own
    center, nonzero only when the source has been block-mean downsampled (its
    center can land up to half a block off the full-resolution center — see
    _prepare_reg_image). Everything is expressed relative to array centers, so
    no per-angle stage coordinate appears anywhere in here.
    """
    Rinv = _rotation_matrix(rotation_deg).T
    A_out = np.array([[0.0, out_pitch_mm[0]], [out_pitch_mm[1], 0.0]])  # [row,col] -> [x,y]
    D = np.array([[0.0, 1.0 / src_dy_mm], [1.0 / src_dx_mm, 0.0]])      # [x,y] -> [row,col]
    b_out = np.asarray(out_origin_ref_mm, dtype=np.float64)
    shift = np.asarray(shift_mm, dtype=np.float64)
    off = np.asarray(src_center_off_mm, dtype=np.float64)

    matrix = D @ Rinv @ A_out
    offset = D @ (Rinv @ (b_out - shift) - off) + np.asarray(src_center_idx, dtype=np.float64)
    return matrix, offset


def apply_alignment(result: AlignmentResult, angle_idx: int, img: np.ndarray,
                    order: int = 0) -> np.ndarray:
    """Resample any already-computed 2D image for `angle_idx` (same shape as
    that angle's raw (n_rows, n_frames)) onto the shared alignment canvas.
    order=0 (nearest) avoids blending real data with zero-padding or with
    masked-out (0-valued) CH1/velocity pixels at mask edges. Channel-
    agnostic: the same per-angle transform (found from the CH4 image) works
    for any channel's image of that angle."""
    t = result.per_angle[angle_idx]
    return scipy_ndimage.affine_transform(
        img.astype(np.float32, copy=False), t.matrix, offset=t.offset,
        output_shape=result.canvas_shape, order=order,
        mode="constant", cval=0.0)


def block_mean_2d(img: np.ndarray, fy: int, fx: int) -> np.ndarray:
    """Block-mean by independent row/column factors. Independent factors matter
    because the raw grid is strongly anisotropic (5 µm along x, 50 µm along y):
    a single square factor would either alias along x or throw away rows."""
    fy, fx = max(1, int(fy)), max(1, int(fx))
    if fy == 1 and fx == 1:
        return img
    h, w = img.shape
    h2, w2 = (h // fy) * fy, (w // fx) * fx
    if h2 == 0 or w2 == 0:
        return img
    trimmed = img[:h2, :w2]
    return trimmed.reshape(h2 // fy, fy, w2 // fx, fx).mean(axis=(1, 3))


# ---- Registration: rigid (rotation + translation) fit against the reference
#
# Everything below works on _RegImage, an angle's image resampled to a shared
# *isotropic* grid centered on its own array center. Nothing in here can see a
# stage coordinate even in principle, which is the point: the fit is decided by
# image content alone.

@dataclass
class _RegImage:
    """One angle's image prepared for registration: block-mean downsampled to
    roughly the registration pitch, carrying the physical pitch it ended up
    with and the local-mm offset of its own array center from the
    full-resolution array center (block-mean trims a partial trailing block, so
    the two centers can differ by up to half a block)."""
    img: np.ndarray
    dx_mm: float
    dy_mm: float
    center_off_mm: tuple[float, float]


def _block_center_off(n_full: int, n_small: int, factor: int,
                      pitch_mm: float) -> float:
    """Local-mm offset of a block-mean-downsampled array's own center from the
    full-resolution array's center along one axis.

    Small pixel k averages full pixels [k*f, k*f + f - 1], so its center sits
    at full index k*f + (f-1)/2; block-mean also trims a partial trailing
    block. Both effects together move the small array's center by up to half a
    block, which has to be accounted for or a downsampled image registers (or
    previews) at a systematically shifted position.
    """
    return (((n_small - 1) / 2.0) * factor + (factor - 1) / 2.0
            - (n_full - 1) / 2.0) * pitch_mm


def _prepare_reg_image(sras: SrasFile, angle_idx: int, img: np.ndarray,
                       pitch_mm: float) -> _RegImage:
    """Block-mean an angle's image down to roughly *pitch_mm* before it is
    resampled onto the registration grid. Pre-averaging matters: the raw grid
    is 10x finer along x than along y, so sampling it directly at the (much
    coarser) isotropic registration pitch would alias badly along x."""
    dx, dy = pixel_pitch_mm(sras, angle_idx)
    fx = max(1, int(pitch_mm / abs(dx)))
    fy = max(1, int(pitch_mm / abs(dy)))
    small = block_mean_2d(np.asarray(img, dtype=np.float32), fy, fx)
    n_rows, n_frames = img.shape
    return _RegImage(
        small, dx * fx, dy * fy,
        (_block_center_off(n_frames, small.shape[1], fx, dx),
         _block_center_off(n_rows, small.shape[0], fy, dy)))


def _embed(reg: _RegImage, pitch_mm: float, n: int, rotation_deg: float,
           order: int) -> np.ndarray:
    """Resample a _RegImage onto the shared n x n isotropic registration grid,
    rotated by *rotation_deg* about the grid center and with no translation
    (translation is what phase correlation then measures)."""
    half = (n - 1) / 2.0 * pitch_mm
    matrix, offset = _affine_out_to_src(
        out_pitch_mm=(pitch_mm, pitch_mm), out_origin_ref_mm=(-half, -half),
        src_dx_mm=reg.dx_mm, src_dy_mm=reg.dy_mm,
        src_center_idx=((reg.img.shape[0] - 1) / 2.0, (reg.img.shape[1] - 1) / 2.0),
        src_center_off_mm=reg.center_off_mm, rotation_deg=rotation_deg)
    return scipy_ndimage.affine_transform(
        reg.img, matrix, offset=offset, output_shape=(n, n), order=order,
        mode="constant", cval=0.0)


def _embed_with_valid(reg: _RegImage, pitch_mm: float, n: int,
                      rotation_deg: float) -> tuple[np.ndarray, np.ndarray]:
    """Embedded image plus the boolean mask of where that angle actually has
    data. The valid mask is what lets registration ignore each angle's
    differently-shaped scan window instead of locking onto its silhouette."""
    img = _embed(reg, pitch_mm, n, rotation_deg, order=1)
    ones = _RegImage(np.ones_like(reg.img), reg.dx_mm, reg.dy_mm, reg.center_off_mm)
    valid = _embed(ones, pitch_mm, n, rotation_deg, order=0) > 0.5
    return img, valid


def _shift_into(img: np.ndarray, dr: int, dc: int) -> np.ndarray:
    """img translated by whole pixels with zero fill (never wrapping, unlike
    np.roll — wrapped content would score as a spurious match)."""
    out = np.zeros_like(img)
    h, w = img.shape
    sr0, sr1 = max(0, dr), min(h, h + dr)
    sc0, sc1 = max(0, dc), min(w, w + dc)
    if sr0 >= sr1 or sc0 >= sc1:
        return out
    out[sr0:sr1, sc0:sc1] = img[sr0 - dr:sr1 - dr, sc0 - dc:sc1 - dc]
    return out


def _overlap_ncc(ref: np.ndarray, ref_valid: np.ndarray,
                 mov: np.ndarray, mov_valid: np.ndarray,
                 min_overlap_frac: float = 0.15) -> float:
    """Zero-mean normalized cross-correlation over the two images' common valid
    region — the score every rotation candidate is ranked by.

    Computed on the overlap only, and rejected outright (-1) when the overlap
    is too small a fraction of the smaller footprint: without that floor a
    candidate that slides the angles almost entirely apart can win on a handful
    of coincidentally-similar pixels.
    """
    both = ref_valid & mov_valid
    n = int(both.sum())
    smaller = min(int(ref_valid.sum()), int(mov_valid.sum()))
    if smaller == 0 or n < min_overlap_frac * smaller or n < 16:
        return -1.0
    a = ref[both].astype(np.float64)
    b = mov[both].astype(np.float64)
    a -= a.mean()
    b -= b.mean()
    denom = np.sqrt((a * a).sum() * (b * b).sum())
    return float((a * b).sum() / denom) if denom > 0 else -1.0


def _masked_shift(ref: np.ndarray, ref_valid: np.ndarray,
                  mov: np.ndarray, mov_valid: np.ndarray) -> tuple[int, int]:
    """Integer (dr, dc) that best registers *mov* onto *ref*, from skimage's
    masked FFT phase correlation (Padfield). The masked variant is the whole
    reason scikit-image is a dependency: plain phase correlation on these
    images locks onto the scan window's rectangular silhouette, which differs
    per angle, instead of onto the sample."""
    pcc = _skimage_phase_cross_correlation()
    result = pcc(ref, mov, reference_mask=ref_valid, moving_mask=mov_valid)
    shift = result[0] if isinstance(result, tuple) else result
    return int(round(float(shift[0]))), int(round(float(shift[1])))


def _subpixel_residual(ref: np.ndarray, mov: np.ndarray,
                       both: np.ndarray) -> tuple[float, float]:
    """Sub-pixel leftover shift between two already integer-aligned images,
    from upsampled phase correlation over their common valid region.

    A separate pass because skimage's *masked* phase correlation has no
    upsample_factor; here both images are zeroed outside the shared overlap and
    mean-subtracted inside it, so the plain upsampled version is well posed.
    Clamped to ±1 px: this only ever polishes an already-good integer fit, and
    a larger "residual" means the peak was spurious.

    normalization=None (plain cross-correlation, not phase correlation) is
    load-bearing. Whitening the spectrum is what makes phase correlation good
    at finding a large unknown shift, but here the two images are already
    aligned to within a pixel and the masked-off surroundings put a hard edge
    in both: whitened, that edge and the high-frequency noise swamp the true
    sub-pixel peak and the default returns a flat zero every time.
    """
    if not both.any():
        return 0.0, 0.0
    a = np.zeros_like(ref)
    b = np.zeros_like(mov)
    a[both] = ref[both] - ref[both].mean()
    b[both] = mov[both] - mov[both].mean()
    if not (np.any(a) and np.any(b)):
        return 0.0, 0.0
    pcc = _skimage_phase_cross_correlation()
    result = pcc(a, b, upsample_factor=20, normalization=None)
    shift = result[0] if isinstance(result, tuple) else result
    dr, dc = float(shift[0]), float(shift[1])
    if abs(dr) > 1.0 or abs(dc) > 1.0:
        return 0.0, 0.0
    return dr, dc


def _reg_pitch_and_size(sras: SrasFile, max_dim: int,
                        margin: float = 1.25) -> tuple[float, int]:
    """Isotropic pitch (mm/px) and side length for the shared registration
    grid: square, big enough for the largest angle's footprint at any rotation
    (hence its diagonal) plus *margin* headroom for the translation search.

    The floor on pitch is the *geometric mean* of the two native pitches, not
    the coarser of them. The raw grid is anisotropic (5 µm along x, 50 µm along
    y): flooring at 50 µm would throw away all the extra x detail, and rotation
    precision depends directly on it — a feature at radius r moves by r·δθ, so
    at 50 µm a 20 mm-wide sample can only resolve rotation to a few tenths of a
    degree. The geometric mean interpolates y up by ~3x rather than discarding
    x, which costs a little memory and buys real angular precision. On a
    full-size scan max_dim binds first and this floor never applies at all.
    """
    diag = max(float(np.hypot(*(2 * v for v in _local_half_extent_mm(sras, a))))
               for a in range(sras.n_angles))
    span = diag * margin
    native = float(np.sqrt(
        abs(sras.pixel_x_mm)
        * max(abs(pixel_pitch_mm(sras, a)[1]) for a in range(sras.n_angles))))
    pitch = max(span / max_dim, native)
    n = int(scipy_fft.next_fast_len(max(16, int(np.ceil(span / pitch)))))
    return pitch, n


def _registration_workers(sras: SrasFile, fine_dim: int) -> int:
    """How many angles may register concurrently.

    Not plan_angle_level: that budgets for waveform chunks, and registration
    never touches a waveform — it works on already-computed DC images and a few
    grid-sized arrays. The real limit is the masked phase correlation, which
    pads to roughly twice the grid and holds several complex128 arrays of that
    size live at once, so the count is derived from the same
    _TOTAL_BYTES_BUDGET the rest of the module honours. The estimate below
    comes out at 262 MB for the default fine_dim=640, against 199 MB measured —
    deliberately on the pessimistic side, since overshooting the budget costs
    swapping while undershooting only costs a little wall time.
    """
    per_worker = 10 * (2 * fine_dim) ** 2 * 16      # ~10 complex128 grids
    return int(max(1, min(_MAX_WORKERS, sras.n_angles,
                          _TOTAL_BYTES_BUDGET // max(1, per_worker))))


def registration_workers(sras: SrasFile) -> int:
    """Public: concurrent-registration cap at the default fine grid."""
    return _registration_workers(sras, _DEFAULT_FINE_DIM)


def default_max_workers() -> int:
    """Public: the module-wide worker cap (SRAS_MAX_WORKERS or cpu count)."""
    return _MAX_WORKERS


def _rotation_candidates(nominal_deg: float, search_deg: float,
                         step_deg: float,
                         signs: tuple[int, ...] = (-1, 1)) -> list[float]:
    """Coarse rotation candidates: a window around each requested sign of the
    stage's reported angle change. Scoring both signs (the default) is what
    makes the stage's sign convention a non-issue — the images decide which way
    the stage turns, and a file whose stage reports the opposite sense
    registers just as well.

    *signs* exists only so a user who has established which way their own stage
    turns can halve the coarse sweep from the alignment wizard. It is an
    override, not an inference: nothing in the file says which sign is right.
    """
    out: list[float] = []
    step_deg = max(abs(step_deg), 1e-9)   # a 0 step would divide by zero below
    for sign in signs:
        center = sign * nominal_deg
        k = int(np.floor(search_deg / step_deg))
        for i in range(-k, k + 1):
            out.append(center + i * step_deg)
    # Dedupe (the two windows coincide when nominal_deg is 0) while keeping order.
    seen: set[float] = set()
    return [t for t in out if not (round(t, 6) in seen or seen.add(round(t, 6)))]


def _score_rotation(ref_img: np.ndarray, ref_valid: np.ndarray,
                    mov: _RegImage, pitch: float, n: int, theta: float,
                    subpixel: bool = False) -> tuple[float, tuple[float, float]]:
    """Best score this rotation can reach, and the translation that reaches it:
    rotate, phase-correlate for the shift, score the overlap.

    *subpixel* also removes the leftover sub-pixel translation before scoring.
    That matters more than it sounds: without it every candidate is scored at
    whole-pixel alignment, so rotations that differ by less than one pixel of
    rim displacement are ranked by quantization noise rather than by fit, and
    the refinement stalls a degree or so off. Left off for the coarse sweep,
    which only has to pick a basin, and on for the refinement.
    """
    img, valid = _embed_with_valid(mov, pitch, n, theta)
    dr, dc = _masked_shift(ref_img, ref_valid, img, valid)
    shifted = _shift_into(img, dr, dc)
    shifted_valid = _shift_into(valid.astype(np.float32), dr, dc) > 0.5
    if subpixel:
        sub_dr, sub_dc = _subpixel_residual(
            ref_img, shifted, ref_valid & shifted_valid)
        if sub_dr or sub_dc:
            shifted = scipy_ndimage.shift(shifted, (sub_dr, sub_dc), order=1,
                                          mode="constant", cval=0.0)
            dr, dc = dr + sub_dr, dc + sub_dc
    return (_overlap_ncc(ref_img, ref_valid, shifted, shifted_valid),
            (float(dr), float(dc)))


def _refine_rotation(ref_img: np.ndarray, ref_valid: np.ndarray,
                     mov: _RegImage, pitch: float, n: int,
                     theta: float, score: float, shift: tuple[float, float],
                     step_deg: float, min_step_deg: float = 0.05,
                     max_evals: int = 80
                     ) -> tuple[float, float, tuple[float, float]]:
    """Hill-climb the rotation from the coarse winner: step out while the score
    improves, halve the step when it doesn't, stop below *min_step_deg*.

    A walking search rather than a fixed grid around the coarse winner, because
    the coarse stage ranks rotations at a coarse pitch where a degree can be
    worth less than the translation quantization — its winner can legitimately
    land a degree or two off, which a fixed ±½° refinement window could never
    recover from.

    Every candidate here is scored with subpixel=True, matching how the
    incoming *score* was measured. Mixing the two is silently fatal: the
    sub-pixel-corrected score is strictly the higher of the two, so a
    subpixel-scored start compared against un-corrected candidates can never be
    beaten and the search sits still at whatever the coarse stage handed it.
    """
    evals = 0
    while step_deg >= min_step_deg and evals < max_evals:
        trials = []
        for cand in (theta - step_deg, theta + step_deg):
            s, sh = _score_rotation(ref_img, ref_valid, mov, pitch, n, cand,
                                    subpixel=True)
            evals += 1
            trials.append((s, cand, sh))
        best_s, best_cand, best_sh = max(trials, key=lambda t: t[0])
        if best_s > score:
            theta, score, shift = best_cand, best_s, best_sh
        else:
            step_deg /= 2.0
    return theta, score, shift


def _source_images(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
                   signal_mv: dict[int, np.ndarray], dc_threshold_mv: float,
                   sources) -> list[tuple[str, np.ndarray, np.ndarray]]:
    """(name, reference image, moving image) per requested registration source.

    "signal" is each angle's own CH4 image minus its own minimum — subtracting
    per-angle rather than globally keeps a nonzero DC baseline from reading as
    a high-contrast edge against the zero padding. "mask" is the binarized
    >= dc_threshold_mv image, the same silhouette the overlay draws. A mask
    that is empty or completely full for either angle carries no registration
    information at all, so that source is dropped rather than scored.
    """
    out = []
    for name in sources:
        pair = []
        for a in (ref_angle_idx, angle_idx):
            img = signal_mv[a]
            if name == "mask":
                m = img >= dc_threshold_mv
                if not m.any() or m.all():
                    pair = []
                    break
                pair.append(m.astype(np.float32))
            else:
                pair.append((img - img.min()).astype(np.float32))
        if pair:
            out.append((name, pair[0], pair[1]))
    return out


def register_angle_to_reference(
        sras: SrasFile, angle_idx: int, ref_angle_idx: int,
        signal_mv: dict[int, np.ndarray], *,
        dc_threshold_mv: float = 0.0,
        sources: tuple[str, ...] = ("signal", "mask"),
        coarse_dim: int = 256, fine_dim: int = _DEFAULT_FINE_DIM,
        search_deg: float = 6.0, coarse_step_deg: float = 2.0,
        seed_deg: float | None = None,
        seed_signs: tuple[int, ...] = (-1, 1),
        refine: bool = True) -> RigidFit:
    """Rigid (rotation + translation, never scale) fit of *angle_idx* onto
    *ref_angle_idx*, found entirely by cross-correlating image content.

    Two stages:
      1. Coarse sweep at *coarse_dim* over a ±*search_deg* window around the
         requested signs of the stage's reported angle change (see
         _rotation_candidates), for each requested source image, scored by
         _overlap_ncc. Only has to pick the right basin.
      2. Hill-climbing refinement of the winning (source, rotation) at
         *fine_dim* with sub-pixel translation folded into every score
         (_refine_rotation, _score_rotation), down to 0.05°.

    The returned rotation_deg/shift_mm are the rigid map from this angle's
    local mm to ref mm (q = R @ l + shift); score is the final NCC, which the
    caller can surface so a bad scan is visible rather than silently fused in.
    Returns identity for the reference angle itself.

    The last three arguments only exist to let the alignment wizard expose the
    rotation search; every default reproduces the search this function has
    always done.

    *seed_deg* replaces the stage's reported angle change as the center of the
    coarse sweep — the pre-rotation the search starts from. None (the default)
    means the stage angle from the file's Angle Table, which is nearly always
    what you want: it puts the sweep within a couple of degrees of the answer.
    Pass 0.0 to search around no rotation at all, which is the honest choice for
    a file whose stage angles are known to be wrong.

    *refine=False* stops after the coarse sweep, leaving rotation on the
    *coarse_step_deg* grid. Combined with search_deg=0.0 and a single seed sign
    it pins rotation to exactly the seed and searches translation only — for a
    scan whose stage angles are trusted more than its image content.
    """
    if angle_idx == ref_angle_idx:
        return RigidFit(0.0, (0.0, 0.0), 1.0, "reference")

    candidates = _source_images(sras, angle_idx, ref_angle_idx, signal_mv,
                                dc_threshold_mv, sources)
    if not candidates:
        return RigidFit(0.0, (0.0, 0.0), -1.0, "none")

    nominal = (nominal_delta_deg(sras, angle_idx, ref_angle_idx)
               if seed_deg is None else float(seed_deg))
    thetas = _rotation_candidates(nominal, search_deg, coarse_step_deg,
                                  seed_signs)

    # ---- Stage 1: coarse sweep, every source ------------------------------
    pitch_c, n_c = _reg_pitch_and_size(sras, coarse_dim)
    best = (-2.0, 0.0, (0, 0), "none")   # score, theta, (dr, dc), source
    for name, ref_raw, mov_raw in candidates:
        ref_reg = _prepare_reg_image(sras, ref_angle_idx, ref_raw, pitch_c)
        mov_reg = _prepare_reg_image(sras, angle_idx, mov_raw, pitch_c)
        ref_img, ref_valid = _embed_with_valid(ref_reg, pitch_c, n_c, 0.0)
        for theta in thetas:
            score, shift = _score_rotation(ref_img, ref_valid, mov_reg,
                                           pitch_c, n_c, theta)
            if score > best[0]:
                best = (score, theta, shift, name)

    if best[3] == "none":
        return RigidFit(0.0, (0.0, 0.0), -1.0, "none")

    # ---- Stage 2: refine the winner at full registration resolution -------
    name = best[3]
    ref_raw, mov_raw = next((r, m) for n_, r, m in candidates if n_ == name)
    pitch_f, n_f = _reg_pitch_and_size(sras, fine_dim)
    ref_reg = _prepare_reg_image(sras, ref_angle_idx, ref_raw, pitch_f)
    mov_reg = _prepare_reg_image(sras, angle_idx, mov_raw, pitch_f)
    ref_img, ref_valid = _embed_with_valid(ref_reg, pitch_f, n_f, 0.0)

    theta = best[1]
    score, shift = _score_rotation(ref_img, ref_valid, mov_reg, pitch_f, n_f,
                                   theta, subpixel=True)
    if refine:
        theta, score, shift = _refine_rotation(
            ref_img, ref_valid, mov_reg, pitch_f, n_f, theta, score, shift,
            step_deg=coarse_step_deg)

    dr, dc = shift
    return RigidFit(float(theta), (float(dc * pitch_f), float(dr * pitch_f)),
                    float(score), name)


# ---- Shared canvas: angle 0's own pixel grid, extended --------------------

def canvas_for_params(sras: SrasFile, ref_angle_idx: int,
                      pitch_mm: tuple[float, float],
                      per_angle_params: dict[int, ManualAngleParams],
                      ) -> tuple[tuple[float, float], tuple[int, int]]:
    """Shared-canvas origin (stage mm) and (n_rows, n_cols) at *pitch_mm* that
    contains every angle's footprint after its own rigid transform. Angles
    missing from per_angle_params default to identity (e.g. a sidecar saved
    before a rescan added more angles).

    The canvas grid is aligned with the reference angle's own pixel grid, so
    the reference lands on integer canvas pixels and is resampled by an exact
    integer translation — the concrete meaning of "the canvas carries angle
    0's X/Y coordinates". This requires pitch_mm to be the reference's own
    pitch.
    """
    dx, dy = pitch_mm
    corners = np.vstack([
        _footprint_corners_ref_mm(
            sras, a,
            per_angle_params.get(a, ManualAngleParams()).rotation_deg,
            per_angle_params.get(a, ManualAngleParams()).shift_mm)
        for a in range(sras.n_angles)])
    x_min, y_min = corners.min(axis=0)
    x_max, y_max = corners.max(axis=0)
    center = ref_center_mm(sras, ref_angle_idx)
    # Express the box in the reference's own pixel indices and grow it
    # outward to whole pixels, so canvas index k lands exactly where the
    # reference's own pixel (k + const) does.
    cy, cx = _center_idx(sras, ref_angle_idx)
    cols = sorted((x_min / dx + cx, x_max / dx + cx))
    rows = sorted((y_min / dy + cy, y_max / dy + cy))
    col0, col1 = int(np.floor(cols[0])), int(np.ceil(cols[1]))
    row0, row1 = int(np.floor(rows[0])), int(np.ceil(rows[1]))
    origin_ref = np.array([(col0 - cx) * dx, (row0 - cy) * dy])
    shape = (row1 - row0 + 1, col1 - col0 + 1)

    origin_stage = origin_ref + center
    return (float(origin_stage[0]), float(origin_stage[1])), shape


def build_canvas_affine(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
                        rotation_deg: float, shift_mm: tuple[float, float],
                        pitch_mm: tuple[float, float],
                        canvas_origin_mm: tuple[float, float],
                        *, src_downsample: tuple[int, int] = (1, 1)
                        ) -> tuple[np.ndarray, np.ndarray]:
    """canvas index -> this angle's raw index, for a canvas whose origin is
    given in *stage* mm (the reference's frame). *src_downsample* is the
    (rows, cols) block-mean factor already applied to the image the caller will
    resample — 1:1 for the raw image, coarser for the manual-alignment
    preview's downsampled masks."""
    dx_a, dy_a = pixel_pitch_mm(sras, angle_idx)
    n_rows, n_frames = sras.image_shape(angle_idx)
    fy, fx = (max(1, int(v)) for v in src_downsample)
    if (fy, fx) != (1, 1):
        # Same block-mean bookkeeping _prepare_reg_image does: the downsampled
        # array's own center can sit up to half a block off the full-resolution
        # center, and that offset has to be undone here or every preview layer
        # lands slightly (and inconsistently) off.
        nr_s, nf_s = n_rows // fy, n_frames // fx
        src_dx, src_dy = dx_a * fx, dy_a * fy
        src_center = ((nr_s - 1) / 2.0, (nf_s - 1) / 2.0)
        src_off = (_block_center_off(n_frames, nf_s, fx, dx_a),
                   _block_center_off(n_rows, nr_s, fy, dy_a))
    else:
        src_dx, src_dy = dx_a, dy_a
        src_center = _center_idx(sras, angle_idx)
        src_off = (0.0, 0.0)

    origin_ref = np.asarray(canvas_origin_mm, dtype=np.float64) \
        - ref_center_mm(sras, ref_angle_idx)
    return _affine_out_to_src(
        out_pitch_mm=pitch_mm, out_origin_ref_mm=origin_ref,
        src_dx_mm=src_dx, src_dy_mm=src_dy, src_center_idx=src_center,
        src_center_off_mm=src_off, rotation_deg=rotation_deg, shift_mm=shift_mm)


def _result_from_params(sras: SrasFile, ref_angle_idx: int,
                        dc_threshold_mv: float,
                        params: dict[int, ManualAngleParams],
                        extra: dict[int, tuple[float, str]] | None = None
                        ) -> AlignmentResult:
    """Assemble the final AlignmentResult from per-angle rigid parameters: pick
    the shared canvas, then build each angle's canvas->raw affine. Pure matrix
    and bbox math, so it is cheap enough to call synchronously on the GUI
    thread on every manual edit."""
    pitch = pixel_pitch_mm(sras, ref_angle_idx)
    canvas_origin_mm, canvas_shape = canvas_for_params(
        sras, ref_angle_idx, pitch, params)

    extra = extra or {}
    per_angle: dict[int, AngleTransform] = {}
    for a in range(sras.n_angles):
        p = params.get(a, ManualAngleParams())
        matrix, offset = build_canvas_affine(
            sras, a, ref_angle_idx, p.rotation_deg, p.shift_mm,
            pitch, canvas_origin_mm)
        score, source = extra.get(a, (1.0, ""))
        per_angle[a] = AngleTransform(p.rotation_deg, p.shift_mm, matrix, offset,
                                      score, source)

    return AlignmentResult(ref_angle_idx, dc_threshold_mv, canvas_shape,
                           pitch[0], pitch[1], canvas_origin_mm, per_angle)


def crop_alignment_result(result: AlignmentResult, row0: int, col0: int,
                          n_rows: int, n_cols: int) -> AlignmentResult:
    """The same alignment restricted to a rectangular window of its canvas —
    canvas pixel (row0, col0) becomes the cropped canvas's (0, 0).

    Cropping folds into each angle's existing affine instead of becoming a
    second transform, because a canvas crop is *pure index translation*. From
    _affine_out_to_src, matrix = D @ Rinv @ A_out depends only on pitch and
    rotation, and A_out @ [row0, col0] is exactly the mm displacement of the new
    origin, so:

        matrix @ [r', c'] + (offset + matrix @ [row0, col0])
            == matrix @ [r' + row0, c' + col0] + offset

    i.e. shifting the offset by matrix @ [row0, col0] reproduces the original
    mapping at the shifted indices, exactly. That equality is what lets
    apply_alignment, reproject_mask and the aligned .sras exporter all keep
    working on a cropped result with no special-casing — and it is why the crop
    preview a user approves is guaranteed to be the same pixels the exporter
    writes.

    Bounds are the caller's responsibility (the wizard's ROI page clamps to the
    canvas): an out-of-range window is geometrically well-defined here and
    simply resamples padding.
    """
    if n_rows <= 0 or n_cols <= 0:
        raise ValueError(f"empty crop: {n_rows} x {n_cols}")

    delta = np.array([float(row0), float(col0)])
    per_angle = {
        a: AngleTransform(t.rotation_deg, t.shift_mm, t.matrix,
                          t.offset + t.matrix @ delta, t.score, t.source)
        for a, t in result.per_angle.items()
    }
    origin = (result.canvas_origin_mm[0] + col0 * result.canvas_dx_mm,
              result.canvas_origin_mm[1] + row0 * result.canvas_dy_mm)
    return AlignmentResult(result.ref_angle_idx, result.dc_threshold_mv,
                           (int(n_rows), int(n_cols)),
                           result.canvas_dx_mm, result.canvas_dy_mm,
                           origin, per_angle)


def coarse_rect_to_canvas(rect: tuple[int, int, int, int],
                          downsample: tuple[int, int],
                          canvas_shape: tuple[int, int],
                          *, inset_blocks: int = 0
                          ) -> tuple[int, int, int, int]:
    """A rectangle in coarse (block-mean preview) indices as canvas pixels,
    both as (row0, col0, n_rows, n_cols) and clamped to the canvas.

    The single place the preview grid's relation to the real canvas is
    written down: the coarse grid samples canvas pixels 0, f, 2f, …, so a
    coarse pixel stands for a whole block and every caller has to agree on
    which canvas pixels that block means, or a crop the user approved on the
    preview lands a few pixels off in the export.

    *inset_blocks* shrinks the rectangle by that many coarse blocks on each
    side. A coarse pixel reported as fully covered stands for a block whose
    far edge may not be, so "fit to full overlap" insets by 1 to stay honestly
    inside the overlap region; a union rectangle insets by 0 because it wants
    to contain the region rather than fit inside it.
    """
    row0, col0, nr, nc = rect
    fy, fx = downsample
    n_rows, n_cols = canvas_shape
    r0 = min(n_rows - 1, (row0 + inset_blocks) * fy)
    c0 = min(n_cols - 1, (col0 + inset_blocks) * fx)
    r1 = min(n_rows, (row0 + nr - inset_blocks) * fy)
    c1 = min(n_cols, (col0 + nc - inset_blocks) * fx)
    return r0, c0, max(1, r1 - r0), max(1, c1 - c0)


def overlap_stats(counts: np.ndarray, n_angles: int) -> dict:
    """Summarize a per-pixel "how many angles cover this pixel" image — the
    number the alignment wizard's mask-stack view is colored by.

    Separated from the drawing code because it is the actual judgement the user
    makes on that screen ("do the angles land on top of each other?"), and a
    plain array-in/dict-out function can be tested without Qt.
    """
    counts = np.asarray(counts)
    union = int(np.count_nonzero(counts))
    full = int(np.count_nonzero(counts >= n_angles))
    return {
        "union_px": union,
        "full_px": full,
        "full_frac": (full / union) if union else 0.0,
        "mean_count": float(counts[counts > 0].mean()) if union else 0.0,
        "max_count": int(counts.max()) if counts.size else 0,
        "empty": union == 0,
    }


def largest_rect_at_least(counts: np.ndarray, min_count: int
                          ) -> tuple[int, int, int, int] | None:
    """Largest axis-aligned rectangle whose every pixel has counts >= min_count,
    as (row0, col0, n_rows, n_cols), or None if no pixel qualifies.

    Backs the wizard's "fit to full overlap" button. A *bounding box* of the
    qualifying pixels would be the obvious thing and is wrong: the full-overlap
    region of several rotated scans is roughly a disc, whose bounding box has
    corners no angle covers at all. Offering that as the crop would hand the
    user padding they explicitly asked to avoid, so this finds a rectangle that
    is entirely inside the region.

    Standard largest-rectangle-in-a-histogram sweep — O(rows * cols) on a
    preview-sized array, so it is instant at interactive rates.
    """
    good = np.asarray(counts) >= min_count
    if not good.any():
        return None

    n_rows, n_cols = good.shape
    best = (0, 0, 0, 0)          # area, row0, col0, ...
    best_area = 0
    heights = np.zeros(n_cols, dtype=np.int64)

    for r in range(n_rows):
        heights = np.where(good[r], heights + 1, 0)
        # Sentinel column of height 0 flushes the stack at the end of the row.
        stack: list[tuple[int, int]] = []       # (start col, height)
        for c in range(n_cols + 1):
            h = int(heights[c]) if c < n_cols else 0
            start = c
            while stack and stack[-1][1] >= h:
                s, sh = stack.pop()
                area = sh * (c - s)
                if area > best_area:
                    best_area = area
                    best = (s, sh, c - s, r)
                start = s
            if h:
                stack.append((start, h))

    col0, height, width, row_end = best
    return (row_end - height + 1, col0, height, width)


def _parallel_map(fn, items, n_workers: int) -> list:
    """fn over items, in order, threaded when it pays."""
    items = list(items)
    if n_workers <= 1 or len(items) <= 1:
        return [fn(x) for x in items]
    with ThreadPoolExecutor(max_workers=min(n_workers, len(items))) as pool:
        return list(pool.map(fn, items))


def compute_angle_alignment(sras: SrasFile, ref_angle_idx: int,
                            dc_threshold_mv: float, progress_cb=None,
                            fine_dim: int = _DEFAULT_FINE_DIM) -> AlignmentResult:
    """Top-level alignment driver: register every angle onto *ref_angle_idx* by
    content, then lay them all out on that angle's own coordinate grid.

    Meant for a background thread — deliberately
    recomputes CH4 DC images from scratch rather than reading the GUI-thread
    _dc_cache dict, since background-thread workers must not touch
    GUI-thread-owned caches.
    """
    n = sras.n_angles   # already the *complete*-angle count for aborted v6 scans
    # Parallel over angles, serial within each — see plan_angle_level.
    n_workers, angle_budget = plan_angle_level(sras)

    # progress_cb fires from pool threads, so the counter behind it must be
    # atomic. list.append is, and len() of a list is a consistent read.
    _ticks: list[int] = []

    def tick(base: int, span: int):
        if progress_cb:
            _ticks.append(1)
            progress_cb(base + int(min(len(_ticks), n) / n * span))

    # ---- Step 1: CH4 DC image per angle, native grid ----------------------
    def dc4_for(a: int) -> np.ndarray:
        dc4 = adc_to_mv(
            compute_dc_image(sras, a, CH4_IDX, max_workers=1, budget=angle_budget),
            *sras.cal(CH4_IDX))
        tick(0, 25)
        return dc4

    dc4_mv = dict(enumerate(_parallel_map(dc4_for, range(n), n_workers)))

    # ---- Step 2: rigid registration of every angle against the reference --
    # Its own worker count: registration is bounded by grid-sized FFT buffers,
    # not by the waveform chunking plan_angle_level budgets for.
    _ticks.clear()

    def fit_for(a: int) -> RigidFit:
        fit = register_angle_to_reference(
            sras, a, ref_angle_idx, dc4_mv, dc_threshold_mv=dc_threshold_mv,
            fine_dim=fine_dim)
        tick(25, 65)
        return fit

    fits = dict(enumerate(_parallel_map(
        fit_for, range(n), _registration_workers(sras, fine_dim))))

    # ---- Step 3/4: shared canvas on the reference's grid, per-angle affines
    params = {a: ManualAngleParams(f.rotation_deg, f.shift_mm) for a, f in fits.items()}
    extra = {a: (f.score, f.source) for a, f in fits.items()}
    result = _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params, extra)
    if progress_cb:
        progress_cb(100)
    return result


# ---------------------------------------------------------------------------
# Manual alignment (Fusion menu -> Manual Alignment... dialog)
#
# Skips register_angle_to_reference's search entirely: every angle's
# rotation_deg/shift_mm is supplied directly by the caller (nudged by eye
# against a live multi-angle mask overlay, or pre-seeded from a registration
# run or a saved sidecar). Building the final AlignmentResult from already-known
# per-angle parameters is pure closed-form matrix math (_result_from_params)
# with no per-pixel image work at all, so build_manual_alignment is cheap enough
# to call synchronously on the GUI thread on every edit. The only genuinely
# expensive per-pixel operation anywhere in this flow is reproject_mask, and
# only the wizard's own downsampled mask stack calls that per keystroke
# — see AlignmentWizard.rebuild_stack for how it limits a nudge to reprojecting
# only the actively-edited angle.
# ---------------------------------------------------------------------------

def reproject_mask(sras: SrasFile, angle_idx: int, ref_angle_idx: int,
                   mask: np.ndarray, rotation_deg: float,
                   shift_mm: tuple[float, float],
                   canvas_pitch_mm: tuple[float, float],
                   canvas_origin_mm: tuple[float, float],
                   canvas_shape: tuple[int, int],
                   *, src_downsample: tuple[int, int] = (1, 1)) -> np.ndarray:
    """Resample one angle's binary/float mask onto an arbitrary canvas via an
    explicit rotation+shift — the single building block the alignment wizard's
    live mask stack repeatedly calls (once per angle per edit). *src_downsample* must match the (rows, cols)
    block-mean factor already applied to *mask*, or the reprojection lands at
    the wrong scale. order=0 (nearest) matches apply_alignment's own reasoning:
    a binary mask must never be blended with zero-padding."""
    matrix, offset = build_canvas_affine(
        sras, angle_idx, ref_angle_idx, rotation_deg, shift_mm,
        canvas_pitch_mm, canvas_origin_mm, src_downsample=src_downsample)
    return scipy_ndimage.affine_transform(
        mask.astype(np.float32, copy=False), matrix, offset=offset,
        output_shape=canvas_shape, order=0, mode="constant", cval=0.0)


def build_manual_alignment(sras: SrasFile, ref_angle_idx: int,
                           dc_threshold_mv: float,
                           per_angle_params: dict[int, ManualAngleParams]
                           ) -> AlignmentResult:
    """Build a full, full-resolution AlignmentResult from user-supplied
    per-angle rotation+shift — the Manual Alignment counterpart to
    compute_angle_alignment, skipping its registration search entirely (every
    angle's transform here is exactly what the caller supplied). The reference
    angle's params are always forced to identity, regardless of what
    per_angle_params holds for it — it defines the shared origin and must never
    be transformed.

    dc_threshold_mv plays no part in the geometry; it's stored on the returned
    AlignmentResult purely as a record of the RF-mask threshold in effect at
    the time.
    """
    params = {a: per_angle_params.get(a, ManualAngleParams())
              for a in range(sras.n_angles)}
    params[ref_angle_idx] = ManualAngleParams()
    return _result_from_params(sras, ref_angle_idx, dc_threshold_mv, params)


# ---- Sidecar persistence (<name>.sras.align.json) -------------------------
# A viewer-computed derived artifact, so it lives with the alignment math
# rather than in sras_format (see docs/design.md, "Manual-alignment sidecar").

@dataclass
class ManualAlignmentSidecar:
    ref_angle_idx: int
    dc_threshold_mv: float
    per_angle: dict[int, ManualAngleParams]


def sidecar_path(sras_path) -> Path:
    """<name>.sras.align.json next to the scan file. A thin, independently
    testable function since save/load/delete and the GUI's status messages
    all need the identical path."""
    p = Path(sras_path)
    return p.with_name(p.name + ".align.json")


# Bumped whenever the frame the stored numbers are measured in changes; older
# sidecars are treated as absent, never migrated. Bump history:
# docs/design.md ("Schema history").
_SIDECAR_SCHEMA_VERSION = 3


def save_manual_alignment(sras: SrasFile, ref_angle_idx: int,
                          dc_threshold_mv: float,
                          per_angle: dict[int, ManualAngleParams]) -> Path:
    """Write the sidecar JSON for sras.path (overwriting any existing one)
    and return the path written. Angle indices become JSON object keys, so
    they round-trip as strings — load_manual_alignment converts them back."""
    path = sidecar_path(sras.path)
    payload = {
        "schema_version": _SIDECAR_SCHEMA_VERSION,
        "ref_angle_idx": ref_angle_idx,
        "dc_threshold_mv": dc_threshold_mv,
        "per_angle": {
            str(a): {"rotation_deg": p.rotation_deg, "shift_mm": list(p.shift_mm)}
            for a, p in per_angle.items()
        },
    }
    path.write_text(json.dumps(payload, indent=2))
    return path


def load_manual_alignment(sras: SrasFile) -> ManualAlignmentSidecar | None:
    """Read <sras.path>'s sidecar JSON if present, else None — never raises:
    a missing, corrupt, foreign, or version-mismatched JSON file must not
    block opening the .sras file itself (see SrasViewerWindow._on_load_done).
    Per-angle entries for an angle index no longer present in *sras* (e.g.
    re-scanned with fewer angles) are silently dropped.
    """
    path = sidecar_path(sras.path)
    if not path.exists():
        return None
    try:
        raw = json.loads(path.read_text())
        if raw.get("schema_version") != _SIDECAR_SCHEMA_VERSION:
            return None
        per_angle = {
            int(a): ManualAngleParams(
                float(v["rotation_deg"]),
                (float(v["shift_mm"][0]), float(v["shift_mm"][1])))
            for a, v in raw.get("per_angle", {}).items()
            if int(a) < sras.n_angles
        }
        return ManualAlignmentSidecar(
            ref_angle_idx=int(raw.get("ref_angle_idx", 0)),
            dc_threshold_mv=float(raw.get("dc_threshold_mv", 0.0)),
            per_angle=per_angle)
    except (OSError, ValueError, KeyError, TypeError, IndexError,
            json.JSONDecodeError):
        return None


def delete_manual_alignment(sras: SrasFile) -> bool:
    """Delete the sidecar if present. Returns whether a file actually existed
    to delete, so Clear Alignment's status message can say so. Genuine I/O
    errors (permission denied, read-only share) propagate — the caller
    (the wizard's Clear path) surfaces them rather than silently
    pretending the destructive action succeeded."""
    path = sidecar_path(sras.path)
    try:
        path.unlink()
        return True
    except FileNotFoundError:
        return False