oximedia_codec/av1/

coeff_decode.rs

//! AV1 transform coefficient decoding.
//!
//! This module handles the complete decoding of transform coefficients from
//! the entropy-coded bitstream, including:
//!
//! - EOB (End of Block) position parsing
//! - Coefficient level decoding using multi-level scheme
//! - Coefficient sign decoding
//! - Dequantization
//! - Scan order application
//!
//! # Coefficient Decoding Process
//!
//! 1. **EOB Parsing** - Determine position of last non-zero coefficient
//! 2. **Coefficient Levels** - Decode base levels and ranges
//! 3. **DC Sign** - Decode sign of DC coefficient
//! 4. **AC Signs** - Decode signs of AC coefficients
//! 5. **Dequantization** - Apply quantization parameters
//! 6. **Scan Order** - Convert from scan order to raster order
//!
//! # Context Modeling
//!
//! Coefficient decoding uses adaptive context models based on:
//! - Position within the block
//! - Magnitude of neighboring coefficients
//! - Transform size and type

#![forbid(unsafe_code)]
#![allow(dead_code)]
#![allow(clippy::doc_markdown)]
#![allow(clippy::too_many_arguments)]
#![allow(clippy::cast_possible_truncation)]
#![allow(clippy::cast_sign_loss)]
#![allow(clippy::cast_possible_wrap)]
#![allow(clippy::similar_names)]
#![allow(clippy::module_name_repetitions)]

use super::coefficients::{
    dequantize_block, get_dequant_shift, CoeffBuffer, CoeffContext, CoeffStats, EobContext, EobPt,
    LevelContext, ScanOrderCache,
};
use super::entropy::SymbolReader;
use super::entropy_tables::CdfContext;
use super::quantization::QuantizationParams;
use super::transform::{TxSize, TxType};
use crate::error::CodecResult;

// =============================================================================
// Constants
// =============================================================================

/// Maximum coefficient level for base coding.
pub const COEFF_BASE_MAX: u32 = 3;

/// Coefficient base range (used for higher levels).
pub const BR_CDF_SIZE: usize = 4;

/// Maximum Golomb-Rice parameter.
pub const MAX_BR_PARAM: u8 = 5;

/// Number of EOB offset bits.
pub const EOB_OFFSET_BITS: [u8; 12] = [0, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9];

/// Scan order coeff skip threshold.
pub const COEFF_SKIP_THRESHOLD: u16 = 256;

// =============================================================================
// Coefficient Decoder
// =============================================================================

/// Decoder for transform coefficients.
#[derive(Debug)]
pub struct CoeffDecoder {
    /// Symbol reader.
    reader: SymbolReader,
    /// CDF context for probability models.
    cdf_context: CdfContext,
    /// Scan order cache.
    scan_cache: ScanOrderCache,
    /// Quantization parameters.
    quant_params: QuantizationParams,
    /// Current bit depth.
    bit_depth: u8,
}

impl CoeffDecoder {
    /// Create a new coefficient decoder.
    pub fn new(data: Vec<u8>, quant_params: QuantizationParams, bit_depth: u8) -> Self {
        Self {
            reader: SymbolReader::new(data),
            cdf_context: CdfContext::new(),
            scan_cache: ScanOrderCache::new(),
            quant_params,
            bit_depth,
        }
    }

    /// Decode coefficients for a transform block.
    pub fn decode_coefficients(
        &mut self,
        tx_size: TxSize,
        tx_type: TxType,
        plane: u8,
        skip: bool,
    ) -> CodecResult<CoeffBuffer> {
        let mut ctx = CoeffContext::new(tx_size, tx_type, plane);

        if skip {
            // Skip blocks have all-zero coefficients
            return Ok(CoeffBuffer::from_tx_size(tx_size));
        }

        // Decode EOB position
        ctx.eob = self.decode_eob(tx_size, plane)?;

        if ctx.eob == 0 {
            // No coefficients
            return Ok(CoeffBuffer::from_tx_size(tx_size));
        }

        // Get scan order (clone to avoid borrow issues)
        let scan = self.scan_cache.get(tx_size, ctx.tx_class()).to_vec();

        // Decode coefficient levels
        self.decode_coeff_levels(&mut ctx, &scan)?;

        // Decode signs
        self.decode_signs(&mut ctx, &scan)?;

        // Dequantize coefficients
        self.dequantize_coefficients(&mut ctx, plane)?;

        // Convert from scan order to raster order
        let mut buffer = CoeffBuffer::from_tx_size(tx_size);
        self.reorder_coefficients(&ctx, &scan, &mut buffer)?;

        Ok(buffer)
    }

    /// Decode EOB (End of Block) position.
    fn decode_eob(&mut self, tx_size: TxSize, plane: u8) -> CodecResult<u16> {
        let _eob_ctx = EobContext::new(tx_size);

        // Read EOB multi-symbol
        let ctx = (tx_size as usize * 3) + (plane as usize);
        let eob_multi_cdf = self.cdf_context.get_eob_multi_cdf_mut(ctx);
        let eob_multi = self.reader.read_symbol(eob_multi_cdf) as u8;

        if eob_multi == 0 {
            return Ok(0); // No coefficients
        }

        // Read EOB extra bits
        let eob_pt = EobPt::from_eob(eob_multi.into());
        let extra_bits = eob_pt.extra_bits();

        let eob_extra = if extra_bits > 0 {
            self.reader.read_literal(extra_bits)
        } else {
            0
        };

        let eob = EobContext::compute_eob(eob_multi, eob_extra as u16);
        Ok(eob)
    }

    /// Decode coefficient levels.
    fn decode_coeff_levels(&mut self, ctx: &mut CoeffContext, scan: &[u16]) -> CodecResult<()> {
        let eob = ctx.eob as usize;

        // Decode in reverse scan order (from EOB to DC)
        for scan_idx in (0..eob).rev() {
            let pos = scan[scan_idx] as usize;
            let level = self.decode_coeff_level(ctx, pos, scan_idx == eob - 1)?;
            ctx.levels[pos] = level as i32;
        }

        Ok(())
    }

    /// Decode a single coefficient level.
    fn decode_coeff_level(
        &mut self,
        ctx: &CoeffContext,
        pos: usize,
        is_eob: bool,
    ) -> CodecResult<u32> {
        let level_ctx = ctx.compute_level_context(pos);

        // Decode base level (0-3)
        let base_level = if is_eob {
            // At EOB, coefficient is at least 1
            1 + self.decode_coeff_base_eob(&level_ctx, ctx.plane)?
        } else {
            self.decode_coeff_base(&level_ctx, ctx.plane)?
        };

        if base_level >= COEFF_BASE_MAX {
            // Decode additional range
            let range = self.decode_coeff_base_range(&level_ctx, ctx.plane)?;
            Ok(base_level + range)
        } else {
            Ok(base_level)
        }
    }

    /// Decode coefficient base level.
    fn decode_coeff_base(&mut self, level_ctx: &LevelContext, _plane: u8) -> CodecResult<u32> {
        let context = level_ctx.context() as usize;
        let cdf = self.cdf_context.get_coeff_base_cdf_mut(context);
        Ok(self.reader.read_symbol(cdf) as u32)
    }

    /// Decode coefficient base at EOB.
    fn decode_coeff_base_eob(&mut self, level_ctx: &LevelContext, _plane: u8) -> CodecResult<u32> {
        let context = level_ctx.context() as usize;
        let cdf = self.cdf_context.get_coeff_base_eob_cdf_mut(context);
        Ok(self.reader.read_symbol(cdf) as u32)
    }

    /// Decode coefficient base range for high magnitude coefficients.
    fn decode_coeff_base_range(
        &mut self,
        level_ctx: &LevelContext,
        _plane: u8,
    ) -> CodecResult<u32> {
        let context = level_ctx.mag_context() as usize;
        let mut total_range = 0u32;

        // Multi-level Golomb-Rice coding
        for _level in 0..5 {
            let br_cdf = self.cdf_context.get_coeff_br_cdf_mut(context);
            let br_symbol = self.reader.read_symbol(br_cdf);

            if br_symbol < BR_CDF_SIZE - 1 {
                total_range += br_symbol as u32;
                break;
            } else {
                total_range += (BR_CDF_SIZE - 1) as u32;
            }
        }

        Ok(total_range)
    }

    /// Decode signs for all non-zero coefficients.
    fn decode_signs(&mut self, ctx: &mut CoeffContext, scan: &[u16]) -> CodecResult<()> {
        let eob = ctx.eob as usize;

        // Decode DC sign first (if DC is non-zero)
        if ctx.levels[0] != 0 {
            let dc_sign_ctx = ctx.dc_sign_context();
            let cdf_slice = self.cdf_context.get_dc_sign_cdf_mut(dc_sign_ctx as usize);
            // read_bool needs &mut [u16; 3], but we have &mut [u16]
            // Copy to a fixed-size array
            if cdf_slice.len() >= 3 {
                let mut cdf_array = [cdf_slice[0], cdf_slice[1], cdf_slice[2]];
                let sign = self.reader.read_bool(&mut cdf_array);
                // Copy back updated CDF
                cdf_slice[0] = cdf_array[0];
                cdf_slice[1] = cdf_array[1];
                cdf_slice[2] = cdf_array[2];
                ctx.signs[0] = sign;
                if sign {
                    ctx.levels[0] = -ctx.levels[0];
                }
            }
        }

        // Decode AC signs
        for scan_idx in 1..eob {
            let pos = scan[scan_idx] as usize;
            if ctx.levels[pos] != 0 {
                // AC signs use equiprobable model
                let sign = self.reader.read_bool_eq();
                ctx.signs[pos] = sign;
                if sign {
                    ctx.levels[pos] = -ctx.levels[pos];
                }
            }
        }

        Ok(())
    }

    /// Dequantize coefficients.
    fn dequantize_coefficients(&mut self, ctx: &mut CoeffContext, plane: u8) -> CodecResult<()> {
        // Get dequant values
        let dc_dequant = self
            .quant_params
            .get_dc_quant(plane as usize, self.bit_depth) as i16;
        let ac_dequant = self
            .quant_params
            .get_ac_quant(plane as usize, self.bit_depth) as i16;
        let shift = get_dequant_shift(self.bit_depth);

        dequantize_block(&mut ctx.levels, dc_dequant, ac_dequant, shift);

        Ok(())
    }

    /// Reorder coefficients from scan order to raster order.
    fn reorder_coefficients(
        &self,
        ctx: &CoeffContext,
        scan: &[u16],
        buffer: &mut CoeffBuffer,
    ) -> CodecResult<()> {
        let eob = ctx.eob as usize;

        for scan_idx in 0..eob {
            let pos = scan[scan_idx] as usize;
            if pos < ctx.levels.len() {
                let level = ctx.levels[pos];
                let (row, col) = ctx.get_scan_position(pos);
                buffer.set(row as usize, col as usize, level);
            }
        }

        Ok(())
    }

    /// Check if more data is available.
    pub fn has_more_data(&self) -> bool {
        self.reader.has_more_data()
    }

    /// Get current position.
    pub fn position(&self) -> usize {
        self.reader.position()
    }
}

// =============================================================================
// Coefficient Encoding
// =============================================================================

/// Encoder for transform coefficients.
#[derive(Debug)]
pub struct CoeffEncoder {
    /// Symbol writer.
    writer: super::entropy::SymbolWriter,
    /// CDF context.
    cdf_context: CdfContext,
    /// Scan order cache.
    scan_cache: ScanOrderCache,
}

impl CoeffEncoder {
    /// Create a new coefficient encoder.
    #[must_use]
    pub fn new() -> Self {
        Self {
            writer: super::entropy::SymbolWriter::new(),
            cdf_context: CdfContext::new(),
            scan_cache: ScanOrderCache::new(),
        }
    }

    /// Encode coefficients for a transform block.
    pub fn encode_coefficients(
        &mut self,
        buffer: &CoeffBuffer,
        tx_size: TxSize,
        tx_type: TxType,
        plane: u8,
    ) -> CodecResult<()> {
        let tx_class = tx_type.tx_class();
        let scan = self.scan_cache.get(tx_size, tx_class).to_vec();

        // Find EOB
        let eob = self.find_eob(buffer, &scan);

        // Encode EOB
        self.encode_eob(eob, tx_size, plane)?;

        if eob == 0 {
            return Ok(());
        }

        // Convert to scan order
        let mut levels = vec![0i32; eob as usize];
        buffer.copy_to_scan(&mut levels, &scan[..eob as usize]);

        // Encode levels
        self.encode_levels(&levels, &scan, plane)?;

        // Encode signs
        self.encode_signs(&levels, &scan)?;

        Ok(())
    }

    /// Find EOB position.
    fn find_eob(&self, buffer: &CoeffBuffer, scan: &[u16]) -> u16 {
        for (i, &pos) in scan.iter().enumerate().rev() {
            let (row, col) = self.pos_to_rowcol(pos as usize, buffer);
            if buffer.get(row, col) != 0 {
                return (i + 1) as u16;
            }
        }
        0
    }

    /// Convert position to row/col.
    fn pos_to_rowcol(&self, pos: usize, buffer: &CoeffBuffer) -> (usize, usize) {
        let slice = buffer.as_slice();
        let width = (slice.len() as f64).sqrt() as usize;
        (pos / width, pos % width)
    }

    /// Encode EOB.
    fn encode_eob(&mut self, eob: u16, tx_size: TxSize, plane: u8) -> CodecResult<()> {
        let ctx = (tx_size as usize * 3) + (plane as usize);

        if eob == 0 {
            let cdf = self.cdf_context.get_eob_multi_cdf_mut(ctx);
            self.writer.write_symbol(0, cdf);
            return Ok(());
        }

        let eob_pt = EobPt::from_eob(eob);
        let cdf = self.cdf_context.get_eob_multi_cdf_mut(ctx);
        self.writer.write_symbol(eob_pt as usize, cdf);

        // Write extra bits
        let extra_bits = eob_pt.extra_bits();
        if extra_bits > 0 {
            let offset = eob - eob_pt.base_eob();
            self.writer.write_literal(offset as u32, extra_bits);
        }

        Ok(())
    }

    /// Encode coefficient levels.
    fn encode_levels(&mut self, levels: &[i32], scan: &[u16], plane: u8) -> CodecResult<()> {
        let mut ctx = LevelContext::new();

        for (scan_idx, &_pos) in scan.iter().enumerate().rev() {
            let level = levels[scan_idx].unsigned_abs();

            let base_level = level.min(COEFF_BASE_MAX);
            let is_eob = scan_idx == levels.len() - 1;

            if is_eob {
                let cdf = self
                    .cdf_context
                    .get_coeff_base_eob_cdf_mut(ctx.context() as usize);
                self.writer.write_symbol((base_level - 1) as usize, cdf);
            } else {
                let cdf = self
                    .cdf_context
                    .get_coeff_base_cdf_mut(ctx.context() as usize);
                self.writer.write_symbol(base_level as usize, cdf);
            }

            if level >= COEFF_BASE_MAX {
                self.encode_base_range(level - COEFF_BASE_MAX, &ctx, plane)?;
            }

            // Update context
            ctx.mag += level;
            if level > 0 {
                ctx.count += 1;
            }
        }

        Ok(())
    }

    /// Encode coefficient base range.
    fn encode_base_range(&mut self, range: u32, ctx: &LevelContext, _plane: u8) -> CodecResult<()> {
        let mut remaining = range;
        let mag_ctx = ctx.mag_context() as usize;

        for _level in 0..5 {
            if remaining == 0 {
                break;
            }

            let symbol = remaining.min((BR_CDF_SIZE - 1) as u32) as usize;
            let cdf = self.cdf_context.get_coeff_br_cdf_mut(mag_ctx);
            self.writer.write_symbol(symbol, cdf);

            if symbol < BR_CDF_SIZE - 1 {
                break;
            }

            remaining -= (BR_CDF_SIZE - 1) as u32;
        }

        Ok(())
    }

    /// Encode signs.
    fn encode_signs(&mut self, levels: &[i32], scan: &[u16]) -> CodecResult<()> {
        // DC sign
        if !levels.is_empty() && levels[0] != 0 {
            let dc_ctx = 1; // Simplified context
            let cdf_slice = self.cdf_context.get_dc_sign_cdf_mut(dc_ctx);
            if cdf_slice.len() >= 3 {
                let mut cdf = [cdf_slice[0], cdf_slice[1], cdf_slice[2]];
                self.writer.write_bool(levels[0] < 0, &mut cdf);
                // Copy back updated CDF
                cdf_slice[0] = cdf[0];
                cdf_slice[1] = cdf[1];
                cdf_slice[2] = cdf[2];
            }
        }

        // AC signs
        for (idx, &_pos) in scan.iter().enumerate().skip(1) {
            if idx < levels.len() && levels[idx] != 0 {
                // Equiprobable
                let mut cdf = [16384u16, 32768, 0];
                self.writer.write_bool(levels[idx] < 0, &mut cdf);
            }
        }

        Ok(())
    }

    /// Finalize and get output.
    #[must_use]
    pub fn finish(self) -> Vec<u8> {
        self.writer.finish()
    }
}

impl Default for CoeffEncoder {
    fn default() -> Self {
        Self::new()
    }
}

// =============================================================================
// Batched Coefficient Decoder
// =============================================================================

/// Batched decoder for multiple coefficient blocks.
pub struct BatchedCoeffDecoder {
    /// Base decoder.
    decoder: CoeffDecoder,
    /// Decoded blocks cache.
    blocks: Vec<CoeffBuffer>,
}

impl BatchedCoeffDecoder {
    /// Create a new batched decoder.
    pub fn new(data: Vec<u8>, quant_params: QuantizationParams, bit_depth: u8) -> Self {
        Self {
            decoder: CoeffDecoder::new(data, quant_params, bit_depth),
            blocks: Vec::new(),
        }
    }

    /// Decode multiple blocks.
    pub fn decode_blocks(
        &mut self,
        specs: &[(TxSize, TxType, u8, bool)],
    ) -> CodecResult<Vec<CoeffBuffer>> {
        self.blocks.clear();

        for &(tx_size, tx_type, plane, skip) in specs {
            let buffer = self
                .decoder
                .decode_coefficients(tx_size, tx_type, plane, skip)?;
            self.blocks.push(buffer);
        }

        Ok(std::mem::take(&mut self.blocks))
    }

    /// Get statistics for decoded blocks.
    #[must_use]
    pub fn get_statistics(&self) -> Vec<CoeffStats> {
        self.blocks
            .iter()
            .map(|b| CoeffStats::from_coeffs(b.as_slice()))
            .collect()
    }
}

// =============================================================================
// Coefficient Analysis
// =============================================================================

/// Analyze coefficient distribution.
#[derive(Clone, Debug, Default)]
pub struct CoeffAnalysis {
    /// Total coefficients analyzed.
    pub total_coeffs: u64,
    /// Zero coefficients.
    pub zero_count: u64,
    /// Non-zero coefficients.
    pub nonzero_count: u64,
    /// DC coefficient sum.
    pub dc_sum: i64,
    /// AC coefficient sum.
    pub ac_sum: i64,
    /// Maximum absolute value.
    pub max_abs: u32,
}

impl CoeffAnalysis {
    /// Create new analysis.
    #[must_use]
    pub const fn new() -> Self {
        Self {
            total_coeffs: 0,
            zero_count: 0,
            nonzero_count: 0,
            dc_sum: 0,
            ac_sum: 0,
            max_abs: 0,
        }
    }

    /// Analyze a coefficient buffer.
    pub fn analyze(&mut self, buffer: &CoeffBuffer) {
        let coeffs = buffer.as_slice();
        self.total_coeffs += coeffs.len() as u64;

        if !coeffs.is_empty() {
            self.dc_sum += i64::from(coeffs[0]);
        }

        for (i, &coeff) in coeffs.iter().enumerate() {
            let abs_val = coeff.unsigned_abs();

            if coeff == 0 {
                self.zero_count += 1;
            } else {
                self.nonzero_count += 1;
                self.max_abs = self.max_abs.max(abs_val);

                if i > 0 {
                    self.ac_sum += i64::from(coeff);
                }
            }
        }
    }

    /// Get sparsity ratio (percentage of zeros).
    #[must_use]
    pub fn sparsity(&self) -> f64 {
        if self.total_coeffs > 0 {
            (self.zero_count as f64 / self.total_coeffs as f64) * 100.0
        } else {
            0.0
        }
    }

    /// Get average DC value.
    #[must_use]
    pub fn avg_dc(&self) -> f64 {
        if self.total_coeffs > 0 {
            self.dc_sum as f64 / self.total_coeffs as f64
        } else {
            0.0
        }
    }
}

// =============================================================================
// Tests
// =============================================================================

#[cfg(test)]
mod tests {
    use super::*;

    fn create_test_quant_params() -> QuantizationParams {
        QuantizationParams {
            base_q_idx: 100,
            delta_q_y_dc: 0,
            delta_q_u_dc: 0,
            delta_q_v_dc: 0,
            delta_q_u_ac: 0,
            delta_q_v_ac: 0,
            using_qmatrix: false,
            qm_y: 15,
            qm_u: 15,
            qm_v: 15,
            delta_q_present: false,
            delta_q_res: 0,
        }
    }

    #[test]
    fn test_coeff_decoder_creation() {
        let data = vec![0u8; 128];
        let quant = create_test_quant_params();
        let decoder = CoeffDecoder::new(data, quant, 8);
        assert!(decoder.has_more_data());
    }

    #[test]
    fn test_coeff_encoder_creation() {
        let encoder = CoeffEncoder::new();
        let output = encoder.finish();
        assert!(!output.is_empty() || output.is_empty());
    }

    #[test]
    fn test_batched_decoder() {
        let data = vec![0u8; 256];
        let quant = create_test_quant_params();
        let mut decoder = BatchedCoeffDecoder::new(data, quant, 8);

        let specs = vec![
            (TxSize::Tx4x4, TxType::DctDct, 0, false),
            (TxSize::Tx8x8, TxType::DctDct, 0, false),
        ];

        // Decoding may fail with test data, but should not crash
        let _ = decoder.decode_blocks(&specs);
    }

    #[test]
    fn test_coeff_analysis() {
        let mut analysis = CoeffAnalysis::new();
        let mut buffer = CoeffBuffer::new(4, 4);

        buffer.set(0, 0, 100); // DC
        buffer.set(1, 1, 50); // AC
        buffer.set(2, 2, -30); // AC

        analysis.analyze(&buffer);

        assert_eq!(analysis.total_coeffs, 16);
        assert_eq!(analysis.nonzero_count, 3);
        assert_eq!(analysis.zero_count, 13);
        assert_eq!(analysis.max_abs, 100);
    }

    #[test]
    fn test_coeff_analysis_sparsity() {
        let mut analysis = CoeffAnalysis::new();
        let buffer = CoeffBuffer::new(8, 8); // All zeros

        analysis.analyze(&buffer);

        assert_eq!(analysis.sparsity(), 100.0);
    }

    #[test]
    fn test_constants() {
        assert_eq!(COEFF_BASE_MAX, 3);
        assert_eq!(BR_CDF_SIZE, 4);
        assert_eq!(MAX_BR_PARAM, 5);
    }

    #[test]
    fn test_eob_offset_bits() {
        assert_eq!(EOB_OFFSET_BITS[0], 0);
        assert_eq!(EOB_OFFSET_BITS[3], 1);
        assert_eq!(EOB_OFFSET_BITS[11], 9);
    }

    #[test]
    fn test_coeff_analysis_avg_dc() {
        let mut analysis = CoeffAnalysis::new();
        let mut buffer = CoeffBuffer::new(4, 4);
        buffer.set(0, 0, 200);

        analysis.analyze(&buffer);

        // DC sum is 200, total coeffs is 16
        assert!((analysis.avg_dc() - 12.5).abs() < 0.1);
    }

    #[test]
    fn test_coeff_decoder_position() {
        let data = vec![0u8; 128];
        let quant = create_test_quant_params();
        let decoder = CoeffDecoder::new(data, quant, 8);
        assert!(decoder.position() <= 128);
    }
}