av-denoise 0.3.1

use std::marker::PhantomData;

use cubecl::prelude::*;
use cubecl::server::Handle;

use super::align::StorageAlign;
use super::kernels::gpu_copy;
use super::motion::{self, MotionCtx, MotionEstimation, build_pyramid_for_slot, run_pyramid_build};
use super::noise::{
    EMA_ALPHA,
    NoiseCtx,
    NoiseEstimator,
    TemporalNoiseSample,
    TemporalStatsCtx,
    aggregate_temporal_noise_stats,
    build_spatial_offset_lut,
    correlation_factor,
    noise_partials_slot_stride_bytes,
    partials_len,
    read_temporal_stats_slot,
    run_noise_estimate,
    run_temporal_noise_stats,
    sigma_block_p25_from_partials,
    sigma_from_abs_sum,
    temporal_stats_buf_bytes,
    zero_temporal_stats_slot,
};
use super::params::{NlmParams, SEPARABLE_THRESHOLD, sigma_eff, validate_dimensions};
use super::pending::Pending;
use super::prefilter::{PrefilterCtx, PrefilterMode, run_prefilter};
use super::{BLOCK_1D, MAX_GRID_1D};

/// Stateful NLMeans denoiser. Maintains a ring of frames in
/// `input_buf`; each `push_frame` uploads one frame, each `denoise`
/// processes the current center frame using its temporal neighbourhood.
pub struct NlmDenoiser<R: Runtime> {
    pub(super) client: ComputeClient<R>,
    pub(super) params: NlmParams,
    pub(super) width: u32,
    pub(super) height: u32,
    /// Byte alignment every per-slot buffer view must start on, read
    /// from `client`'s runtime at construction. See [`StorageAlign`].
    pub(super) align: StorageAlign,

    /// Monotonic count of frames pushed; `% total_frames` is the next
    /// physical slot in `input_buf` to overwrite.
    pub(super) ring_head: usize,
    /// Frames loaded so far, capped at `total_frames`.
    pub(super) frames_loaded: usize,
    /// Number of real `push_frame` / `push_frame_with_reference` calls
    /// in the current stream (does not count internal leading/trailing
    /// duplicates). Reset by [`Self::reset_stream_state`].
    pub(super) real_pushes: usize,

    /// `[total_frames * height * width * stored_ch]` ring buffer.
    pub(super) input_buf: Handle,
    /// Reference ring buffer with the same shape as `input_buf`. Used
    /// only when `params.prefilter != None`; supplies the distance
    /// signal for the `_ref` kernel variants.
    pub(super) reference_buf: Option<Handle>,
    /// CPU scratch for YUV-→4-lane repacking. Empty when no padding needed.
    pub(super) padding_scratch: Vec<f32>,
    /// `[pixels * stored_ch]` weighted-pixel accumulator.
    pub(super) accum: Handle,
    /// `[pixels]` total weight per pixel.
    pub(super) weight_sum: Handle,
    /// `[pixels]` max neighbour weight per pixel.
    pub(super) max_weight: Handle,
    /// `[pixels]` weight scratch used by the symmetric (k=0) path.
    pub(super) weight_buf: Handle,
    /// `[pixels]` raw fwd distance (separable path).
    pub(super) raw_fwd: Handle,
    /// `[pixels]` raw bwd distance (separable path).
    pub(super) raw_bwd: Handle,
    /// `[pixels]` hsum intermediate, fwd direction (separable path).
    pub(super) tmp_hsum: Handle,
    /// `[pixels]` hsum intermediate, bwd direction (separable path).
    pub(super) tmp_hsum_bwd: Handle,
    /// Double-buffered `[pixels * stored_ch]` denoised output. A new
    /// `denoise_submit()` writes into `outputs[next_output_slot]` while
    /// the previous slot may still be draining via `read_async`, letting
    /// frame N+1's kernels overlap with frame N's readback.
    pub(super) outputs: [Handle; 2],
    /// Index of the next output slot to write into.
    pub(super) next_output_slot: usize,
    /// CPU scratch reused by the sync `denoise()` path via
    /// `Pending::wait_into`. Avoids a per-frame allocation.
    pub(super) output_scratch: Vec<f32>,

    pub(super) h2_inv_norm: f32,
    /// Distance floor fed to the main-pass weighting kernels. Zero for
    /// `NlmSpatial`, because pilot-vs-pilot distances no longer carry
    /// the 2σ² floor, so subtracting it there would overweight
    /// mismatched patches. Equal to `input_noise_offset` for every
    /// other prefilter mode.
    pub(super) noise_offset: f32,
    /// Distance floor for comparisons against noisy input pixels. The
    /// pilot pass always uses this value, since its own inputs still
    /// carry the full noise floor even when `noise_offset` has been
    /// zeroed for the main pass.
    pub(super) input_noise_offset: f32,
    pub use_separable: bool,
    pub(super) use_reference: bool,

    /// Smoothed grain autocorrelation, same EMA cadence as the sigma
    /// estimate. `None` before the first temporal sample of a stream,
    /// same seeding convention as `NoiseEstimator`. Its first sample
    /// sets the state directly instead of blending from an assumed 0,
    /// so a stream doesn't spend its opening frames under-attenuated.
    /// Only ever updated at fold time when a temporal sample exists,
    /// so the fast path and a fixed `sigma_override` leave it at `None`
    /// (read as 0, white-noise attenuation, i.e. no attenuation) for
    /// the whole stream.
    pub(super) rho_smoothed: Option<f32>,
    /// Per-candidate spatial noise-floor offset for the k=0 windowed
    /// and separable weighting kernels, `[(2*search_radius+1)^2]`
    /// f32s, row-major `(dy+r)*(2r+1)+(dx+r)`. Rebuilt from
    /// `noise_offset` and `rho_smoothed` every `denoise_submit` (see
    /// `Self::rebuild_spatial_offset_lut`). At `rho_smoothed` unset
    /// this is numerically identical to the flat `noise_offset` scalar
    /// it replaced.
    pub(super) spatial_offset_lut: Handle,

    /// Stage-1 noise-estimate scratch ring, `total_frames` slots each
    /// [`noise_partials_slot_stride_bytes`] bytes wide. `Some`
    /// only when the noise level is measured automatically (`hq` is
    /// set and `sigma_override` is `None`). One slot per ring position
    /// keeps a slot's partials intact between the push that queues
    /// them and the later fold that reads the centre slot back for the
    /// low chain's block-p25 statistic, unlike a single shared region
    /// every push would overwrite before that slot ever reaches centre.
    pub(super) noise_partials: Option<Handle>,
    /// Per-ring-slot Immerkær totals, `[total_frames * 4]` f32s. Same
    /// gating as `noise_partials`.
    pub(super) noise_results: Option<Handle>,
    /// Per-ring-slot temporal-residual block stats. One region per
    /// ring slot, each region `blocks_x * blocks_y` block records
    /// (row-major), each record `[sum_d(ch0..stored_ch-1),
    /// sum_d2(ch0..stored_ch-1), sum_lag]`. `Some` only when auto
    /// noise estimation is active (same as `noise_partials`) and
    /// `temporal_radius >= 1` (the r=0 path has no temporal neighbour
    /// to diff against).
    pub(super) temporal_stats_buf: Option<Handle>,
    /// Smooths the median chain's raw per-frame noise estimate into a
    /// stable per-channel sigma, feeding `h2_inv_norm` and `sigma_y`.
    /// Inert (never updated) when noise is not measured automatically.
    pub(super) noise_estimator: NoiseEstimator,
    /// Smooths the low chain's raw per-frame noise estimate into a
    /// stable per-channel sigma, feeding `input_noise_offset` and
    /// `noise_offset` only. The low chain uses block-p25 Immerkær
    /// maxed with the temporal lower quartile rather than the median
    /// chain's frame-mean Immerkær maxed with the temporal median, a
    /// more conservative read for a consumer where an over-read is
    /// destructive. Inert under the same condition as `noise_estimator`.
    pub(super) noise_estimator_low: NoiseEstimator,

    /// Cached motion-compensation context. `Some` when MC is active.
    pub(super) mc_ctx: Option<MotionCtx>,
    /// `[total_frames * height * width * stored_ch]` warped input ring,
    /// matching `input_buf`. Temporal (k≠0) kernels read neighbours
    /// from here; the centre slot is a straight copy of `input_buf`.
    pub(super) compensated_input_buf: Option<Handle>,
    /// Same shape as `compensated_input_buf`, mirroring the reference
    /// ring when a prefilter is active.
    pub(super) compensated_reference_buf: Option<Handle>,
    /// Per-neighbour MV field. Layout:
    /// `[2·temporal_radius][blocks_y * blocks_x * 2]` `i32`. Neighbour
    /// indices `0..R` are the backward k = -R..-1; `R..2R` are forward
    /// k = +1..+R.
    pub(super) mv_field_buf: Option<Handle>,
    /// Adjacent-frame pair-motion ring, laid out
    /// `[2·temporal_radius][2][blocks_y][blocks_x][2]` `i32` (outer
    /// index is the pair slot, next is direction, 0 = older→newer, 1 =
    /// newer→older). `Some` only when `MotionEstimation::Chained` is
    /// active. The direct path never touches this buffer. See
    /// `motion::pair_ring_slot_count` for why `2·temporal_radius` slots
    /// is exactly enough, and `Self::pair_slot` for how a slot is
    /// resolved from a frame's position in the push sequence.
    pub(super) pair_ring_buf: Option<Handle>,
    /// Luma-only pyramid storage:
    /// `[pyramid_levels][total_frames][level_w * level_h]` `f32`.
    pub(super) pyramid_input: Option<Handle>,
    /// Same shape as `pyramid_input`, built from the reference ring
    /// when a prefilter is active.
    pub(super) pyramid_reference: Option<Handle>,

    /// Block geometry for the no-MC confidence pass. `Some` only when
    /// confidence weighting is active (HQ, `temporal_confidence: true`,
    /// `temporal_radius > 0`) and motion compensation is not. The
    /// MC-active case reuses `mc_ctx`'s geometry instead.
    pub(super) confidence_ctx: Option<MotionCtx>,
    /// Per-neighbour block-match confidence. Layout mirrors
    /// `mv_field_buf`, `[2·temporal_radius][blocks_y * blocks_x]`
    /// `f32`. `Some` only when confidence weighting is active (see
    /// `confidence_ctx`), whether the block geometry comes from
    /// `mc_ctx` or from the no-MC confidence pass.
    pub(super) confidence_buf: Option<Handle>,
    /// Luma-only single-level pyramid ring feeding the no-MC
    /// confidence pass. `Some` only alongside `confidence_ctx`.
    pub(super) confidence_pyramid: Option<Handle>,
    /// Discard sink for the no-MC confidence pass's mandatory MV
    /// write. Nothing warps by it without motion compensation. `Some`
    /// only alongside `confidence_ctx`.
    pub(super) confidence_mv_scratch: Option<Handle>,
    /// Small placeholder buffer passed as the fine block-match
    /// kernel's `confidence` argument whenever confidence weighting is
    /// inactive but motion compensation still runs. The kernel's
    /// `write_confidence` comptime flag skips indexing into it
    /// entirely in that case, so its size never matters. Always
    /// allocated (trivially small), unlike the confidence-specific
    /// buffers above.
    pub(super) confidence_dummy: Handle,
    /// Smoothed sigma for channel 0 (the plane motion estimation
    /// treats as luma), feeding the confidence noise floor. Zero
    /// unless HQ is active. A fixed `sigma_override` seeds it once at
    /// construction, auto estimation refreshes it every submit.
    pub(super) sigma_y: f32,
}

impl<R: Runtime> NlmDenoiser<R> {
    /// Build a new denoiser.
    ///
    /// **Panics** if `params.validate()` or the frame-dimension check
    /// fails, the high-level [`crate::Denoiser`] runs both first and
    /// surfaces errors as `Result`, so most callers should prefer that.
    pub fn new(client: &ComputeClient<R>, params: NlmParams, width: u32, height: u32) -> Self {
        params
            .validate()
            .expect("invalid NlmParams; call params.validate() first to surface this as Result");
        validate_dimensions(width, height)
            .expect("unsupported frame dimensions; call validate_dimensions first to surface this as Result");

        let align = StorageAlign::from_client(client);
        let stored_ch = params.channels.storage_count();
        let total_frames = params.total_frames();
        let pixels = (width * height) as usize;
        let frame_bytes = pixels * stored_ch as usize * size_of::<f32>();
        let scalar_bytes = pixels * size_of::<f32>();

        let input_buf = client.empty(frame_bytes * total_frames as usize);
        let reference_buf = if params.prefilter.needs_reference_buf() {
            Some(client.empty(frame_bytes * total_frames as usize))
        } else {
            None
        };

        let padding_scratch = if params.channels.count() != stored_ch {
            vec![0.0f32; pixels * stored_ch as usize]
        } else {
            Vec::new()
        };

        let accum = client.empty(frame_bytes);
        let weight_sum = client.empty(scalar_bytes);
        let max_weight = client.empty(scalar_bytes);
        let weight_buf = client.empty(scalar_bytes);
        let raw_fwd = client.empty(scalar_bytes);
        let raw_bwd = client.empty(scalar_bytes);
        let tmp_hsum = client.empty(scalar_bytes);
        let tmp_hsum_bwd = client.empty(scalar_bytes);
        let outputs = [client.empty(frame_bytes), client.empty(frame_bytes)];

        let h2_inv_norm = params.h2_inv_norm();
        let input_noise_offset = params.noise_offset();
        // The pilot pass compares noisy input pixels, so it always
        // keeps the full noise floor. Main-pass distances under
        // `NlmSpatial` are pilot-vs-pilot, which no longer carries
        // that floor, so subtracting it there would overweight
        // mismatched patches.
        let noise_offset = match params.prefilter {
            PrefilterMode::NlmSpatial { .. } => 0.0,
            _ => input_noise_offset,
        };
        let use_separable = params.patch_radius > SEPARABLE_THRESHOLD;
        let use_reference = params.prefilter.needs_reference_buf();
        let output_scratch_cap = pixels * params.channels.count() as usize;

        // Unset until the first temporal sample lands, so the initial
        // LUT is numerically identical to the flat `noise_offset`
        // scalar it replaces.
        let rho_smoothed: Option<f32> = None;
        let spatial_offset_lut = client.create_from_slice(f32::as_bytes(&build_spatial_offset_lut(
            params.search_radius,
            0.0,
            noise_offset,
        )));

        // Auto noise estimation only runs when HQ is on and the caller
        // hasn't pinned a fixed sigma. The fast path and the
        // sigma-override path allocate neither buffer nor ever launch
        // the estimate kernels.
        let auto_noise = params.hq.is_some_and(|hq| hq.sigma_override.is_none());
        let (noise_partials, noise_results) = if auto_noise {
            let partials_ring_bytes =
                noise_partials_slot_stride_bytes(width, height, align) * total_frames as u64;
            let n_results = (total_frames * 4) as usize;
            (
                Some(client.empty(partials_ring_bytes as usize)),
                Some(client.empty(n_results * size_of::<f32>())),
            )
        } else {
            (None, None)
        };

        // The temporal residual estimator additionally needs a real
        // temporal neighbour to diff against, so it stays inert at
        // temporal_radius = 0 even when auto noise estimation is on.
        let temporal_stats_buf = if auto_noise && params.temporal_radius >= 1 {
            Some(client.empty(temporal_stats_buf_bytes(
                width,
                height,
                stored_ch,
                total_frames,
                align,
            )))
        } else {
            None
        };

        // Motion-compensation buffers. Only allocated when MC is
        // active *and* the temporal window is non-trivial (k=0 path
        // would never touch them).
        let mc_ctx = if params.motion_compensation.is_active() && params.temporal_radius > 0 {
            MotionCtx::new(params.motion_compensation, width, height, align)
        } else {
            None
        };

        let (
            compensated_input_buf,
            compensated_reference_buf,
            mv_field_buf,
            pyramid_input,
            pyramid_reference,
        ) = if let Some(ctx) = mc_ctx.as_ref() {
            let comp_in = client.empty(frame_bytes * total_frames as usize);
            let comp_ref = if use_reference {
                Some(client.empty(frame_bytes * total_frames as usize))
            } else {
                None
            };
            let neighbours = (2 * params.temporal_radius) as u64;
            let mv_field = client.empty((neighbours * ctx.mv_field_bytes_per_neighbour()) as usize);
            let pyramid_pixels =
                motion::pyramid_pixels_per_frame(width, height, ctx.pyramid_levels, ctx.align);
            let pyr_in_bytes = pyramid_pixels * total_frames as usize * size_of::<f32>();
            let pyr_in = client.empty(pyr_in_bytes);
            let pyr_ref = if use_reference {
                Some(client.empty(pyr_in_bytes))
            } else {
                None
            };
            (Some(comp_in), comp_ref, Some(mv_field), Some(pyr_in), pyr_ref)
        } else {
            (None, None, None, None, None)
        };

        // The pair ring is allocated only when `Chained` estimation is
        // active (explicitly, or via `Auto` resolving to it at this
        // temporal radius), on top of `mc_ctx` already being `Some`.
        // The direct path never reads or writes it.
        let is_chained = matches!(
            params
                .motion_compensation
                .resolved_estimation(params.temporal_radius),
            Some(MotionEstimation::Chained { .. })
        );
        let pair_ring_buf = if is_chained {
            mc_ctx.as_ref().map(|ctx| {
                let pair_ring_slots = motion::pair_ring_slot_count(params.temporal_radius) as u64;
                let bytes = pair_ring_slots * ctx.pair_slot_bytes();
                client.empty(bytes as usize)
            })
        } else {
            None
        };

        // Confidence weighting (in either its MC-active or its no-MC
        // form) is active only when HQ has `temporal_confidence: true`
        // and the temporal window is non-trivial. This gate applies
        // even when MC is active. Without it, every MC-active submit
        // would pay for the fine kernel's confidence write whether or
        // not anything consumes it.
        let confidence_active =
            params.hq.is_some_and(|hq| hq.temporal_confidence) && params.temporal_radius > 0;

        // Confidence-only geometry, only needed when MC isn't already
        // supplying block geometry (and thus MVs and confidence via
        // its own analyse pass). This incurs real extra work beyond
        // the MC-active case. It needs its own luma pyramid ring and a
        // block-match kernel per neighbour.
        let confidence_only_active = confidence_active && mc_ctx.is_none();
        let confidence_ctx = confidence_only_active.then(|| MotionCtx::confidence_only(width, height, align));

        // The confidence buffer piggybacks on whichever block geometry
        // is available, but only when confidence weighting is active.
        let confidence_geometry = if confidence_active {
            mc_ctx.as_ref().or(confidence_ctx.as_ref())
        } else {
            None
        };
        let confidence_buf = confidence_geometry.map(|ctx| {
            let neighbours = (2 * params.temporal_radius) as u64;
            client.empty((neighbours * ctx.confidence_bytes_per_neighbour()) as usize)
        });
        // Always allocated, trivially small, and reused whenever the
        // fine block-match kernel runs with `write_confidence: false`.
        let confidence_dummy = client.empty(size_of::<f32>());

        let (confidence_pyramid, confidence_mv_scratch) = if let Some(ctx) = confidence_ctx.as_ref() {
            let pyramid_pixels =
                motion::pyramid_pixels_per_frame(width, height, ctx.pyramid_levels, ctx.align);
            let pyr_bytes = pyramid_pixels * total_frames as usize * size_of::<f32>();
            let mv_scratch_len = ctx.mv_slots_per_neighbour() * 2 * size_of::<i32>();
            (Some(client.empty(pyr_bytes)), Some(client.empty(mv_scratch_len)))
        } else {
            (None, None)
        };

        // `sigma_override` is the only source before the first noise
        // estimate lands. Auto estimation refreshes this every submit
        // (see `update_noise_estimate`). The fast path (`hq: None`)
        // leaves it at zero, which `motion::sad_noise_floor` turns into
        // a zero floor exactly as callers with no estimate should get.
        let sigma_y = params.hq.and_then(|hq| hq.sigma_override).unwrap_or(0.0);

        Self {
            client: client.clone(),
            params,
            width,
            height,
            align,
            ring_head: 0,
            frames_loaded: 0,
            real_pushes: 0,
            input_buf,
            reference_buf,
            padding_scratch,
            accum,
            weight_sum,
            max_weight,
            weight_buf,
            raw_fwd,
            raw_bwd,
            tmp_hsum,
            tmp_hsum_bwd,
            outputs,
            next_output_slot: 0,
            output_scratch: Vec::with_capacity(output_scratch_cap),
            h2_inv_norm,
            noise_offset,
            input_noise_offset,
            use_separable,
            use_reference,
            rho_smoothed,
            spatial_offset_lut,
            noise_partials,
            noise_results,
            temporal_stats_buf,
            noise_estimator: NoiseEstimator::default(),
            noise_estimator_low: NoiseEstimator::default(),
            mc_ctx,
            compensated_input_buf,
            compensated_reference_buf,
            mv_field_buf,
            pair_ring_buf,
            pyramid_input,
            pyramid_reference,
            confidence_ctx,
            confidence_buf,
            confidence_pyramid,
            confidence_mv_scratch,
            confidence_dummy,
            sigma_y,
        }
    }

    /// Push a new frame into the ring buffer. `frame` must hold
    /// `width * height * channels` f32 values normalised to [0, 1].
    /// YUV padding (3→4 lanes) is repacked through a reused CPU scratch.
    ///
    /// For `PrefilterMode::External` use
    /// [`Self::push_frame_with_reference`] instead.
    pub fn push_frame(&mut self, frame: &[f32]) {
        assert!(
            !matches!(self.params.prefilter, PrefilterMode::External),
            "push_frame_with_reference is required when prefilter == External"
        );

        let slot = self.upload_into(&self.input_buf.clone(), frame);

        self.run_noise_estimate_for_slot(slot as u32);
        self.run_temporal_stats_for_slot(slot as u32);
        self.seed_noise_estimate_if_first_frame(slot as u32);

        if let PrefilterMode::NlmSpatial { strength_scale } = self.params.prefilter {
            self.run_nlm_spatial_pilot(slot as u32, strength_scale)
                .expect("nlm spatial pilot dispatch failed");
        } else if self.params.prefilter.is_gpu_internal() {
            self.run_prefilter_for_slot(slot);
        }

        self.build_pyramids_for_slot(slot as u32);
        self.build_confidence_pyramid_for_slot(slot as u32);
        self.run_pair_analyse_for_slot(slot as u32);

        self.advance_ring();
        self.prime_leading_edge_if_first();
    }

    /// Push a new frame together with an externally-prefiltered
    /// reference. Required when `prefilter == External`; both slices
    /// must hold `width * height * channels` f32 values in [0, 1].
    pub fn push_frame_with_reference(&mut self, frame: &[f32], reference: &[f32]) {
        assert!(
            matches!(self.params.prefilter, PrefilterMode::External),
            "push_frame_with_reference requires prefilter == External"
        );

        let slot = self.upload_into(&self.input_buf.clone(), frame);
        let reference_buf = self
            .reference_buf
            .as_ref()
            .expect("reference buffer must exist for External prefilter")
            .clone();
        self.upload_into_slot(&reference_buf, reference, slot);

        // Same order as push_frame. The noise estimate and its
        // first-frame seed run before anything that could read σ
        // (build_pyramids_for_slot only needs the reference upload
        // just above, not the noise estimate, so this reordering
        // doesn't change what either step reads).
        self.run_noise_estimate_for_slot(slot as u32);
        self.run_temporal_stats_for_slot(slot as u32);
        self.seed_noise_estimate_if_first_frame(slot as u32);

        self.build_pyramids_for_slot(slot as u32);
        self.build_confidence_pyramid_for_slot(slot as u32);
        self.run_pair_analyse_for_slot(slot as u32);

        self.advance_ring();
        self.prime_leading_edge_if_first();
    }

    /// Upload `frame` into the next ring slot of `dst`. Returns the
    /// physical slot index written.
    fn upload_into(&mut self, dst: &Handle, frame: &[f32]) -> usize {
        let total_frames = self.params.total_frames() as usize;
        let slot = self.ring_head % total_frames;
        self.upload_into_slot(dst, frame, slot);
        slot
    }

    fn upload_into_slot(&mut self, dst: &Handle, frame: &[f32], slot: usize) {
        let channels = self.params.channels.count() as usize;
        let stored_ch = self.params.channels.storage_count() as usize;
        let pixels = self.width as usize * self.height as usize;
        let expected = pixels * channels;

        assert_eq!(
            frame.len(),
            expected,
            "frame size mismatch: expected {expected}, got {}",
            frame.len()
        );

        let staging = if channels == stored_ch {
            self.client.create_from_slice(f32::as_bytes(frame))
        } else {
            for i in 0..pixels {
                let dst_off = i * stored_ch;
                let src_off = i * channels;
                self.padding_scratch[dst_off..dst_off + channels]
                    .copy_from_slice(&frame[src_off..src_off + channels]);
            }
            self.client
                .create_from_slice(f32::as_bytes(&self.padding_scratch))
        };

        self.copy_frame_into_slot(dst, slot, &staging, 0, 1);
    }

    fn run_prefilter_for_slot(&self, slot: usize) {
        let reference_buf = self
            .reference_buf
            .as_ref()
            .expect("reference buffer must exist for GPU prefilter");

        let ctx = PrefilterCtx {
            width: self.width,
            height: self.height,
            channels: self.params.channels.count(),
            stored_ch: self.params.channels.storage_count(),
            frame_count: self.params.total_frames(),
            frame: slot as u32,
            input_buf: &self.input_buf,
            reference_buf,
        };

        run_prefilter::<R>(self.params.prefilter, &self.client, &ctx).expect("prefilter dispatch failed");
    }

    /// Build the per-frame motion-estimation pyramid for `slot` on
    /// both the input and (when present) the reference rings. No-op
    /// when MC is disabled.
    fn build_pyramids_for_slot(&self, slot: u32) {
        let Some(ctx) = self.mc_ctx.as_ref() else {
            return;
        };

        let stored_ch = self.params.channels.storage_count();
        let frame_count = self.params.total_frames();

        if let Some(pyr) = self.pyramid_input.as_ref() {
            build_pyramid_for_slot::<R>(
                &self.client,
                ctx,
                self.width,
                self.height,
                frame_count,
                slot,
                &self.input_buf,
                pyr,
                stored_ch,
            )
            .expect("input pyramid build dispatch failed");
        }

        if let (Some(pyr_ref), Some(ref_buf)) = (self.pyramid_reference.as_ref(), self.reference_buf.as_ref())
        {
            build_pyramid_for_slot::<R>(
                &self.client,
                ctx,
                self.width,
                self.height,
                frame_count,
                slot,
                ref_buf,
                pyr_ref,
                stored_ch,
            )
            .expect("reference pyramid build dispatch failed");
        }
    }

    /// Extract the level-0 luma plane for `slot` into the no-MC
    /// confidence pyramid. No-op unless the no-MC confidence pass is
    /// active. Always reads `input_buf`, even under a prefilter. The
    /// no-MC path keeps confidence simple by comparing raw input
    /// rather than duplicating the reference ring's pyramid.
    ///
    /// Calls `run_pyramid_build` directly rather than going through
    /// [`Self::build_pyramids_for_slot`]'s `build_pyramid_for_slot`
    /// helper. That helper only ever touches `mc_ctx`'s own pyramid
    /// buffers (`pyramid_input`, `pyramid_reference`), and
    /// `confidence_ctx` is only `Some` when `mc_ctx` is `None`, so it
    /// would return immediately without ever building this method's
    /// own `confidence_pyramid`.
    fn build_confidence_pyramid_for_slot(&self, slot: u32) {
        let (Some(ctx), Some(pyr)) = (self.confidence_ctx.as_ref(), self.confidence_pyramid.as_ref()) else {
            return;
        };

        run_pyramid_build::<R>(
            &self.client,
            ctx,
            self.width,
            self.height,
            self.params.total_frames(),
            slot,
            &self.input_buf,
            pyr,
            self.params.channels.storage_count(),
        )
        .expect("confidence pyramid build dispatch failed");
    }

    /// Whether `Chained` motion estimation is configured, explicitly or
    /// via `Auto` resolving to it at this denoiser's temporal radius
    /// (see `resolved_estimation`, the single source every
    /// estimation-dependent decision goes through). Orthogonal to
    /// `mc_ctx.is_some()`, which callers still need to check
    /// separately, since `mc_ctx` also requires `temporal_radius > 0`.
    pub(super) fn is_chained(&self) -> bool {
        matches!(
            self.params
                .motion_compensation
                .resolved_estimation(self.params.temporal_radius),
            Some(MotionEstimation::Chained { .. })
        )
    }

    /// Run the adjacent-frame pair analyse for the physical input-ring
    /// slot just written by `push_frame`/`push_frame_with_reference`,
    /// storing both directions' motion fields into the pair ring at
    /// `Self::pair_slot(0)`. No-op unless `Chained` estimation is
    /// active, and for the very first frame of a stream (`ring_head ==
    /// 0`), which has no older partner to pair against. Composition
    /// for that gap instead reads the priming duplicate's zero-filled
    /// pair (see [`Self::zero_pair_slot_for_duplicate`]).
    fn run_pair_analyse_for_slot(&self, newer_slot: u32) {
        if self.ring_head == 0 {
            return;
        }
        let Some(mc) = self.mc_ctx.as_ref() else {
            return;
        };
        if !self.is_chained() {
            return;
        }
        let pair_ring = self
            .pair_ring_buf
            .as_ref()
            .expect("pair_ring allocated when Chained is active");

        // Use the cleaner of the two buffers for motion estimation,
        // exactly as `run_motion_compensation` does for the direct path.
        let pyramid = self.pyramid_reference.as_ref().unwrap_or_else(|| {
            self.pyramid_input
                .as_ref()
                .expect("pyramid_input allocated when mc_ctx is Some")
        });

        let total_frames = self.params.total_frames();
        let older_slot = (newer_slot + total_frames - 1) % total_frames;
        let pair_slot = self.pair_slot(0);

        motion::run_pair_analyse::<R>(
            &self.client,
            mc,
            self.width,
            self.height,
            total_frames,
            older_slot,
            newer_slot,
            pair_slot,
            pyramid,
            pair_ring,
            &self.confidence_dummy,
        )
        .expect("pair analyse dispatch failed");
    }

    /// Zero-fill the pair-ring slot for a duplicated ring slot (stream
    /// priming or end-of-stream flush). No-op unless `Chained`
    /// estimation is active.
    fn zero_pair_slot_for_duplicate(&self) {
        let Some(mc) = self.mc_ctx.as_ref() else {
            return;
        };
        if !self.is_chained() {
            return;
        }
        let pair_ring = self
            .pair_ring_buf
            .as_ref()
            .expect("pair_ring allocated when Chained is active");
        let pair_slot = self.pair_slot(0);
        motion::zero_pair_slot::<R>(&self.client, mc, pair_ring, pair_slot);
    }

    /// Queue the Immerkær noise estimate for `slot` on the input ring.
    /// No-op unless auto noise estimation is active. The read of these
    /// results normally happens later in [`Self::denoise_submit`], once
    /// `slot` reaches the centre of the temporal window. The stream's
    /// very first frame also gets an immediate read, see
    /// [`Self::seed_noise_estimate_if_first_frame`].
    fn run_noise_estimate_for_slot(&self, slot: u32) {
        let (Some(partials_buf), Some(results_buf)) =
            (self.noise_partials.as_ref(), self.noise_results.as_ref())
        else {
            return;
        };

        let stride = noise_partials_slot_stride_bytes(self.width, self.height, self.align);
        let partials_slot = partials_buf.clone().offset_start((slot as u64) * stride);

        let ctx = NoiseCtx {
            width: self.width,
            height: self.height,
            channels: self.params.channels.count(),
            stored_ch: self.params.channels.storage_count(),
            frame_count: self.params.total_frames(),
            frame: slot,
            slot,
            input_buf: &self.input_buf,
            partials_buf: &partials_slot,
            results_buf,
        };

        run_noise_estimate::<R>(&self.client, &ctx).expect("noise estimate dispatch failed");
    }

    /// Queue the temporal-residual noise-stats kernel for `slot`,
    /// diffing it against the ring's immediately preceding physical
    /// slot. No-op unless the temporal estimator is active
    /// (`temporal_stats_buf` allocated), and for the very first frame
    /// of a stream (`ring_head == 0`), which has no predecessor to
    /// diff against — mirrors [`Self::run_pair_analyse_for_slot`]'s
    /// same gate. The centre slot's stats are read back and aggregated
    /// later in [`Self::update_noise_estimate`].
    fn run_temporal_stats_for_slot(&self, slot: u32) {
        let Some(stats_buf) = self.temporal_stats_buf.as_ref() else {
            return;
        };
        if self.ring_head == 0 {
            return;
        }

        let total_frames = self.params.total_frames();
        let slot_prev = (slot + total_frames - 1) % total_frames;

        let ctx = TemporalStatsCtx {
            width: self.width,
            height: self.height,
            stored_ch: self.params.channels.storage_count(),
            frame_count: total_frames,
            slot_new: slot,
            slot_prev,
            input_buf: &self.input_buf,
            stats_buf,
            align: self.align,
        };

        run_temporal_noise_stats::<R>(&self.client, &ctx).expect("temporal noise stats dispatch failed");
    }

    /// Zero-fill the duplicated slot's temporal-stats region. A
    /// duplicate mirrors its predecessor's pixels exactly, so a real
    /// diff against it would just recompute an all-zero record; this
    /// is the cheaper equivalent. No-op unless the temporal estimator
    /// is active.
    fn zero_temporal_stats_for_slot(&self, slot: u32) {
        let Some(stats_buf) = self.temporal_stats_buf.as_ref() else {
            return;
        };
        zero_temporal_stats_slot::<R>(
            &self.client,
            stats_buf,
            self.width,
            self.height,
            self.params.channels.storage_count(),
            slot,
            self.align,
        );
    }

    /// One-time σ bootstrap for the very first frame of a stream. Auto
    /// noise estimation normally updates `h2_inv_norm` / `noise_offset`
    /// / `input_noise_offset` from [`Self::update_noise_estimate`] at
    /// submit time, but any push-time GPU work that reads them (the
    /// nlm-spatial pilot) runs before the first submit ever happens.
    /// Without this, that work would run on the absolute-strength
    /// fallback set at construction for every frame up to the first
    /// submit. One blocking read of the estimate this push just queued
    /// for `slot` fixes that from frame one onward. Detects "first
    /// frame of the stream" from `frames_loaded`, the same counter
    /// [`Self::prime_leading_edge_if_first`] checks, but reads it here
    /// before [`Self::advance_ring`] increments it, and applies for
    /// every `temporal_radius` rather than only when priming happens.
    /// The first submit's [`Self::update_noise_estimate`] folds the
    /// same frame's estimate into the EMA a second time, which only
    /// reproduces this seed's values up to floating-point rounding,
    /// not bit-exactly.
    fn seed_noise_estimate_if_first_frame(&mut self, slot: u32) {
        if self.frames_loaded != 0 {
            return;
        }
        let Some(results_buf) = self.noise_results.as_ref() else {
            return;
        };

        let bytes = self
            .client
            .read_one(results_buf.clone())
            .expect("noise-estimate seed readback failed");
        let data = f32::from_bytes(&bytes);

        // The stream's first frame has no predecessor, so
        // `run_temporal_stats_for_slot` never ran for it (see its
        // `ring_head == 0` gate) and this slot's stats region is
        // unwritten. Seed from Immerkær alone, exactly as before the
        // temporal estimator existed.
        let imm_low = self
            .read_noise_partials_low(slot)
            .expect("noise-partials seed readback failed");
        self.fold_noise_estimate(data, slot as usize, None, imm_low);
    }

    /// Fold one physical ring slot's already-read-back noise totals
    /// into the median and low chains' EMAs and recompute the filter
    /// parameters derived from them (`h2_inv_norm`, `noise_offset`,
    /// `input_noise_offset`, `sigma_y`). Shared by
    /// [`Self::seed_noise_estimate_if_first_frame`] and
    /// [`Self::update_noise_estimate`], which differ in how they obtain
    /// `data`/`temporal`/`imm_low` and which slot they pass in.
    ///
    /// Both chains start from an Immerkær read of `data`, maxed with a
    /// temporal-residual read when `temporal` carries a sample — the
    /// temporal residual estimator sees correlated grain the Immerkær
    /// mask underestimates, but scarce static content (motion, scene
    /// changes) makes its sample unreliable, so it can only push either
    /// chain's estimate up, never down. The two chains differ only in
    /// which statistic they read at each step. The median chain reads
    /// `data`'s frame-mean Immerkær total and `temporal.sigma`'s
    /// per-block median, while the low chain reads `imm_low` (Immerkær's
    /// own block-p25) and `temporal.sigma_low`'s per-block lower
    /// quartile. `noise_offset` weighs patch distances quadratically in
    /// sigma, so an over-read there is destructive to fine texture, and
    /// the low chain's conservative statistics keep it from over-reading
    /// on shots where texture leaks into the temporal residuals. The
    /// strength and confidence floor keep the median chain instead,
    /// since that's what the dark-footage calibration validated.
    ///
    /// The temporal sample's `rho` also folds into `rho_smoothed`, the
    /// spatial-offset LUT's attenuation input, once regardless of which
    /// chain reads it. The first fold of a stream sets it directly from
    /// `sample.rho` instead of blending with an assumed 0, the same
    /// convention `NoiseEstimator` uses for its own first sample.
    /// `rho_smoothed` stays unset for the fast path and a fixed
    /// `sigma_override`, since `temporal` is always `None` there.
    fn fold_noise_estimate(
        &mut self,
        data: &[f32],
        slot: usize,
        temporal: Option<TemporalNoiseSample>,
        imm_low: [f32; 3],
    ) {
        let channels = self.params.channels.count() as usize;
        let base = slot * 4;

        let mut raw = [0.0f32; 3];
        for (c, s) in raw.iter_mut().enumerate().take(channels) {
            *s = sigma_from_abs_sum(data[base + c], self.width, self.height);
        }
        let mut raw_low = imm_low;

        if let Some(sample) = temporal {
            let factor = correlation_factor(sample.rho);
            for c in 0..channels {
                raw[c] = raw[c].max(sample.sigma[c] * factor);
                raw_low[c] = raw_low[c].max(sample.sigma_low[c] * factor);
            }
            self.rho_smoothed = Some(match self.rho_smoothed {
                None => sample.rho,
                Some(prev) => EMA_ALPHA * sample.rho + (1.0 - EMA_ALPHA) * prev,
            });
        }

        // User nudge on the measured noise level. Applied after the
        // temporal/Immerkær blend and before the EMA fold, so it scales
        // the smoothed estimate and everything derived from it
        // (`h2_inv_norm`, `noise_offset`, `sigma_y`). Only reached when
        // `sigma_override` is `None` (see `auto_noise` at construction),
        // so `self.params.hq` is always `Some` here.
        let sigma_scale = self.params.hq.map_or(1.0, |hq| hq.sigma_scale);
        for c in 0..channels {
            raw[c] *= sigma_scale;
            raw_low[c] *= sigma_scale;
        }

        let updated = self.noise_estimator.update(&raw[..channels]);
        let mut smoothed = [0.0f32; 3];
        smoothed[..channels].copy_from_slice(updated);

        let updated_low = self.noise_estimator_low.update(&raw_low[..channels]);
        let mut smoothed_low = [0.0f32; 3];
        smoothed_low[..channels].copy_from_slice(updated_low);

        let eff = sigma_eff(&smoothed[..channels], self.params.channels);
        self.h2_inv_norm = self.params.h2_inv_norm_with(Some(eff));
        self.input_noise_offset = self.params.noise_offset_with(Some(&smoothed_low[..channels]));
        self.noise_offset = match self.params.prefilter {
            PrefilterMode::NlmSpatial { .. } => 0.0,
            _ => self.input_noise_offset,
        };
        // Channel 0 is whatever motion estimation already treats as
        // luma (see `nlm_mc_extract_luma`), so the confidence floor
        // uses the median chain's same-plane noise estimate.
        self.sigma_y = smoothed[0];
    }

    fn advance_ring(&mut self) {
        let total_frames = self.params.total_frames() as usize;
        self.ring_head += 1;
        if self.frames_loaded < total_frames {
            self.frames_loaded += 1;
        }
        self.real_pushes += 1;
    }

    /// GPU→GPU copy of one frame from `src`'s slot `src_slot` into
    /// `dst`'s slot `slot`. `dst` must have ring-buffer layout matching
    /// `input_buf` (`total_frames * height * width * stored_ch`).
    /// `src_slots` is how many frames `src` holds, `1` for a
    /// frame-sized staging buffer.
    ///
    /// Both handles are bound whole and the slots addressed by the
    /// kernel's own offset arguments. Binding a slot directly would
    /// need its byte offset to be a multiple of the GPU's
    /// `min_storage_buffer_offset_alignment`, which a
    /// `width * height * stored_ch` frame stride meets only by luck.
    fn copy_frame_into_slot(
        &self,
        dst: &Handle,
        slot: usize,
        src: &Handle,
        src_slot: usize,
        src_slots: usize,
    ) {
        let stored_ch = self.params.channels.storage_count();
        let frame_size = self.width * self.height * stored_ch;
        let dst_slots = self.params.total_frames() as usize;

        let grid = frame_size.div_ceil(BLOCK_1D).min(MAX_GRID_1D);
        let total_threads = grid * BLOCK_1D;

        unsafe {
            gpu_copy::launch_unchecked::<R>(
                &self.client,
                CubeCount::new_1d(grid),
                CubeDim::new_1d(BLOCK_1D),
                ArrayArg::from_raw_parts(src.clone(), src_slots * frame_size as usize),
                ArrayArg::from_raw_parts(dst.clone(), dst_slots * frame_size as usize),
                src_slot as u32 * frame_size,
                slot as u32 * frame_size,
                frame_size,
                total_threads,
            )
        };
    }

    /// Mirror the very first pushed frame into the `R` leading ring
    /// slots so the temporal window starts symmetric instead of dropping
    /// the first `R` logical frames. Mirrors the trailing-edge logic in
    /// [`Self::flush`].
    fn prime_leading_edge_if_first(&mut self) {
        let r = self.params.temporal_radius as usize;

        if r == 0 || self.frames_loaded != 1 {
            return;
        }

        for _ in 0..r {
            self.duplicate_last_frame();
            self.frames_loaded += 1;
        }
    }

    /// Duplicate the most recently pushed frame into the next ring slot.
    /// Used at end-of-stream to keep the window full while future
    /// context shrinks. Slots never overlap, so the in-buffer copy is
    /// well-defined. The reference ring is duplicated in lockstep when
    /// active, so weight calculation never falls back to a stale slot.
    fn duplicate_last_frame(&mut self) {
        let total_frames = self.params.total_frames() as usize;
        let last_slot = (self.ring_head - 1) % total_frames;
        let next_slot = self.ring_head % total_frames;

        let input_buf = self.input_buf.clone();
        self.copy_frame_into_slot(&input_buf, next_slot, &input_buf, last_slot, total_frames);

        // Skipped for `NlmSpatial`: the pilot dispatch below recomputes
        // this slot's reference from scratch, so the byte copy would
        // just be overwritten immediately.
        if !matches!(self.params.prefilter, PrefilterMode::NlmSpatial { .. })
            && let Some(reference_buf) = self.reference_buf.clone()
        {
            self.copy_frame_into_slot(&reference_buf, next_slot, &reference_buf, last_slot, total_frames);
        }

        // Keep the motion-estimation pyramid and noise estimate for the
        // duplicated slot in lockstep so a subsequent denoise sees valid
        // state for every ring slot it visits, not whatever an older
        // frame left behind at this physical position. The nlm-spatial
        // pilot needs the same treatment, otherwise the duplicated
        // slot's reference would keep whatever an older frame at this
        // physical position last wrote there.
        if let PrefilterMode::NlmSpatial { strength_scale } = self.params.prefilter {
            self.run_nlm_spatial_pilot(next_slot as u32, strength_scale)
                .expect("nlm spatial pilot dispatch failed");
        }
        self.build_pyramids_for_slot(next_slot as u32);
        self.build_confidence_pyramid_for_slot(next_slot as u32);
        self.run_noise_estimate_for_slot(next_slot as u32);
        self.zero_temporal_stats_for_slot(next_slot as u32);
        // Runs before `ring_head` advances, so `pair_slot(0)` reads the
        // same pre-advance `ring_head` as `run_pair_analyse_for_slot`
        // (see `Self::pair_slot`).
        self.zero_pair_slot_for_duplicate();

        self.ring_head += 1;
    }

    /// Queue denoise kernels for the current window and kick off an
    /// async readback. Returns a [`Pending`] whose `wait()` produces the
    /// denoised frame.
    ///
    /// Output handles are double-buffered (`outputs: [Handle; 2]`), so
    /// the caller may keep up to `self.outputs.len()` (= 2) `Pending`s
    /// in flight at once, so frame N+1's kernels overlap frame N's
    /// readback. A third concurrent submit would alias the oldest
    /// pending's output handle and silently corrupt results, so the
    /// high-level [`crate::Denoiser`] enforces that cap via its
    /// `MAX_PENDING` constant.
    ///
    /// Returns `Ok(None)` if the temporal window is not yet filled.
    pub fn denoise_submit(&mut self) -> Result<Option<Pending<R>>, anyhow::Error> {
        let total_frames = self.params.total_frames() as usize;
        if self.frames_loaded < total_frames {
            return Ok(None);
        }

        if self.noise_results.is_some() {
            self.update_noise_estimate()?;
        }
        self.rebuild_spatial_offset_lut();

        let slot = self.next_output_slot;
        self.next_output_slot = (slot + 1) % self.outputs.len();

        self.run_denoise_kernels(slot)?;

        // Call `read_async` eagerly so the GPU-side copy is queued before
        // the caller dispatches the next frame's kernels. The future is
        // wrapped in an `async move` that owns a cloned `ComputeClient`
        // (cheap: it's `Arc`-shared internally). That owned client lives
        // inside the future's state machine, so the resulting future is
        // genuinely `'static` and the `Pending` may outlive the
        // `NlmDenoiser` without any lifetime gymnastics.
        let handle = self.outputs[slot].clone();
        let client = self.client.clone();
        let fut = Box::pin(async move { client.read_async(vec![handle]).await });

        let pixels = (self.width * self.height) as usize;
        Ok(Some(Pending {
            fut,
            channels: self.params.channels.count(),
            stored_ch: self.params.channels.storage_count(),
            pixels,
            _marker: PhantomData,
        }))
    }

    /// Refresh `h2_inv_norm` / `noise_offset` from the centre slot's
    /// noise estimate. The centre slot's estimate was queued
    /// `temporal_radius` pushes earlier (see
    /// [`Self::run_noise_estimate_for_slot`]), so this blocking read
    /// lands on work the GPU already finished instead of stalling the
    /// pipeline behind a fresh dispatch.
    fn update_noise_estimate(&mut self) -> Result<(), anyhow::Error> {
        let results_buf = self
            .noise_results
            .as_ref()
            .expect("noise_results allocated when auto noise is active")
            .clone();

        let bytes = self
            .client
            .read_one(results_buf)
            .map_err(|e| anyhow::anyhow!("noise-estimate results readback failed: {e}"))?;
        let data = f32::from_bytes(&bytes);

        let center_t = self.params.temporal_radius;
        let center_slot = self.phys_frame(center_t as i32) as usize;

        let temporal = self.read_temporal_noise_sample(center_slot as u32)?;
        let imm_low = self.read_noise_partials_low(center_slot as u32)?;

        self.fold_noise_estimate(data, center_slot, temporal, imm_low);

        Ok(())
    }

    /// Reads back one ring slot's stage-1 noise partials and reduces
    /// them to the low chain's per-channel block-p25 Immerkær estimate.
    /// Slices the shared ring handle by byte offset so the transfer is
    /// proportional to one slot instead of the whole ring, mirroring
    /// [`read_temporal_stats_slot`].
    fn read_noise_partials_low(&self, slot: u32) -> Result<[f32; 3], anyhow::Error> {
        let partials_buf = self
            .noise_partials
            .as_ref()
            .expect("noise_partials allocated when auto noise is active");

        let slot_len_bytes = partials_len(self.width, self.height) as u64 * size_of::<f32>() as u64;
        let stride = noise_partials_slot_stride_bytes(self.width, self.height, self.align);
        let total_bytes = self.params.total_frames() as u64 * stride;
        let start = (slot as u64) * stride;
        let end_trim = total_bytes - start - slot_len_bytes;

        let sliced = partials_buf.clone().offset_start(start).offset_end(end_trim);
        let bytes = self
            .client
            .read_one(sliced)
            .map_err(|e| anyhow::anyhow!("noise partials readback failed: {e}"))?;
        let data = f32::from_bytes(&bytes);

        Ok(sigma_block_p25_from_partials(
            data,
            self.params.channels.count(),
            self.width,
            self.height,
        ))
    }

    /// Reads back and aggregates the centre slot's temporal-residual
    /// stats. `None` when the temporal estimator is inactive
    /// (`temporal_stats_buf` unallocated: `temporal_radius == 0` or a
    /// fixed `sigma_override`) or when aggregation itself falls back
    /// (see [`aggregate_temporal_noise_stats`]).
    fn read_temporal_noise_sample(&self, slot: u32) -> Result<Option<TemporalNoiseSample>, anyhow::Error> {
        let Some(stats_buf) = self.temporal_stats_buf.as_ref() else {
            return Ok(None);
        };

        let stored_ch = self.params.channels.storage_count();
        let channels = self.params.channels.count();
        let frame_count = self.params.total_frames();

        let records = read_temporal_stats_slot::<R>(
            &self.client,
            stats_buf,
            self.width,
            self.height,
            stored_ch,
            frame_count,
            slot,
            self.align,
        )?;

        Ok(aggregate_temporal_noise_stats(
            &records,
            channels,
            stored_ch,
            self.width,
            self.height,
        ))
    }

    /// Rebuilds `spatial_offset_lut` from the current `noise_offset`
    /// and `rho_smoothed`. Called once per `denoise_submit`, after
    /// `noise_offset` has been refreshed for this submit (when auto
    /// estimation is active) so the LUT and the scalar it's derived
    /// from never disagree. The rebuild is cheap, at most
    /// `(2*8+1)^2` floats.
    fn rebuild_spatial_offset_lut(&mut self) {
        let lut = build_spatial_offset_lut(
            self.params.search_radius,
            self.rho_smoothed.unwrap_or(0.0),
            self.noise_offset,
        );
        self.spatial_offset_lut = self.client.create_from_slice(f32::as_bytes(&lut));
    }

    /// Synchronous convenience wrapper: submits + immediately waits.
    /// Prefer [`Self::denoise_submit`] when the caller can hold one frame
    /// in flight, letting frame N+1's kernels overlap with frame N's
    /// readback.
    ///
    /// Returns `Ok(None)` if not enough frames have been pushed yet.
    /// On success returns `Ok(Some(&[f32]))` borrowing a reusable
    /// internal buffer; copy it out (e.g. `to_vec()`) if you need to
    /// hold the data across another `denoise`/`flush`/`push_frame` call.
    pub fn denoise(&mut self) -> Result<Option<&[f32]>, anyhow::Error> {
        let Some(pending) = self.denoise_submit()? else {
            return Ok(None);
        };
        pending.wait_into(&mut self.output_scratch)?;
        Ok(Some(self.output_scratch.as_slice()))
    }

    /// Flush remaining frames at end-of-stream. For the last `d` frames
    /// the temporal window is clamped by duplicating the last frame.
    /// `sink` is invoked once per produced frame; the borrowed slice is
    /// only valid for that call.
    pub fn flush(&mut self, mut sink: impl FnMut(&[f32])) -> Result<(), anyhow::Error> {
        let temporal_radius = self.params.temporal_radius as usize;
        let total_frames = self.params.total_frames() as usize;

        // Spatial mode has no trailing context to drain.
        if temporal_radius == 0 || self.real_pushes == 0 {
            self.reset_stream_state();
            return Ok(());
        }

        // During pushes the backend submits `real_pushes - R` denoises
        // (zero when `real_pushes <= R`). flush must produce the
        // remainder so the caller gets exactly `real_pushes` outputs.
        let target = self.real_pushes.min(temporal_radius);
        let mut emitted = 0usize;

        // Partial window: pad with trailing duplicates of the last
        // pushed frame so the temporal neighbourhood is complete, then
        // emit centred denoises. Each padded step shifts the centre
        // forward by one, so we may emit several outputs before
        // crossing into the regular trailing-tail loop.
        while self.frames_loaded < total_frames && emitted < target {
            self.duplicate_last_frame();
            self.frames_loaded += 1;
            if self.frames_loaded == total_frames
                && let Some(pending) = self.denoise_submit()?
            {
                pending.wait_into(&mut self.output_scratch)?;
                sink(self.output_scratch.as_slice());
                emitted += 1;
            }
        }

        // Trailing window: full ring, shrinking future context. Each
        // iteration duplicates the most recent frame and emits one more
        // centred denoise.
        while emitted < target {
            self.duplicate_last_frame();
            if let Some(pending) = self.denoise_submit()? {
                pending.wait_into(&mut self.output_scratch)?;
                sink(self.output_scratch.as_slice());
                emitted += 1;
            }
        }

        // Leave the denoiser ready for a fresh stream of the same shape.
        // GPU buffers stay allocated, they get overwritten slot-by-slot
        // as new frames arrive, and `prime_leading_edge_if_first` re-fills
        // the leading edge once `frames_loaded == 1` on the new stream.
        self.reset_stream_state();

        Ok(())
    }

    /// Reset stream-tracking indices so the next `push_frame` begins a
    /// fresh temporal stream. GPU buffers are intentionally not cleared.
    /// `pair_ring_buf` relies on the same write-before-read convention as
    /// the pyramid and noise-estimate buffers. A fresh stream's first
    /// pushes fully overwrite every slot they touch before anything
    /// reads it back, so leftover content from the previous stream is
    /// never observed.
    pub(crate) fn reset_stream_state(&mut self) {
        self.ring_head = 0;
        self.frames_loaded = 0;
        self.next_output_slot = 0;
        self.real_pushes = 0;
        self.noise_estimator.reset();
        self.noise_estimator_low.reset();
        self.rho_smoothed = None;
    }

    /// Physical slot of logical frame 0 (oldest frame in the window).
    /// Defined only once a full window has been pushed.
    pub(super) fn ring_start(&self) -> u32 {
        let total_frames = self.params.total_frames() as usize;
        (self.ring_head % total_frames) as u32
    }

    /// Resolve a logical frame index in `[0, total_frames)` to its
    /// physical slot inside `input_buf`.
    pub(super) fn phys_frame(&self, logical: i32) -> u32 {
        let total_frames = self.params.total_frames() as i32;
        let wrapped = logical.rem_euclid(total_frames);
        ((self.ring_start() as i32 + wrapped).rem_euclid(total_frames)) as u32
    }

    /// Pair-ring slot for the gap between window-relative logical
    /// frames `gap_index` and `gap_index + 1`. Reduces the current
    /// `ring_head`, the monotonic count of frames pushed so far
    /// (including duplicates), by `2 * temporal_radius` instead of the
    /// `total_frames` modulus `Self::phys_frame` uses for the frame
    /// ring.
    ///
    /// Called two ways that resolve to the same slot for the same
    /// physical pair. At push time, with `gap_index = 0` and the
    /// pre-advance `ring_head` (the generation of the frame just
    /// written), it gives the slot that frame's pair with its
    /// immediate predecessor belongs in. At compose time, with the
    /// post-push `ring_head` and `gap_index` measured from the
    /// window's centre, it gives the slot a past push already wrote
    /// to. The two calls only differ in how far `ring_head` has
    /// advanced since the pair was created, and `gap_index` exactly
    /// offsets that advance, so `ring_head + gap_index` lands on the
    /// same value mod `2 * temporal_radius` either way.
    pub(super) fn pair_slot(&self, gap_index: i32) -> u32 {
        let radius = self.params.temporal_radius as i32;
        debug_assert!(
            radius > 0,
            "pair ring is only meaningful when temporal_radius > 0"
        );
        let n = 2 * radius;
        ((self.ring_head as i32 + gap_index).rem_euclid(n)) as u32
    }
}