oxideav-prores 0.0.8

//! ProRes encoder following SMPTE RDD 36 §5 + §7.
//!
//! Reads a `Yuv422P` or `Yuv444P` `VideoFrame`, walks 16x16 macroblocks,
//! forward DCTs each 8x8 block, quantises by `qmat * qScale / 8`,
//! per-component slice-scans, and emits the entropy-coded slice payload
//! per RDD 36.

use std::collections::VecDeque;

use oxideav_core::Encoder;
use oxideav_core::{
    CodecId, CodecParameters, Error, Frame, MediaType, Packet, PixelFormat, Result, TimeBase,
    VideoFrame,
};

use crate::alpha::{encode_scanned_alpha, AlphaChannelType};
use crate::dct::fdct8x8;
use crate::decoder::BitDepth;
use crate::frame::{
    compute_slice_sizes, frame_rate_code_from_rational, write_frame_with_meta,
    write_picture_header, write_slice_header, ChromaFormat, FrameMeta, Profile,
};
use crate::quant::{qscale, QuantMatrices, DEFAULT_QMAT};
use crate::slice::{blocks_per_mb, chroma_blocks_per_mb, encode_slice_components};

/// Encoder-side configuration. Defaults match the legacy behaviour
/// (flat all-4s quantisation matrices, `load_luma_qmat = 0`,
/// `load_chroma_qmat = 0`, per-profile default quantisation index).
#[derive(Clone, Debug, Default)]
pub struct EncoderConfig {
    /// Per-component quantisation weight matrices. `None` is identical
    /// to `Some(QuantMatrices::flat())` — the encoder writes
    /// `load_luma_qmat = load_chroma_qmat = 0` and uses the spec's
    /// default matrix internally for quantisation. When non-default
    /// matrices are supplied the encoder writes the matrices into the
    /// frame header (setting `load_*_qmat` to 1) so any RDD 36 decoder
    /// can dequantise correctly.
    pub quant_matrices: Option<QuantMatrices>,
    /// Per-slice `quantization_index` (RDD 36 §7.3 / Table 15) used for
    /// every slice in every encoded frame. Lower index → finer step →
    /// higher quality + larger packet. Range `1..=224`.
    ///
    /// `None` (the default) selects the per-profile default returned by
    /// [`Profile::default_quant_index`] — currently `8 / 6 / 4 / 2 /
    /// 2 / 1` for Proxy / LT / Standard / HQ / 4444 / 4444 XQ. Set this
    /// when the caller wants a different point on the rate/quality
    /// curve without re-mapping the profile selection.
    ///
    /// When `rate_control` is `true` this field is the *starting point*
    /// for the binary search; `None` uses the profile default as seed.
    pub quantization_index: Option<u8>,
    /// Descriptive metadata fields written into the RDD 36 frame header
    /// (`aspect_ratio_information`, `frame_rate_code`,
    /// `color_primaries`, `transfer_characteristic`,
    /// `matrix_coefficients` — all per §5.1.1 / §6.2). `None` lets
    /// [`make_encoder_with_config`] derive `frame_rate_code` from
    /// `CodecParameters::frame_rate` and leave the rest at 0
    /// ("unknown"); `Some(meta)` overrides everything verbatim.
    pub meta: Option<FrameMeta>,
    /// Enable two-pass per-frame rate control.
    ///
    /// When `true` and the encoder was constructed with a
    /// `CodecParameters::bit_rate` and `CodecParameters::frame_rate`,
    /// each call to `send_frame` performs a binary search over
    /// `quantization_index` (up to [`RATE_CTRL_MAX_PASSES`] trial
    /// encodes) to hit the per-frame byte target derived from the
    /// nominal bit-rate within [`RATE_CTRL_TOLERANCE`] (5 %). The
    /// search starts from the profile default qi (or the explicit
    /// `quantization_index` if set) and respects the full 1..=224
    /// range.
    ///
    /// The overhead is bounded: at most `RATE_CTRL_MAX_PASSES` full
    /// encodes per frame. For constant-content sequences the search
    /// typically converges in 2-3 passes. Set `false` (the default)
    /// to preserve the original single-pass behaviour.
    pub rate_control: bool,
}

/// Maximum number of trial encodes per frame when rate control is active.
/// Covers the full qi range (1..=224) in log2(224) ≈ 8 steps.
pub const RATE_CTRL_MAX_PASSES: usize = 10;

/// Fractional tolerance for the rate-control target (0.05 = ±5 %).
pub const RATE_CTRL_TOLERANCE: f64 = 0.05;

impl EncoderConfig {
    /// Construct a config that emits the flat all-4s matrices and
    /// `load_*_qmat = 0` (back-compat with the pre-config encoder).
    pub fn flat() -> Self {
        Self::default()
    }

    /// Construct a config that emits perceptual JPEG-derived quant
    /// matrices (see [`QuantMatrices::perceptual`]). The matrices are
    /// written into the frame header so cross-decoders pick them up.
    pub fn perceptual() -> Self {
        Self {
            quant_matrices: Some(QuantMatrices::perceptual()),
            ..Self::default()
        }
    }

    /// Use the supplied per-component matrices. Both must have weights
    /// in `2..=63` per RDD 36 §7.3 (validated at encode time).
    pub fn with_quant_matrices(mut self, qm: QuantMatrices) -> Self {
        self.quant_matrices = Some(qm);
        self
    }

    /// Override the per-profile default `quantization_index` (RDD 36
    /// §7.3 / Table 15). Must be in `1..=224`; validated at encoder
    /// construction. Lower index = finer step = higher quality.
    pub fn with_quantization_index(mut self, qi: u8) -> Self {
        self.quantization_index = Some(qi);
        self
    }

    /// Override the descriptive frame-header metadata
    /// (aspect_ratio_information, frame_rate_code, color_primaries,
    /// transfer_characteristic, matrix_coefficients). Equivalent to
    /// setting the [`Self::meta`] field directly.
    pub fn with_meta(mut self, meta: FrameMeta) -> Self {
        self.meta = Some(meta);
        self
    }

    /// Enable two-pass per-frame rate control (see [`Self::rate_control`]).
    /// Requires `CodecParameters::bit_rate` and `frame_rate` to be set at
    /// encoder construction; silently degrades to single-pass otherwise.
    pub fn with_rate_control(mut self) -> Self {
        self.rate_control = true;
        self
    }
}

/// Default `quantization_index` used for 422 Standard. Lower = higher quality.
pub const DEFAULT_QUANT_INDEX: u8 = 4;

/// Internal cap on encoded packet size — bounds the output Vec against
/// `width * height * bytes-per-sample * a small constant`. Prevents a
/// pathological caller (e.g. a header that lies about dimensions) from
/// driving an unbounded allocation.
fn output_capacity_cap(width: u16, height: u16, chroma: ChromaFormat) -> usize {
    let pixels = width as usize * height as usize;
    // A worst-case ProRes packet for 8-bit YUV is ~10 bytes per pixel
    // before container overhead. Pad to 16 + a slack for headers.
    let bpp = match chroma {
        ChromaFormat::Y422 => 16,
        ChromaFormat::Y444 => 24,
    };
    pixels.saturating_mul(bpp).saturating_add(1 << 16)
}

const MB_SIDE_PX: usize = 16;
const SLICE_MB_WIDTH_LOG2: u8 = 3; // 8 MBs per slice (typical)

/// Pick a profile from `bit_rate` when the caller expresses a target rate.
///
/// Public so callers can preview the encoder's profile selection without
/// running an encode (the chosen profile is no longer carried in the
/// RDD 36 bitstream — it lives at the container level via FourCC).
pub fn pick_profile(chroma: ChromaFormat, bit_rate: Option<u64>) -> Profile {
    match (chroma, bit_rate) {
        (ChromaFormat::Y422, Some(br)) if br <= 70_000_000 => Profile::Proxy,
        (ChromaFormat::Y422, Some(br)) if br <= 125_000_000 => Profile::Lt,
        (ChromaFormat::Y422, Some(br)) if br <= 180_000_000 => Profile::Standard,
        (ChromaFormat::Y422, Some(_)) => Profile::Hq,
        (ChromaFormat::Y422, None) => Profile::Standard,
        (ChromaFormat::Y444, Some(br)) if br >= 400_000_000 => Profile::Prores4444Xq,
        (ChromaFormat::Y444, _) => Profile::Prores4444,
    }
}

pub fn make_encoder(params: &CodecParameters) -> Result<Box<dyn Encoder>> {
    make_encoder_with_config(params, EncoderConfig::default())
}

/// Build a ProRes encoder with explicit [`EncoderConfig`] — wires
/// optional perceptual quantisation matrices through to the frame
/// header (setting `load_luma_qmat = load_chroma_qmat = 1` when the
/// matrices differ from the spec default of all-4s).
pub fn make_encoder_with_config(
    params: &CodecParameters,
    config: EncoderConfig,
) -> Result<Box<dyn Encoder>> {
    if let Some(qm) = &config.quant_matrices {
        if !qm.weights_valid() {
            return Err(Error::invalid(
                "prores encoder: quant matrix weight outside RDD 36 range 2..=63",
            ));
        }
    }
    if let Some(qi) = config.quantization_index {
        if !(1..=224).contains(&qi) {
            return Err(Error::invalid(
                "prores encoder: EncoderConfig::quantization_index out of range \
                 (must be 1..=224 per RDD 36 §7.3 / Table 15)",
            ));
        }
    }
    let width = params
        .width
        .ok_or_else(|| Error::invalid("prores encoder: missing width"))?;
    let height = params
        .height
        .ok_or_else(|| Error::invalid("prores encoder: missing height"))?;
    let pix = params.pixel_format.unwrap_or(PixelFormat::Yuv422P);

    let (chroma, bit_depth) = match pix {
        PixelFormat::Yuv422P => (ChromaFormat::Y422, BitDepth::Eight),
        PixelFormat::Yuv444P => (ChromaFormat::Y444, BitDepth::Eight),
        PixelFormat::Yuv422P10Le => (ChromaFormat::Y422, BitDepth::Ten),
        PixelFormat::Yuv444P10Le => (ChromaFormat::Y444, BitDepth::Ten),
        PixelFormat::Yuv422P12Le => (ChromaFormat::Y422, BitDepth::Twelve),
        PixelFormat::Yuv444P12Le => (ChromaFormat::Y444, BitDepth::Twelve),
        other => {
            return Err(Error::unsupported(format!(
                "prores encoder: pixel format {other:?} not supported \
                 (expected Yuv4(2|4)4P / Yuv4(2|4)4P10Le / Yuv4(2|4)4P12Le)"
            )));
        }
    };
    let profile = pick_profile(chroma, params.bit_rate);

    let mut output_params = params.clone();
    output_params.media_type = MediaType::Video;
    output_params.codec_id = CodecId::new(super::CODEC_ID_STR);
    output_params.width = Some(width);
    output_params.height = Some(height);
    output_params.pixel_format = Some(pix);

    let quant_index = config
        .quantization_index
        .unwrap_or_else(|| profile.default_quant_index());

    // Resolve the metadata block once at construction. When the caller
    // doesn't supply an explicit `FrameMeta`, derive `frame_rate_code`
    // from `params.frame_rate` (per RDD 36 §6.2 / Table 4) and leave
    // every other field at 0 ("unknown / unspecified").
    let meta = config.meta.unwrap_or_else(|| FrameMeta {
        frame_rate_code: params.frame_rate.map_or(0, frame_rate_code_from_rational),
        ..FrameMeta::default()
    });

    // Compute per-frame byte target for rate control. We need both
    // bit_rate and frame_rate; either missing → rate control disabled.
    let target_bytes = if config.rate_control {
        if let (Some(br), Some(fr)) = (params.bit_rate, params.frame_rate) {
            if fr.num > 0 && fr.den > 0 {
                // bytes_per_frame = (bit_rate / 8) * (den / num)
                let bits_per_frame = (br * fr.den as u64).saturating_div(fr.num as u64);
                (bits_per_frame / 8) as usize
            } else {
                0
            }
        } else {
            0
        }
    } else {
        0
    };

    Ok(Box::new(ProResEncoder {
        output_params,
        width,
        height,
        chroma,
        bit_depth,
        profile,
        quant_index,
        meta,
        config,
        time_base: params
            .frame_rate
            .map_or(TimeBase::new(1, 90_000), |r| TimeBase::new(r.den, r.num)),
        target_bytes,
        pending: VecDeque::new(),
        eof: false,
    }))
}

struct ProResEncoder {
    output_params: CodecParameters,
    width: u32,
    height: u32,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    profile: Profile,
    quant_index: u8,
    meta: FrameMeta,
    config: EncoderConfig,
    time_base: TimeBase,
    /// Target bytes per frame for rate control, or 0 when disabled.
    target_bytes: usize,
    pending: VecDeque<Packet>,
    eof: bool,
}

impl Encoder for ProResEncoder {
    fn codec_id(&self) -> &CodecId {
        &self.output_params.codec_id
    }

    fn output_params(&self) -> &CodecParameters {
        &self.output_params
    }

    fn send_frame(&mut self, frame: &Frame) -> Result<()> {
        match frame {
            Frame::Video(v) => {
                let data = if self.target_bytes > 0 {
                    encode_frame_with_rate_control(
                        v,
                        self.width,
                        self.height,
                        self.chroma,
                        self.bit_depth,
                        self.profile,
                        self.quant_index,
                        self.config.quant_matrices,
                        self.meta,
                        self.target_bytes,
                    )?
                } else {
                    encode_frame_full(
                        v,
                        self.width,
                        self.height,
                        self.chroma,
                        self.bit_depth,
                        self.profile,
                        self.quant_index,
                        None,
                        0,
                        self.config.quant_matrices,
                        self.meta,
                    )?
                };
                let mut pkt = Packet::new(0, self.time_base, data);
                pkt.pts = v.pts;
                pkt.dts = v.pts;
                pkt.flags.keyframe = true;
                self.pending.push_back(pkt);
                Ok(())
            }
            _ => Err(Error::invalid("prores encoder: video frames only")),
        }
    }

    fn receive_packet(&mut self) -> Result<Packet> {
        self.pending.pop_front().ok_or(Error::NeedMore)
    }

    fn flush(&mut self) -> Result<()> {
        self.eof = true;
        Ok(())
    }
}

/// Back-compat wrapper that encodes a 4:2:2 frame with the same API
/// shape as the pre-RDD 36 implementation.
pub fn encode_frame_422(
    frame: &VideoFrame,
    width: u32,
    height: u32,
    profile: Profile,
    quant_index: u8,
) -> Result<Vec<u8>> {
    encode_frame(
        frame,
        width,
        height,
        ChromaFormat::Y422,
        profile,
        quant_index,
    )
}

/// Encode a single picture (4:2:2 or 4:4:4) to a complete RDD 36 frame.
/// 8-bit input only; for 10-bit see [`encode_frame_with_depth`].
pub fn encode_frame(
    frame: &VideoFrame,
    img_w: u32,
    img_h: u32,
    chroma: ChromaFormat,
    profile: Profile,
    quantization_index: u8,
) -> Result<Vec<u8>> {
    encode_frame_with_depth(
        frame,
        img_w,
        img_h,
        chroma,
        BitDepth::Eight,
        profile,
        quantization_index,
    )
}

/// Encode a single picture to an RDD 36 frame at the requested bit depth.
///
/// `BitDepth::Eight` reads each sample as one byte; the deeper-bit paths
/// read each sample as a little-endian `u16` whose value is bounded by
/// the depth (`[0, 1023]` for 10-bit, `[0, 4095]` for 12-bit; high bits
/// ignored). Internal DCT precision is the same for all depths — input
/// samples are level-shifted into the spec's centred range
/// `v = s / 2^(b-9) - 256` (RDD 36 §7.5.1) so the quant-matrix and
/// qScale tables apply identically across depths.
#[allow(clippy::too_many_arguments)]
pub fn encode_frame_with_depth(
    frame: &VideoFrame,
    img_w: u32,
    img_h: u32,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    profile: Profile,
    quantization_index: u8,
) -> Result<Vec<u8>> {
    encode_frame_with_alpha(
        frame,
        img_w,
        img_h,
        chroma,
        bit_depth,
        profile,
        quantization_index,
        None,
    )
}

/// Encode a single picture using explicit per-component quantisation
/// weight matrices (RDD 36 §7.3). When `qmats` differs from the spec
/// default of all-4s the encoder loads the matrices into the frame
/// header (`load_luma_qmat = load_chroma_qmat = 1`) so any RDD 36
/// decoder reconstructs them correctly.
///
/// Equivalent to [`encode_frame_with_depth`] when
/// `qmats == QuantMatrices::flat()`.
#[allow(clippy::too_many_arguments)]
pub fn encode_frame_with_qmats(
    frame: &VideoFrame,
    img_w: u32,
    img_h: u32,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    profile: Profile,
    quantization_index: u8,
    qmats: QuantMatrices,
) -> Result<Vec<u8>> {
    encode_frame_full(
        frame,
        img_w,
        img_h,
        chroma,
        bit_depth,
        profile,
        quantization_index,
        None,
        0,
        Some(qmats),
        FrameMeta::default(),
    )
}

/// Encode a single picture to an RDD 36 frame with optional alpha
/// channel coding (RDD 36 §5.3.3 + §7.1.2).
///
/// When `alpha_channel_type` is `Some`, the input frame must carry a
/// 4th `VideoPlane` with a per-pixel alpha array at full luma resolution.
/// Each sample is read as one byte (`Eight`) — alpha values are
/// promoted to the spec's 16-bit internal representation and emitted as
/// `scanned_alpha()` blobs at the tail of every slice. The frame header
/// `alpha_channel_type` field is set accordingly.
///
/// When `alpha_channel_type` is `None`, behaviour is identical to
/// [`encode_frame_with_depth`] (3-plane input, no alpha emission).
#[allow(clippy::too_many_arguments)]
pub fn encode_frame_with_alpha(
    frame: &VideoFrame,
    img_w: u32,
    img_h: u32,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    profile: Profile,
    quantization_index: u8,
    alpha_channel_type: Option<AlphaChannelType>,
) -> Result<Vec<u8>> {
    encode_frame_full(
        frame,
        img_w,
        img_h,
        chroma,
        bit_depth,
        profile,
        quantization_index,
        alpha_channel_type,
        0,
        None,
        FrameMeta::default(),
    )
}

/// Encode an interlaced RDD 36 frame. `interlace_mode` selects the
/// field order (1 = top-field-first, 2 = bottom-field-first). The
/// supplied frame's planes are sliced into top + bottom field pictures
/// per §7.5.3 (rows {0, 2, …} → top, rows {1, 3, …} → bottom) and each
/// field is encoded as a separate `picture()` per §5.1, sharing one
/// frame_header().
#[allow(clippy::too_many_arguments)]
pub fn encode_frame_interlaced(
    frame: &VideoFrame,
    img_w: u32,
    img_h: u32,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    profile: Profile,
    quantization_index: u8,
    alpha_channel_type: Option<AlphaChannelType>,
    interlace_mode: u8,
) -> Result<Vec<u8>> {
    if interlace_mode != 1 && interlace_mode != 2 {
        return Err(Error::invalid(
            "prores encoder: encode_frame_interlaced requires interlace_mode in {1, 2}",
        ));
    }
    encode_frame_full(
        frame,
        img_w,
        img_h,
        chroma,
        bit_depth,
        profile,
        quantization_index,
        alpha_channel_type,
        interlace_mode,
        None,
        FrameMeta::default(),
    )
}

/// Two-pass per-frame rate control: binary-search `quantization_index` to
/// hit `target_bytes` within [`RATE_CTRL_TOLERANCE`] (±5 %).
///
/// Strategy:
/// 1. Encode once at `seed_qi` (the profile default or caller's qi).
/// 2. If the size is already within tolerance, return immediately.
/// 3. Binary-search the qi range [1, 224], converging in at most
///    [`RATE_CTRL_MAX_PASSES`] further trials.
///
/// Invariant: larger qi → coarser quantisation → smaller frame.
/// So `lo` is the qi that produced the largest recent frame and `hi`
/// is the qi that produced the smallest recent frame. We pick midpoints
/// until the target is hit or the range collapses.
#[allow(clippy::too_many_arguments)]
fn encode_frame_with_rate_control(
    frame: &VideoFrame,
    img_w: u32,
    img_h: u32,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    profile: Profile,
    seed_qi: u8,
    qmats: Option<QuantMatrices>,
    meta: FrameMeta,
    target_bytes: usize,
) -> Result<Vec<u8>> {
    let tol_lo = (target_bytes as f64 * (1.0 - RATE_CTRL_TOLERANCE)) as usize;
    let tol_hi = (target_bytes as f64 * (1.0 + RATE_CTRL_TOLERANCE)) as usize;

    // First encode at seed qi.
    let seed = encode_frame_full(
        frame, img_w, img_h, chroma, bit_depth, profile, seed_qi, None, 0, qmats, meta,
    )?;
    if seed.len() >= tol_lo && seed.len() <= tol_hi {
        return Ok(seed);
    }

    // Decide search direction.
    // If seed is too large (above target+tol) we need higher qi (coarser).
    // If seed is too small (below target-tol) we need lower qi (finer).
    let (mut lo, mut hi): (u8, u8) = if seed.len() > tol_hi {
        // Too large → need coarser quantisation → higher qi
        (seed_qi, 224)
    } else {
        // Too small → need finer quantisation → lower qi
        (1, seed_qi)
    };

    let mut best = seed;

    for _ in 0..RATE_CTRL_MAX_PASSES {
        if lo >= hi {
            break;
        }
        let mid = lo + (hi - lo) / 2;
        let candidate = encode_frame_full(
            frame, img_w, img_h, chroma, bit_depth, profile, mid, None, 0, qmats, meta,
        )?;
        let sz = candidate.len();
        if sz >= tol_lo && sz <= tol_hi {
            return Ok(candidate);
        }
        // Track the closest candidate by absolute distance to target.
        let best_dist = (best.len() as i64 - target_bytes as i64).unsigned_abs();
        let cand_dist = (sz as i64 - target_bytes as i64).unsigned_abs();
        if cand_dist < best_dist {
            best = candidate;
        }
        if sz > tol_hi {
            // Frame too large → raise qi (coarser)
            lo = mid + 1;
        } else {
            // Frame too small → lower qi (finer)
            // Safe because mid >= lo >= 1; if mid == 1 the loop exits.
            if mid == 0 {
                break;
            }
            hi = mid - 1;
        }
    }
    // Return the best candidate found (closest to target).
    Ok(best)
}

/// Internal entrypoint shared by progressive and interlaced encodes.
/// `interlace_mode == 0` builds a single picture; `1` (TFF) or `2`
/// (BFF) builds two field pictures per §5.1.
///
/// `qmats == None` reproduces the legacy behaviour: flat all-4s
/// matrices, `load_luma_qmat = load_chroma_qmat = 0`, frame_header_size
/// = 20. `qmats == Some(QuantMatrices::flat())` is treated identically
/// (no point loading the default matrix into the bitstream). For any
/// other matrices, the encoder writes `load_luma_qmat = 1` and (when
/// the chroma matrix differs from the luma matrix) `load_chroma_qmat
/// = 1`, growing the frame header by 64 or 128 bytes per §7.3.
#[allow(clippy::too_many_arguments)]
fn encode_frame_full(
    frame: &VideoFrame,
    img_w: u32,
    img_h: u32,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    profile: Profile,
    quantization_index: u8,
    alpha_channel_type: Option<AlphaChannelType>,
    interlace_mode: u8,
    qmats: Option<QuantMatrices>,
    meta: FrameMeta,
) -> Result<Vec<u8>> {
    let expected_planes = if alpha_channel_type.is_some() { 4 } else { 3 };
    if frame.planes.len() != expected_planes {
        return Err(Error::invalid(format!(
            "prores encoder: expected {expected_planes} planes (got {})",
            frame.planes.len()
        )));
    }
    if !(1..=224).contains(&quantization_index) {
        return Err(Error::invalid(
            "prores encoder: quantization_index out of range",
        ));
    }
    if profile.chroma_format() != chroma {
        return Err(Error::invalid(
            "prores encoder: profile chroma_format does not match requested chroma",
        ));
    }
    if let Some(qm) = &qmats {
        if !qm.weights_valid() {
            return Err(Error::invalid(
                "prores encoder: quant matrix weight outside RDD 36 range 2..=63",
            ));
        }
    }
    let width = img_w as usize;
    let height = img_h as usize;

    // Bound output capacity against header-declared dimensions.
    let cap = output_capacity_cap(img_w as u16, img_h as u16, chroma);

    // Resolve the per-component matrices used for quantisation. When
    // the caller passed flat (or no) matrices we keep load_*_qmat = 0
    // for byte-exact compatibility with the pre-config encoder.
    let qmat_pair = qmats.unwrap_or_default();
    let load_luma = !qmat_pair.is_default();
    let load_chroma = load_luma && qmat_pair.chroma != qmat_pair.luma;
    let luma_qmat = if load_luma {
        &qmat_pair.luma
    } else {
        &DEFAULT_QMAT
    };
    let chroma_qmat = if load_luma {
        &qmat_pair.chroma
    } else {
        &DEFAULT_QMAT
    };

    // Per §6.2 picture_vertical_size derivation. Each interlaced field
    // is a separate picture sized at half the frame height (rounded
    // appropriately for top vs. bottom).
    let pictures: Vec<(usize, FieldStride)> = if interlace_mode == 0 {
        vec![(height, FieldStride::progressive())]
    } else {
        let top_h = height.div_ceil(2);
        let bot_h = height / 2;
        // interlace_mode 1: first picture is top field (offset 0)
        // interlace_mode 2: first picture is bottom field (offset 1)
        if interlace_mode == 1 {
            vec![
                (top_h, FieldStride::new(2, 0)),
                (bot_h, FieldStride::new(2, 1)),
            ]
        } else {
            vec![
                (bot_h, FieldStride::new(2, 1)),
                (top_h, FieldStride::new(2, 0)),
            ]
        }
    };

    let interlaced = interlace_mode != 0;
    let mut picture_blobs: Vec<Vec<u8>> = Vec::with_capacity(pictures.len());
    for (picture_height, field) in &pictures {
        let blob = encode_one_picture(
            frame,
            width,
            height,
            *picture_height,
            chroma,
            bit_depth,
            quantization_index,
            luma_qmat,
            chroma_qmat,
            alpha_channel_type,
            interlaced,
            *field,
        )?;
        picture_blobs.push(blob);
    }

    // Per §5.1.1 frame_header_size: 20 + 64 (load_luma) + 64 (load_chroma).
    let frame_header_size =
        20usize + if load_luma { 64 } else { 0 } + if load_chroma { 64 } else { 0 };
    let pictures_total: usize = picture_blobs.iter().map(|p| p.len()).sum();
    let total_frame_size_no_padding = 4 + 4 + frame_header_size + pictures_total;
    if total_frame_size_no_padding > cap {
        return Err(Error::invalid(
            "prores encoder: encoded size exceeds internal cap",
        ));
    }

    let mut out = Vec::with_capacity(total_frame_size_no_padding);
    write_frame_with_meta(
        &mut out,
        total_frame_size_no_padding as u32,
        img_w as u16,
        img_h as u16,
        chroma,
        interlace_mode,
        luma_qmat,
        chroma_qmat,
        load_luma,
        load_chroma,
        alpha_channel_type.map_or(0, |a| a.code()),
        meta,
    );
    for blob in &picture_blobs {
        out.extend_from_slice(blob);
    }
    debug_assert_eq!(out.len(), total_frame_size_no_padding);
    Ok(out)
}

/// Field-row mapping for source-plane reads on the encoder side.
/// Mirrors `decoder::FieldStride`.
#[derive(Copy, Clone, Debug)]
struct FieldStride {
    step: usize,
    offset: usize,
}

impl FieldStride {
    fn new(step: usize, offset: usize) -> Self {
        Self { step, offset }
    }
    fn progressive() -> Self {
        Self { step: 1, offset: 0 }
    }
    fn map(self, picture_row: usize) -> usize {
        self.step * picture_row + self.offset
    }
}

/// Build one `picture()` blob (picture_header + slice_table +
/// concatenated slice payloads). For interlaced encodes the caller
/// invokes this twice (once per field).
///
/// `luma_qmat` is applied to all four luma blocks per macroblock;
/// `chroma_qmat` is applied to both Cb and Cr blocks. Both matrices
/// must be the same matrices written into the frame header so the
/// decoder dequantises with the matching W[][].
#[allow(clippy::too_many_arguments)]
fn encode_one_picture(
    frame: &VideoFrame,
    frame_w: usize,
    frame_h: usize,
    picture_height: usize,
    chroma: ChromaFormat,
    bit_depth: BitDepth,
    quantization_index: u8,
    luma_qmat: &[u8; 64],
    chroma_qmat: &[u8; 64],
    alpha_channel_type: Option<AlphaChannelType>,
    interlaced: bool,
    field: FieldStride,
) -> Result<Vec<u8>> {
    let c_w = match chroma {
        ChromaFormat::Y422 => frame_w.div_ceil(2),
        ChromaFormat::Y444 => frame_w,
    };
    let mbs_x = frame_w.div_ceil(MB_SIDE_PX);
    let mbs_y = picture_height.div_ceil(MB_SIDE_PX);
    let slice_sizes_template = compute_slice_sizes(mbs_x, SLICE_MB_WIDTH_LOG2);
    let slices_per_row = slice_sizes_template.len();
    let slice_count = slices_per_row * mbs_y;
    let _cb_per_mb = chroma_blocks_per_mb(chroma);
    let per_mb = blocks_per_mb(chroma);

    const LUMA_OFFSETS: [(usize, usize); 4] = [(0, 0), (1, 0), (0, 1), (1, 1)];
    let chroma_offsets: &[(usize, usize)] = match chroma {
        ChromaFormat::Y422 => &[(0, 0), (0, 1)],
        ChromaFormat::Y444 => &LUMA_OFFSETS,
    };

    let mut slice_payloads: Vec<Vec<u8>> = Vec::with_capacity(slice_count);
    for my in 0..mbs_y {
        let mut mx = 0usize;
        for &mbs_this_slice in &slice_sizes_template {
            let mbs_this_slice = mbs_this_slice.min(mbs_x - mx);
            if mbs_this_slice == 0 {
                break;
            }
            let mut blocks: Vec<[i32; 64]> = Vec::with_capacity(mbs_this_slice * per_mb);
            for mb_within in 0..mbs_this_slice {
                let mb_x = mx + mb_within;
                for (bx, by) in LUMA_OFFSETS {
                    let x0 = mb_x * MB_SIDE_PX + bx * 8;
                    let y0 = my * MB_SIDE_PX + by * 8;
                    blocks.push(encode_block(
                        &frame.planes[0].data,
                        frame.planes[0].stride,
                        frame_w,
                        frame_h,
                        x0,
                        y0,
                        luma_qmat,
                        quantization_index,
                        bit_depth,
                        field,
                    ));
                }
                for plane_idx in [1usize, 2] {
                    for (bx, by) in chroma_offsets.iter().copied() {
                        let (x0, y0) = match chroma {
                            ChromaFormat::Y422 => (mb_x * 8, my * MB_SIDE_PX + by * 8),
                            ChromaFormat::Y444 => {
                                (mb_x * MB_SIDE_PX + bx * 8, my * MB_SIDE_PX + by * 8)
                            }
                        };
                        blocks.push(encode_block(
                            &frame.planes[plane_idx].data,
                            frame.planes[plane_idx].stride,
                            c_w,
                            frame_h,
                            x0,
                            y0,
                            chroma_qmat,
                            quantization_index,
                            bit_depth,
                            field,
                        ));
                    }
                }
            }
            let (y_data, cb_data, cr_data) =
                encode_slice_components(mbs_this_slice, chroma, interlaced, &blocks)?;
            if y_data.len() > u16::MAX as usize
                || cb_data.len() > u16::MAX as usize
                || cr_data.len() > u16::MAX as usize
            {
                return Err(Error::invalid(
                    "prores encoder: slice component exceeded u16 size limit",
                ));
            }

            let alpha_blob: Vec<u8> = if let Some(act) = alpha_channel_type {
                // Emit alpha for the FULL macroblock-row height (16
                // sample rows) regardless of visible picture clipping.
                // Decoders MUST allocate the padded MB-aligned plane
                // and crop after decode (RDD 36 §7.5.2 — alphaValues is
                // the padded picture size); ffmpeg's prores_ks behaves
                // the same way. Edge-pixels for the partially-visible
                // last MB row are clamped to the last visible row so
                // the stream stays self-roundtrippable.
                let slice_vertical_size = MB_SIDE_PX;
                let cols = MB_SIDE_PX * mbs_this_slice;
                let mut samples: Vec<u16> = Vec::with_capacity(cols * slice_vertical_size);
                let a_plane = &frame.planes[3];
                let a_stride = a_plane.stride;
                for r in 0..slice_vertical_size {
                    let frame_row = field
                        .map(my * MB_SIDE_PX + r)
                        .min(frame_h.saturating_sub(1));
                    for c in 0..cols {
                        let x = (mx * MB_SIDE_PX + c).min(frame_w.saturating_sub(1));
                        let v: u16 = match act {
                            AlphaChannelType::Eight => {
                                a_plane.data[frame_row * a_stride + x] as u16
                            }
                            AlphaChannelType::Sixteen => {
                                let off = frame_row * a_stride + x * 2;
                                u16::from_le_bytes([a_plane.data[off], a_plane.data[off + 1]])
                            }
                        };
                        samples.push(v);
                    }
                }
                encode_scanned_alpha(&samples, act)?
            } else {
                Vec::new()
            };

            let cr_field = if alpha_channel_type.is_some() {
                Some(cr_data.len() as u16)
            } else {
                None
            };
            let mut slice_buf = Vec::with_capacity(
                8 + y_data.len() + cb_data.len() + cr_data.len() + alpha_blob.len(),
            );
            write_slice_header(
                &mut slice_buf,
                quantization_index,
                y_data.len() as u16,
                cb_data.len() as u16,
                cr_field,
            );
            slice_buf.extend_from_slice(&y_data);
            slice_buf.extend_from_slice(&cb_data);
            slice_buf.extend_from_slice(&cr_data);
            slice_buf.extend_from_slice(&alpha_blob);
            slice_payloads.push(slice_buf);
            mx += mbs_this_slice;
        }
    }
    debug_assert_eq!(slice_payloads.len(), slice_count);
    if slice_payloads.iter().any(|p| p.len() > u16::MAX as usize) {
        return Err(Error::invalid(
            "prores encoder: slice exceeded u16 size table limit",
        ));
    }

    let slice_table_size = slice_count * 2;
    let slice_bytes: usize = slice_payloads.iter().map(|p| p.len()).sum();
    let picture_header_size = 8usize;
    let picture_size = (picture_header_size + slice_table_size + slice_bytes) as u32;
    let mut blob = Vec::with_capacity(picture_size as usize);
    write_picture_header(
        &mut blob,
        picture_size,
        if slice_count <= u16::MAX as usize {
            slice_count as u16
        } else {
            0
        },
        SLICE_MB_WIDTH_LOG2,
    );
    for p in &slice_payloads {
        blob.extend_from_slice(&(p.len() as u16).to_be_bytes());
    }
    for p in &slice_payloads {
        blob.extend_from_slice(p);
    }
    debug_assert_eq!(blob.len(), picture_size as usize);
    Ok(blob)
}

/// Sample one IDCT input value from the source plane at sample
/// coordinate `(x, y)`, applying the spec's level-shift to a centred
/// `v` in the range `[-256, 256)` per RDD 36 §7.5.1. The inverse of
/// the decoder formula `s = 2^b * (v + 256) / 512` is
/// `v = s * 512 / 2^b - 256 = s / 2^(b-9) - 256`. `stride` is in
/// **bytes**; for 10/12-bit planes that's `2 * samples_per_row`.
fn read_sample(plane: &[u8], stride: usize, x: usize, y: usize, bit_depth: BitDepth) -> f32 {
    match bit_depth {
        BitDepth::Eight => (plane[y * stride + x] as f32) * 2.0 - 256.0,
        BitDepth::Ten => {
            let off = y * stride + x * 2;
            let lo = plane[off] as u16;
            let hi = plane[off + 1] as u16;
            let s = (lo | (hi << 8)) & 0x03FF;
            (s as f32) / 2.0 - 256.0
        }
        BitDepth::Twelve => {
            let off = y * stride + x * 2;
            let lo = plane[off] as u16;
            let hi = plane[off + 1] as u16;
            let s = (lo | (hi << 8)) & 0x0FFF;
            (s as f32) / 8.0 - 256.0
        }
    }
}

#[allow(clippy::too_many_arguments)]
fn encode_block(
    plane: &[u8],
    stride: usize,
    plane_w: usize,
    plane_h: usize,
    x0: usize,
    y0: usize,
    qmat: &[u8; 64],
    quantization_index: u8,
    bit_depth: BitDepth,
    field: FieldStride,
) -> [i32; 64] {
    let mut blk = [0.0f32; 64];
    for j in 0..8 {
        // Map per-picture row to per-frame row; this is the identity for
        // progressive (`step=1, offset=0`) and 2*r+offset for interlaced.
        let frame_row = field.map(y0 + j).min(plane_h.saturating_sub(1));
        for i in 0..8 {
            let x = (x0 + i).min(plane_w.saturating_sub(1));
            blk[j * 8 + i] = read_sample(plane, stride, x, frame_row, bit_depth);
        }
    }
    fdct8x8(&mut blk);
    // Quantisation: F[v][u] = (QF[v][u] * W[v][u] * qScale) / 8
    // Inverse: QF = round(F * 8 / (W * qScale)).
    let qs = qscale(quantization_index) as f32;
    let mut out = [0i32; 64];
    for k in 0..64 {
        let denom = qmat[k] as f32 * qs;
        let v = blk[k] * 8.0 / denom;
        out[k] = if v >= 0.0 {
            (v + 0.5) as i32
        } else {
            -((-v + 0.5) as i32)
        };
    }
    out
}