rusty_mp3 0.5.1

//! Short-block coding (brick **Q5**) — the forward reorder and the short-block
//! quantizer.
//!
//! Short blocks give the filterbank finer time resolution, which suppresses the
//! pre-echo a long block smears before a transient. The forward MDCT (`mdct.rs`)
//! already produces the three short windows interleaved in *subband* order; the
//! bitstream stores them in *scalefactor-band* order (sfb-major, then window, then
//! frequency). [`reorder_subband_to_bitstream`] is the exact inverse of the
//! decoder's requantize reorder, so a coded short block round-trips.
//!
//! This first version quantises uniformly (flat per-window scalefactors,
//! `subblock_gain = 0`) under global-gain rate control — valid short blocks that
//! fix pre-echo. Per-window perceptual shaping is a later refinement.

use crate::frame::{BlockType, GRANULE_LINES};
use crate::header::FrameHeader;
use crate::tables;

use super::quantize::QuantizedGranule;

const MAX_UNCLIPPED: i32 = 8191;

/// Reorder a short-block spectrum from the MDCT's subband-interleaved order into
/// the bitstream's `(sfb, window, freq)` order — the forward of the decoder's
/// requantize reorder (`out[dst] = coeffs[src]`).
pub fn reorder_subband_to_bitstream(
    sample_rate: u32,
    subband: &[f32; GRANULE_LINES],
) -> [f32; GRANULE_LINES] {
    let off = tables::sfb_short_offsets(sample_rate);
    let mut out = [0f32; GRANULE_LINES];
    for sfb in 0..13 {
        let start = off[sfb] as usize;
        let width = (off[sfb + 1] - off[sfb]) as usize;
        for window in 0..3 {
            for f in 0..width {
                let src = start * 3 + window * width + f;
                let dst = start * 3 + window + f * 3;
                if src < GRANULE_LINES && dst < GRANULE_LINES {
                    out[src] = subband[dst];
                }
            }
        }
    }
    out
}

/// Uniformly quantize bitstream-order short-block lines at `global_gain`
/// (`subblock_gain = 0`, flat scalefactors), the forward of the decoder's short
/// requantization with those fields zero. Uses the precomputed `|freq|^(3/4)`
/// (A1): each gain probe is a multiply-and-round, no per-line `powf`.
fn quantize_uniform(
    freq: &[f32; GRANULE_LINES],
    xrp: &[f64; GRANULE_LINES],
    gain: i32,
) -> [i32; GRANULE_LINES] {
    // step = scale_inv^(3/4), scale_inv = 2^(-0.25·(gain−210)).
    let step = 2f64.powf(0.75 * -0.25 * (gain - 210) as f64);
    let mut coeffs = [0i32; GRANULE_LINES];
    for (i, &x) in freq.iter().enumerate() {
        let mag = super::quantize::level_from(xrp[i] * step);
        coeffs[i] = if x < 0.0 { -mag } else { mag };
    }
    coeffs
}

/// **Q5 short quantizer** — quantize a short-block granule (already reordered to
/// bitstream order) to fit `bit_budget`, with flat per-window scalefactors. Picks
/// the smallest non-clipping `global_gain` whose Huffman cost fits.
pub fn quantize_short(
    header: &FrameHeader,
    freq_bitstream: &[f32; GRANULE_LINES],
    bit_budget: usize,
) -> QuantizedGranule {
    let xrp = super::quantize::xrpow(freq_bitstream); // A1: hoist |freq|^(3/4)
    let ok = |g: i32| {
        let coeffs = quantize_uniform(freq_bitstream, &xrp, g);
        coeffs.iter().all(|&c| c.abs() <= MAX_UNCLIPPED)
            && super::huffman::cost_short(header, &coeffs) <= bit_budget
    };
    let (mut lo, mut hi) = (0i32, 255i32);
    while lo < hi {
        let mid = (lo + hi) / 2;
        if ok(mid) {
            hi = mid;
        } else {
            lo = mid + 1;
        }
    }
    let coeffs = quantize_uniform(freq_bitstream, &xrp, lo);
    let (mut side, _) = super::huffman::select(header, &coeffs, BlockType::Short);
    side.global_gain = lo as u8;
    side.scalefac_compress = 0; // flat scalefactors → zero scalefactor bits
    QuantizedGranule {
        coeffs,
        side,
        scalefactors: [0; 39],
    }
}

/// **Q5 block-type FSM.** Given the previous granule's window type and this
/// frame's per-granule attack flags, choose valid window types that bracket every
/// attack with the required transition windows. Produces sequences like
/// `LONG…LONG | START, SHORT | (SHORT, SHORT)* | STOP, LONG` — a `START` always
/// precedes a `SHORT` and a `STOP` always follows it, the constraint the decoder's
/// overlapping windows require. MPEG-1 (2 granules per frame).
pub fn decide_block_types(prev: BlockType, attacks: &[bool]) -> (Vec<BlockType>, BlockType) {
    if attacks.len() != 2 {
        // MPEG-2 (single granule) not yet block-switched.
        return (vec![BlockType::Long; attacks.len()], BlockType::Long);
    }
    let any = attacks[0] || attacks[1];
    let types = if prev == BlockType::Short {
        // Inside a short run: continue or wind down with a STOP.
        if any {
            [BlockType::Short, BlockType::Short]
        } else {
            [BlockType::Stop, BlockType::Long]
        }
    } else if any {
        // Long/Stop and an attack this frame: transition in (START → SHORT).
        [BlockType::Start, BlockType::Short]
    } else {
        [BlockType::Long, BlockType::Long]
    };
    let new_prev = types[1];
    (types.to_vec(), new_prev)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::decode::scalefactors::ScaleFactors;
    use crate::frame::{ChannelMode, SideInfo};
    use crate::header::MpegVersion;

    fn hdr() -> FrameHeader {
        FrameHeader {
            version: MpegVersion::V1,
            crc_protected: false,
            bitrate_kbps: 128,
            sample_rate: 44100,
            padding: false,
            channel_mode: ChannelMode::Mono,
            copyright: false,
            original: true,
            emphasis: 0,
        }
    }

    #[test]
    fn short_coefficients_round_trip_through_decoder() {
        // A short-block spectrum (subband order) → reorder → short quantize →
        // Huffman → decoder Huffman + reorder-requantize must recover it.
        let header = hdr();
        let mut s = 0xABCD_1234u32;
        let mut freq = [0f32; GRANULE_LINES];
        for (i, v) in freq.iter_mut().enumerate() {
            s = s.wrapping_mul(1_664_525).wrapping_add(1_013_904_223);
            // taper high lines so there's an rzero tail (realistic short block)
            let taper = 1.0 - (i as f32 / GRANULE_LINES as f32);
            *v = ((s >> 8) as f32 / (1u32 << 24) as f32 - 0.5) * 8.0 * taper;
        }

        let freq_bs = reorder_subband_to_bitstream(header.sample_rate, &freq);
        let quant = quantize_short(&header, &freq_bs, 100_000);

        // Emit the Huffman spectrum, then decode it back.
        let mut w = crate::bitio::BitWriter::new();
        let hbits = super::super::huffman::encode(&quant, &header, &mut w);
        let bits = w.finish();

        let mut gi = quant.side.clone();
        gi.part2_3_length = hbits as u16;
        let mut pos = 0;
        let (coeffs, nz) = crate::decode::huffman::decode(&bits, &mut pos, hbits, &header, &gi);
        assert_eq!(
            coeffs, quant.coeffs,
            "short Huffman must round-trip exactly"
        );

        // Requantize (with the decoder's short reorder) → subband order.
        let mut si = SideInfo::default();
        si.granules[0][0] = gi.clone();
        let mut out = [0f32; GRANULE_LINES];
        crate::decode::requantize::apply(
            &header,
            &gi,
            &ScaleFactors::default(),
            &coeffs,
            nz,
            &mut out,
        );

        // out (subband order) ≈ freq (subband order), to quantization error.
        let mut sig = 0f64;
        let mut err = 0f64;
        for i in 0..GRANULE_LINES {
            sig += (freq[i] as f64).powi(2);
            err += ((freq[i] - out[i]) as f64).powi(2);
        }
        let snr = 10.0 * (sig / err).log10();
        eprintln!("[Q5] short-block coefficient round-trip SNR {snr:.1} dB");
        assert!(snr > 40.0, "short-block coding SNR too low: {snr:.1} dB");
    }

    #[test]
    fn reorder_inverts_the_decoder_mapping() {
        // Forward reorder then the decoder's src/dst mapping must return the
        // original subband-order spectrum.
        let subband: [f32; GRANULE_LINES] = std::array::from_fn(|i| (i as f32 * 0.013).sin());
        let bitstream = reorder_subband_to_bitstream(44100, &subband);

        let off = tables::sfb_short_offsets(44100);
        let mut recovered = [0f32; GRANULE_LINES];
        for sfb in 0..13 {
            let start = off[sfb] as usize;
            let width = (off[sfb + 1] - off[sfb]) as usize;
            for window in 0..3 {
                for f in 0..width {
                    let src = start * 3 + window * width + f;
                    let dst = start * 3 + window + f * 3;
                    recovered[dst] = bitstream[src]; // decoder's placement
                }
            }
        }
        for i in 0..GRANULE_LINES {
            assert!((recovered[i] - subband[i]).abs() < 1e-9, "line {i}");
        }
    }
}

/// Map the LONG-block masking thresholds onto the short-block band grid.
///
/// The psychoacoustic model only produces long-band thresholds (its
/// `block_type` is hard-wired to `Long`), so a short granule has no masking
/// curve of its own. Short blocks have 3x coarser frequency resolution — 192
/// lines per window against 576 — so short band `b`, spanning per-window lines
/// `[s0, s1)`, covers the same frequencies as long lines `[3·s0, 3·s1)`.
///
/// The overlapping long bands are combined by taking the MINIMUM. That is the
/// conservative direction on purpose: a short band straddling a quiet long band
/// must not inherit a loud neighbour's noise allowance, because the audible
/// failure mode here is pre-echo on exactly the transients short blocks exist
/// to code.
///
/// This is an approximation, and it is the reason short-block VBR is coarser
/// than long-block VBR. The principled fix is a short-block psymodel.
fn short_thresholds(sample_rate: u32, long_thr: &[f32; crate::frame::SFB_LONG]) -> [f32; 13] {
    let long_off = crate::tables::sfb_long_offsets(sample_rate);
    let short_off = crate::tables::sfb_short_offsets(sample_rate);
    let mut out = [f32::MAX; 13];
    for (b, o) in out.iter_mut().enumerate() {
        let (lo, hi) = (3 * short_off[b] as usize, 3 * short_off[b + 1] as usize);
        let mut m = f32::MAX;
        for lb in 0..21 {
            let (llo, lhi) = (long_off[lb] as usize, long_off[lb + 1] as usize);
            if llo < hi && lo < lhi {
                m = m.min(long_thr[lb]);
            }
        }
        *o = if m == f32::MAX { long_thr[20] } else { m };
    }
    out
}

/// Quantization noise per short band, summed across the three windows.
///
/// In bitstream order a short band's three windows are adjacent, so band `b`
/// occupies reordered indices `[3·off[b], 3·off[b+1])` — the same span the
/// threshold mapping above uses.
fn short_band_noise(
    sample_rate: u32,
    freq_bs: &[f32; GRANULE_LINES],
    coeffs: &[i32; GRANULE_LINES],
    gain: i32,
) -> [f32; 13] {
    let off = crate::tables::sfb_short_offsets(sample_rate);
    let scale = 2f64.powf(0.25 * (gain - 210) as f64);
    let mut noise = [0f32; 13];
    for (b, n) in noise.iter_mut().enumerate() {
        let (lo, hi) = (
            (3 * off[b] as usize).min(GRANULE_LINES),
            (3 * off[b + 1] as usize).min(GRANULE_LINES),
        );
        let mut e = 0f64;
        for i in lo..hi {
            let xr = coeffs[i].signum() as f64
                * super::quantize::requant_magnitude(coeffs[i])
                * scale;
            let d = freq_bs[i] as f64 - xr;
            e += d * d;
        }
        *n = e as f32;
    }
    noise
}

/// **VBR short blocks.** Coarsest gain (fewest bits) whose peak NMR still meets
/// `target_nmr`, bounded by the frame's physical capacity.
///
/// Short granules previously ignored `-q:a` entirely and always took the CBR bit
/// budget, so on transient-heavy content — where short blocks are 15-37% of
/// granules — the quality knob was partly dead even after the VBR rate control
/// was fixed.
///
/// Mirrors `quantize::loops_vbr`: same search shape, and the same domain
/// rescaling, because the thresholds live in the FFT power domain while the
/// noise is measured in the MDCT domain.
pub fn quantize_short_vbr(
    header: &FrameHeader,
    freq_bitstream: &[f32; GRANULE_LINES],
    long_thresholds: &[f32; crate::frame::SFB_LONG],
    signal_energy: f32,
    target_nmr: f32,
    bit_ceiling: usize,
) -> QuantizedGranule {
    let xrp = super::quantize::xrpow(freq_bitstream);
    let thr = short_thresholds(header.sample_rate, long_thresholds);
    let mdct_energy: f32 = freq_bitstream.iter().map(|x| x * x).sum();
    let domain_scale = if signal_energy > 1e-20 && mdct_energy > 1e-20 {
        mdct_energy / signal_energy
    } else {
        1.0
    };
    let ok = |g: i32| {
        let coeffs = quantize_uniform(freq_bitstream, &xrp, g);
        if !coeffs.iter().all(|&c| c.abs() <= MAX_UNCLIPPED) {
            return false;
        }
        // A quality target is never a licence to overflow the frame.
        if super::huffman::cost_short(header, &coeffs) > bit_ceiling {
            return false;
        }
        let noise = short_band_noise(header.sample_rate, freq_bitstream, &coeffs, g);
        noise
            .iter()
            .zip(thr.iter())
            .map(|(&n, &t)| n / (t * domain_scale).max(1e-20))
            .fold(0f32, f32::max)
            <= target_nmr
    };
    // Largest gain (coarsest, fewest bits) that still satisfies everything.
    let (mut lo, mut hi) = (0i32, 255i32);
    while lo < hi {
        let mid = (lo + hi + 1) / 2;
        if ok(mid) {
            lo = mid;
        } else {
            hi = mid - 1;
        }
    }
    // If nothing met the target, fall back to the budget-driven path rather than
    // emitting an over-target or clipped granule.
    if !ok(lo) {
        return quantize_short(header, freq_bitstream, bit_ceiling);
    }
    let coeffs = quantize_uniform(freq_bitstream, &xrp, lo);
    let (mut side, _) = super::huffman::select(header, &coeffs, BlockType::Short);
    side.global_gain = lo as u8;
    side.scalefac_compress = 0;
    QuantizedGranule {
        coeffs,
        side,
        scalefactors: [0; 39],
    }
}