av-denoise 0.3.1

use super::helpers::*;
use crate::nlmeans::*;

/// Shared HQ auto-estimation params, k=0 only, so the temporal chain
/// stays inert (`temporal_stats_buf` needs `temporal_radius >= 1`) and
/// the two Immerkær spatial statistics (frame mean vs block p25) are
/// the only thing that can tell the median and low chains apart.
fn auto_params_k0() -> NlmParams {
    NlmParams {
        temporal_radius: 0,
        search_radius: 3,
        patch_radius: 2,
        strength: 1.2,
        self_weight: 1.0,
        channels: ChannelMode::Luma,
        prefilter: PrefilterMode::None,
        motion_compensation: MotionCompensationMode::None,
        hq: Some(HqParams {
            auto_strength: true,
            noise_floor: true,
            sigma_override: None,
            temporal_confidence: false,
            thsad_scale: 1.0,
            sigma_scale: 1.0,
        }),
    }
}

/// A frame whose top half carries `sigma_a` noise and whose bottom
/// half carries `sigma_b`, both generated by `make_noisy_gaussian_frame`
/// at the same base level and spliced row-wise.
fn block_heterogeneous_frame(w: u32, h: u32, base: f32, sigma_a: f32, sigma_b: f32) -> Vec<f32> {
    let top = make_noisy_gaussian_frame(w, h, 1, base, &[sigma_a]);
    let bottom = make_noisy_gaussian_frame(w, h, 1, base, &[sigma_b]);
    let row_len = w as usize;
    let half = (h / 2) as usize;
    let mut out = top;
    for row in half..h as usize {
        let start = row * row_len;
        out[start..start + row_len].copy_from_slice(&bottom[start..start + row_len]);
    }
    out
}

/// Spatially uniform synthetic noise. The low chain's block-p25 sits
/// close to the median chain's frame mean (uniform noise means p25 is
/// nearly the median), so the derived `noise_offset` must stay close
/// to what the pre-split single-chain design would have produced from
/// the median chain's own smoothed sigma.
#[test]
fn uniform_noise_low_chain_matches_median_chain() {
    let client = make_client();
    let w = 256;
    let h = 256;
    let sigma = 8.0 / 255.0;
    let frame = make_noisy_gaussian_frame(w, h, 1, 0.5, &[sigma]);

    let mut denoiser = NlmDenoiser::<R>::new(&client, auto_params_k0(), w, h);
    denoiser.push_frame(&frame);
    denoiser.denoise().unwrap();

    let median = denoiser.noise_estimator.current().expect("seeded on first push")[0];
    let low = denoiser
        .noise_estimator_low
        .current()
        .expect("seeded on first push")[0];

    let rel_err = (low - median).abs() / median;
    assert!(
        rel_err <= 0.15,
        "uniform noise: low chain {low} should sit close to median chain {median} (rel err {rel_err:.3})"
    );

    let expected_offset = denoiser.params.noise_offset_with(Some(&[median]));
    let offset_rel_err = (denoiser.noise_offset - expected_offset).abs() / expected_offset;
    assert!(
        offset_rel_err <= 0.3,
        "uniform noise: noise_offset {} should stay close to the pre-split median-based offset \
         {expected_offset} (rel err {offset_rel_err:.3})",
        denoiser.noise_offset
    );
}

/// Block-heterogeneous noise (half the frame at a low sigma, half at a
/// much higher sigma). The low chain must read below the median chain,
/// so `noise_offset` tracks the conservative statistic strictly below
/// what the median chain's own smoothed sigma would produce, while
/// `h2_inv_norm` (strength) keeps tracking the median chain exactly.
#[test]
fn split_noise_offset_tracks_low_chain_strength_tracks_median() {
    let client = make_client();
    let w = 256;
    let h = 256;
    let sigma_a = 2.0 / 255.0;
    let sigma_b = 20.0 / 255.0;
    let frame = block_heterogeneous_frame(w, h, 0.5, sigma_a, sigma_b);

    let mut denoiser = NlmDenoiser::<R>::new(&client, auto_params_k0(), w, h);
    denoiser.push_frame(&frame);
    denoiser.denoise().unwrap();

    let median = denoiser.noise_estimator.current().expect("seeded on first push")[0];
    let low = denoiser
        .noise_estimator_low
        .current()
        .expect("seeded on first push")[0];

    assert!(
        low < median,
        "block-heterogeneous noise: low chain {low} should read below median chain {median}"
    );

    let expected_median_offset = denoiser.params.noise_offset_with(Some(&[median]));
    assert!(
        denoiser.noise_offset < expected_median_offset,
        "noise_offset {} should track the low chain strictly below the median-based offset {expected_median_offset}",
        denoiser.noise_offset
    );

    let expected_h2 = denoiser.params.h2_inv_norm_with(Some(median));
    assert_eq!(
        denoiser.h2_inv_norm, expected_h2,
        "h2_inv_norm must keep tracking the median chain exactly"
    );
}

/// Ring isolation. A slot's own stage-1 partials must survive untouched
/// from the push that queues them until that slot reaches the centre
/// and gets folded, even though several other slots get pushed (and
/// dispatch their own noise estimate) in between.
///
/// `temporal_radius = 2` (a 5-slot window) primes the leading edge with
/// two duplicates of frame 0 (low sigma), then two further high-sigma
/// pushes fill the window and trigger the first denoise. Its centre
/// slot is one of those low-sigma duplicates, untouched since priming.
/// A single shared (non-ring) partials scratch buffer would instead
/// have been overwritten by the most recent high-sigma dispatch by the
/// time it's read back, so the low chain's estimate would jump toward
/// the high sigma instead of staying low.
///
/// One more high-sigma push then rotates the centre onto the first
/// real high-sigma frame, and the low chain must rise to reflect it,
/// confirming the ring keeps serving each slot's own fresh data as it
/// becomes centre rather than latching onto whichever slot happened to
/// isolate correctly first.
#[test]
fn partials_ring_isolates_slots_between_push_and_fold() {
    let client = make_client();
    let w = 128;
    let h = 128;
    let sigma_low = 2.0 / 255.0;
    let sigma_high = 30.0 / 255.0;

    let params = NlmParams {
        temporal_radius: 2,
        ..auto_params_k0()
    };
    let mut denoiser = NlmDenoiser::<R>::new(&client, params, w, h);

    let frame_low = make_noisy_gaussian_frame(w, h, 1, 0.5, &[sigma_low]);
    let frame_high_1 = make_noisy_gaussian_frame(w, h, 1, 0.5, &[sigma_high]);
    let frame_high_2 = make_noisy_gaussian_frame(w, h, 1, 0.5, &[sigma_high]);
    let frame_high_3 = make_noisy_gaussian_frame(w, h, 1, 0.5, &[sigma_high]);

    denoiser.push_frame(&frame_low);
    assert!(denoiser.denoise().unwrap().is_none(), "window not full yet");
    denoiser.push_frame(&frame_high_1);
    assert!(denoiser.denoise().unwrap().is_none(), "window not full yet");
    denoiser.push_frame(&frame_high_2);
    assert!(
        denoiser.denoise().unwrap().is_some(),
        "window should be full after the third push"
    );

    let low1 = denoiser.noise_estimator_low.current().expect("folded by now")[0];
    assert!(
        low1 < sigma_high * 0.15,
        "the first centred slot duplicates the low-sigma frame and its own partials \
         must still be intact, got low chain estimate {low1} (high sigma is {sigma_high})"
    );

    denoiser.push_frame(&frame_high_3);
    assert!(denoiser.denoise().unwrap().is_some());

    let low2 = denoiser.noise_estimator_low.current().expect("folded by now")[0];
    assert!(
        low2 > low1 * 1.5,
        "the centre should now be the first high-sigma push, so the low chain must rise \
         to reflect it (low1={low1}, low2={low2})"
    );
}