maroontree 0.1.8

/*
 * Copyright (c) Radzivon Bartoshyk 6/2026. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without modification,
 * are permitted provided that the following conditions are met:
 *
 * 1.  Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 *
 * 2.  Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 *
 * 3.  Neither the name of the copyright holder nor the names of its
 * contributors may be used to endorse or promote products derived from
 * this software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

use crate::util::FastRound;

pub(crate) trait AqLuma: Copy {
    fn to_f32(self) -> f32;
}
impl AqLuma for f32 {
    #[inline(always)]
    fn to_f32(self) -> f32 {
        self
    }
}
impl AqLuma for i32 {
    #[inline(always)]
    fn to_f32(self) -> f32 {
        self as f32
    }
}

/// The representative variance of a superblock for Variance Boost: the value at the
/// requested `octile` (1..=8) of the 64 sorted 8x8 variances. Octile 1 = the most
/// low-variance-biased pick (boost readily), octile 8 = only the maximum (boost only
/// when the whole SB is low-variance).
pub(crate) fn sb_octile_variance(subvars: &mut [f32; 64], octile: u8) -> f32 {
    subvars.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
    // Octile o in 1..=8 maps to sorted index o*8 - 1 (o=1 -> 7, o=4 -> 31 (median-ish),
    // o=6 -> 47 (SVT-AV1-PSY default), o=8 -> 63 (max)).
    let o = octile.clamp(1, 8) as usize;
    let idx = (o * 8 - 1).min(63);
    subvars[idx]
}

#[inline(always)]
pub(crate) fn f_fmlaf(a: f32, b: f32, c: f32) -> f32 {
    #[cfg(any(
        all(
            any(target_arch = "x86", target_arch = "x86_64"),
            target_feature = "fma"
        ),
        target_arch = "aarch64"
    ))]
    {
        f32::mul_add(a, b, c)
    }
    #[cfg(not(any(
        all(
            any(target_arch = "x86", target_arch = "x86_64"),
            target_feature = "fma"
        ),
        target_arch = "aarch64"
    )))]
    {
        a * b + c
    }
}

#[inline]
pub(crate) fn dirty_log1pf(d: f32) -> f32 {
    let ix = (1.0 + d).to_bits();
    let exponent = ix & 0x7f80_0000;
    let n = (exponent >> 23) as i32 - 0x7f;

    // Replacing the exponent with 127 normalizes 1+d to 1+t without division.
    // For n=0 use d itself so tiny inputs do not lose precision in (1+d)-1.
    let mantissa = f32::from_bits((ix & 0x007f_ffff) | 0x3f80_0000);
    let t = if n == 0 { d } else { mantissa - 1.0 };

    // Direct minimax ln(1+t) ~= t*P(t), t in [0, 1]. See tools/log1p.sollya.
    let mut p = 0.014539075084030628;
    p = f_fmlaf(p, t, -0.0675969123840332);
    p = f_fmlaf(p, t, 0.15056970715522766);
    p = f_fmlaf(p, t, -0.23573730885982513);
    p = f_fmlaf(p, t, 0.33125850558280945);
    p = f_fmlaf(p, t, -0.4998837411403656);
    p = f_fmlaf(p, t, 0.999998927116394);
    f_fmlaf(n as f32, std::f32::consts::LN_2, t * p)
}

/// Variance Boost qindex delta for one superblock. `picked_var` is the octile pick
/// (see [`sb_octile_variance`]); `ref_log` is the tile mean log-variance.
#[inline(never)]
pub(crate) fn variance_boost_delta(
    picked_var: f32,
    ref_log: f32,
    strength: f32,
    boost_only: bool,
) -> i32 {
    // Work in log-variance: compresses the huge dynamic range of variance and matches
    // the reference (which is a mean of log-variances).
    let v_log = dirty_log1pf(picked_var);
    // Low-variance threshold (curve 0): ln(1 + 256).
    const LOW_LOG: f32 = 5.549_076; // (1.0 + 256.0).ln()
    const MAX_BOOST: f32 = 18.0; // max qindex *reduction* for the flattest SBs
    const MAX_CUT: f32 = 10.0; // max qindex *increase* for the busiest SBs
    // qindex per unit log-variance for each side.
    const BOOST_SLOPE: f32 = 5.0;
    const CUT_SLOPE: f32 = 3.0;

    if v_log < LOW_LOG {
        // Low contrast: boost (negative delta). Deeper below threshold => stronger.
        let d = ((LOW_LOG - v_log) * BOOST_SLOPE * strength).min(MAX_BOOST);
        -(d.fast_round() as i32)
    } else if boost_only {
        0
    } else {
        // Higher contrast: coarsen relative to the tile reference, capped. Using the
        // reference (not the threshold) keeps well-textured frames near zero-mean.
        let over = (v_log - ref_log.max(LOW_LOG)).max(0.0);
        let d = (over * CUT_SLOPE * strength).min(MAX_CUT);
        d.fast_round() as i32
    }
}

#[derive(Clone, Copy, Debug)]
pub struct DarkAq {
    pub enabled: bool,
    /// Only active at base_q >= this (qstep-extinction range, roughly q<=45 in 8-bit).
    pub min_q: i32,
    pub mean_floor: f32,
    pub dark_ref: f32,
    pub gamma: f32,
    pub max_weight: f32,
    /// qindex units per unit of `log1p(mid_energy * dark_weight)`.
    pub scale: f32,
    /// Cap on the extra boost (qindex reduction).
    pub max_qidx: i32,
}

impl Default for DarkAq {
    fn default() -> Self {
        DarkAq {
            enabled: false,
            min_q: 150,
            mean_floor: 16.0,
            dark_ref: 56.0,
            gamma: 1.2,
            // `darkness = max_weight - 1` is the effective multiplier for the darkest SBs.
            max_weight: 4.5,
            scale: 4.0,
            max_qidx: 16,
        }
    }
}

impl DarkAq {
    /// Disabled — no dark protection.
    pub(crate) fn off() -> Self {
        DarkAq {
            enabled: false,
            ..DarkAq::default()
        }
    }

    /// Enabled with the calibrated defaults.
    pub(crate) fn on() -> Self {
        DarkAq {
            enabled: true,
            ..DarkAq::default()
        }
    }
}

#[inline]
fn laplacian_abs_sum<const STRIDE: usize>(buf: &[[f32; STRIDE]], h: usize, w: usize) -> (f32, u32) {
    let mut sum = 0.0f32;
    let mut n = 0u32;
    for rows in buf[..h].array_windows::<3>() {
        let [top, middle, bottom] = rows;
        for ((&up, &down), &[left, center, right]) in top[1..w - 1]
            .iter()
            .zip(bottom[1..w - 1].iter())
            .zip(middle[..w].array_windows::<3>())
        {
            let l = 4.0 * center - up - down - left - right;
            sum += l.abs();
            n += 1;
        }
    }
    (sum, n)
}

#[inline]
fn box_downsample_2x<const SRC_STRIDE: usize, const DST_STRIDE: usize>(
    src: &[[f32; SRC_STRIDE]],
    h: usize,
    w: usize,
    dst: &mut [[f32; DST_STRIDE]],
) -> (usize, usize) {
    let (hh, ww) = (h / 2, w / 2);
    for (dst_row, rows) in dst[..hh]
        .iter_mut()
        .zip(src[..h].array_windows::<2>().step_by(2))
    {
        let [top, bottom] = rows;
        for (out, (&[top_left, top_right], &[bottom_left, bottom_right])) in
            dst_row[..ww].iter_mut().zip(
                top[..w]
                    .array_windows::<2>()
                    .step_by(2)
                    .zip(bottom[..w].array_windows::<2>().step_by(2)),
            )
        {
            *out = 0.25 * (top_left + top_right + bottom_left + bottom_right);
        }
    }
    (hh, ww)
}

pub(crate) fn dark_structure_stats<T: AqLuma>(
    yp: &[T],
    pw: usize,
    sb_y: usize,
    sb_x: usize,
    width: usize,
    height: usize,
    scale: f32,
) -> (f32, f32) {
    let h = height.saturating_sub(sb_y).min(64);
    let w = width.saturating_sub(sb_x).min(64);
    if h == 0 || w == 0 {
        return (0.0, 0.0);
    }
    let mut buf = [[0f32; 64]; 64];
    let mut sum = 0f32;
    for (r, row) in buf.iter_mut().enumerate().take(h) {
        let base = (sb_y + r) * pw + sb_x;
        for (yp, dst) in yp[base..base + w].iter().zip(row.iter_mut()) {
            let v = yp.to_f32() * scale;
            *dst = v;
            sum += v;
        }
    }
    let mean = sum / (h * w) as f32;
    if h < 3 || w < 3 {
        return (mean, 0.0);
    }
    // Full-resolution |Laplacian| (interior 3×3).
    let (lap_full, nf) = laplacian_abs_sum(&buf, h, w);
    let lap_full = lap_full / nf as f32;
    // 2× box downsample, then |Laplacian| at the coarse scale.
    let (hh, ww) = (h / 2, w / 2);
    if hh < 3 || ww < 3 {
        return (mean, 0.0);
    }
    let mut half = [[0f32; 32]; 32];
    let (hh, ww) = box_downsample_2x(&buf, h, w, &mut half);
    if hh < 3 || ww < 3 {
        return (mean, 0.0);
    }
    let (lap_half, nh) = laplacian_abs_sum(&half, hh, ww);
    let lap_half = lap_half / nh as f32;
    (mean, (lap_full * lap_half).sqrt())
}

/// Extra qindex reduction (>= 0) for a dark, structured SB. 0 when disabled, out of the
/// gated quality range, or the SB carries no cross-scale structure. `scale` normalizes
/// the plane to 8-bit range for [`dark_structure_stats`] (AV2: `1.0`; AV1: `1/(1<<(bd-8))`).
#[allow(clippy::too_many_arguments)]
pub(crate) fn dark_protection<T: AqLuma>(
    d: &DarkAq,
    base_q: i32,
    yp: &[T],
    pw: usize,
    sb_y: usize,
    sb_x: usize,
    width: usize,
    height: usize,
    scale: f32,
) -> i32 {
    if !d.enabled || base_q < d.min_q {
        return 0;
    }
    let (mean, mid_energy) = dark_structure_stats(yp, pw, sb_y, sb_x, width, height, scale);
    if mid_energy <= 0.0 {
        return 0;
    }
    // Darker SBs get a heavier weight; `dark_weight` is 1.0 at/above `dark_ref` and
    // rises toward `max_weight` as the SB darkens. Subtracting 1 makes the protection
    // vanish for bright/mid SBs (so it doesn't just boost all texture) and scale with
    // how far below the reference the SB sits.
    let dark_weight = ((d.mean_floor + d.dark_ref) / (d.mean_floor + mean))
        .powf(d.gamma)
        .clamp(1.0, d.max_weight);
    let darkness = dark_weight - 1.0;
    if darkness <= 0.0 {
        return 0;
    }
    let dark_structure = dirty_log1pf(mid_energy * darkness);
    (dark_structure * d.scale)
        .min(d.max_qidx as f32)
        .max(0.0)
        .fast_round() as i32
}