tx2_iff/

wavelet.rs

//! Integer Wavelet Transform (CDF 5/3) for Layer 1 skeleton
//!
//! Implements the Cohen-Daubechies-Feauveau 5/3 biorthogonal wavelet
//! using the lifting scheme with all operations in integers for perfect
//! reconstruction and deterministic behavior.

use crate::error::{IffError, Result};
use crate::prime::QuantizationTable;
use crate::compression::{compress_rle, decompress_rle};
use serde::{Deserialize, Serialize};

/// Wavelet subband type
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SubBand {
    LL,
    LH,
    HL,
    HH,
}

/// Wavelet decomposition structure (Compressed)
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct WaveletDecomposition {
    /// Original image dimensions
    pub width: u32,
    pub height: u32,
    /// Number of decomposition levels
    pub levels: usize,
    /// Number of channels
    pub channels: u8,
    /// Compressed coefficients (RLE encoded)
    /// Layout: [Channel 0 Data] [Channel 1 Data] ...
    pub data: Vec<u8>,
}

impl WaveletDecomposition {
    /// Create a new empty decomposition
    pub fn new(width: u32, height: u32, levels: usize, channels: u8) -> Self {
        WaveletDecomposition {
            width,
            height,
            levels,
            channels,
            data: Vec::new(),
        }
    }

    /// Create from dense coefficient buffers
    pub fn from_dense(
        width: u32,
        height: u32,
        levels: usize,
        channel_data: &[Vec<i32>],
    ) -> Result<Self> {
        let mut all_data = Vec::new();
        for channel in channel_data {
            all_data.extend_from_slice(channel);
        }

        let compressed = compress_rle(&all_data)?;

        Ok(WaveletDecomposition {
            width,
            height,
            levels,
            channels: channel_data.len() as u8,
            data: compressed,
        })
    }

    /// Decompress to dense coefficient buffers
    pub fn to_dense(&self) -> Result<Vec<Vec<i32>>> {
        let total_pixels = (self.width * self.height) as usize;
        let expected_len = total_pixels * self.channels as usize;

        let all_data = decompress_rle(&self.data, Some(expected_len))?;

        let mut channels = Vec::with_capacity(self.channels as usize);
        for i in 0..self.channels as usize {
            let start = i * total_pixels;
            let end = start + total_pixels;
            if end > all_data.len() {
                return Err(IffError::Other("Insufficient data for channels".to_string()));
            }
            channels.push(all_data[start..end].to_vec());
        }

        Ok(channels)
    }
}

/// CDF 5/3 Integer Wavelet Transform
pub struct Cdf53Transform {
    /// Number of decomposition levels
    levels: usize,
}

impl Cdf53Transform {
    /// Create a new CDF 5/3 transform with specified levels
    pub fn new(levels: usize) -> Self {
        Cdf53Transform { levels }
    }

    /// Forward transform: image → wavelet coefficients (dense)
    pub fn forward(&self, image: &[i32], width: usize, height: usize) -> Result<Vec<i32>> {
        if image.len() != width * height {
            return Err(IffError::Other(
                "Image dimensions don't match data length".to_string(),
            ));
        }

        // Working buffer (will be modified in-place)
        let mut buffer = image.to_vec();

        // Apply transform levels
        let mut current_width = width;
        let mut current_height = height;

        for _ in 0..self.levels {
            if current_width < 2 || current_height < 2 {
                break; // Can't decompose further
            }

            // Transform rows
            for y in 0..current_height {
                self.forward_1d_row(&mut buffer, y, current_width, width);
            }

            // Transform columns
            for x in 0..current_width {
                self.forward_1d_col(&mut buffer, x, current_height, width);
            }

            // Next level operates on LL subband
            current_width /= 2;
            current_height /= 2;
        }

        Ok(buffer)
    }

    /// Inverse transform: wavelet coefficients (dense) → image
    pub fn inverse(&self, coefficients: &[i32], width: usize, height: usize) -> Result<Vec<i32>> {
        if coefficients.len() != width * height {
            return Err(IffError::Other(
                "Coefficient dimensions don't match data length".to_string(),
            ));
        }

        let mut buffer = coefficients.to_vec();

        // Apply inverse transform levels (in reverse order)

        // Adjust starting level if image is too small

        // We need to reconstruct the correct sequence of widths/heights
        // Or just iterate backwards
        // The forward loop goes: (w,h) -> (w/2, h/2) ...
        // The inverse loop should go: ... (w/2, h/2) -> (w,h)

        // Let's pre-calculate the dimensions for each level
        let mut dims = Vec::new();
        let mut w = width;
        let mut h = height;
        for _ in 0..self.levels {
            dims.push((w, h));
            w /= 2;
            h /= 2;
        }

        for (w, h) in dims.iter().rev() {
             let current_width = *w;
             let current_height = *h;
             
             if current_width < 2 || current_height < 2 {
                 continue;
             }

            // Inverse transform columns
            for x in 0..current_width {
                self.inverse_1d_col(&mut buffer, x, current_height, width);
            }

            // Inverse transform rows
            for y in 0..current_height {
                self.inverse_1d_row(&mut buffer, y, current_width, width);
            }
        }

        Ok(buffer)
    }

    /// Quantize coefficients using 360-prime pattern (in-place)
    pub fn quantize(&self, buffer: &mut [i32], width: usize, height: usize, table: &QuantizationTable) {
        // We need to iterate over subbands to know coordinates
        // But the buffer is dense.
        // We can iterate over levels and subbands similar to extract_coefficients
        
        let mut current_width = width;
        let mut current_height = height;

        for level in 0..self.levels {
            let half_w = current_width / 2;
            let half_h = current_height / 2;
            
            if half_w == 0 || half_h == 0 { break; }

            // Process subbands
            // LL is skipped (it's processed in next level, or at the end)
            // But wait, LL of level N is the input to level N+1.
            // We only quantize the details (LH, HL, HH) of each level.
            // And the final LL.

            // LH subband
            for y in 0..half_h {
                for x in half_w..current_width {
                    self.quantize_pixel(buffer, x, y, width, table);
                }
            }

            // HL subband
            for y in half_h..current_height {
                for x in 0..half_w {
                    self.quantize_pixel(buffer, x, y, width, table);
                }
            }

            // HH subband
            for y in half_h..current_height {
                for x in half_w..current_width {
                    self.quantize_pixel(buffer, x, y, width, table);
                }
            }
            
            // If this is the last level, also quantize the LL band
            if level == self.levels - 1 {
                for y in 0..half_h {
                    for x in 0..half_w {
                        self.quantize_pixel(buffer, x, y, width, table);
                    }
                }
            }

            current_width = half_w;
            current_height = half_h;
        }
    }
    
    fn quantize_pixel(&self, buffer: &mut [i32], x: usize, y: usize, width: usize, table: &QuantizationTable) {
        let idx = y * width + x;
        let val = buffer[idx];
        let step = table.get_step(x, y);
        
        let quantized = val / step as i32;
        
        // Aggressive sparsification
        if quantized.abs() < 2 { // Reduced threshold from 3 to 2 to be less aggressive?
             buffer[idx] = 0;
        } else {
             buffer[idx] = quantized * step as i32;
        }
    }

    /// Forward 1D transform on a row (in-place lifting scheme)
    fn forward_1d_row(&self, data: &mut [i32], y: usize, width: usize, stride: usize) {
        if width < 2 {
            return;
        }

        let offset = y * stride;
        let mut temp = vec![0i32; width];

        // Predict step: d[i] = s[2i+1] - floor((s[2i] + s[2i+2]) / 2)
        for i in 0..(width / 2) {
            let left = data[offset + 2 * i];
            let right = if 2 * i + 2 < width {
                data[offset + 2 * i + 2]
            } else {
                data[offset + 2 * i] // Mirror boundary
            };
            temp[width / 2 + i] = data[offset + 2 * i + 1] - ((left + right) / 2);
        }

        // Update step: s[i] = s[2i] + floor((d[i-1] + d[i] + 2) / 4)
        for i in 0..(width / 2) {
            let d_left = if i > 0 {
                temp[width / 2 + i - 1]
            } else {
                temp[width / 2] // Mirror boundary
            };
            let d_right = if i < width / 2 {
                temp[width / 2 + i]
            } else {
                temp[width / 2 + i - 1] // Mirror boundary
            };
            temp[i] = data[offset + 2 * i] + ((d_left + d_right + 2) / 4);
        }

        // Copy back
        for i in 0..width {
            data[offset + i] = temp[i];
        }
    }

    /// Forward 1D transform on a column
    fn forward_1d_col(&self, data: &mut [i32], x: usize, height: usize, stride: usize) {
        if height < 2 {
            return;
        }

        let mut temp = vec![0i32; height];

        // Predict step
        for i in 0..(height / 2) {
            let top = data[2 * i * stride + x];
            let bottom = if 2 * i + 2 < height {
                data[(2 * i + 2) * stride + x]
            } else {
                data[2 * i * stride + x] // Mirror boundary
            };
            temp[height / 2 + i] = data[(2 * i + 1) * stride + x] - ((top + bottom) / 2);
        }

        // Update step
        for i in 0..(height / 2) {
            let d_top = if i > 0 {
                temp[height / 2 + i - 1]
            } else {
                temp[height / 2] // Mirror boundary
            };
            let d_bottom = if i < height / 2 {
                temp[height / 2 + i]
            } else {
                temp[height / 2 + i - 1] // Mirror boundary
            };
            temp[i] = data[2 * i * stride + x] + ((d_top + d_bottom + 2) / 4);
        }

        // Copy back
        for i in 0..height {
            data[i * stride + x] = temp[i];
        }
    }

    /// Inverse 1D transform on a row
    fn inverse_1d_row(&self, data: &mut [i32], y: usize, width: usize, stride: usize) {
        if width < 2 {
            return;
        }

        let offset = y * stride;
        let mut temp = vec![0i32; width];

        // Copy to temp
        for i in 0..width {
            temp[i] = data[offset + i];
        }

        // Undo update step
        for i in 0..(width / 2) {
            let d_left = if i > 0 {
                temp[width / 2 + i - 1]
            } else {
                temp[width / 2]
            };
            let d_right = if i < width / 2 {
                temp[width / 2 + i]
            } else {
                temp[width / 2 + i - 1]
            };
            data[offset + 2 * i] = temp[i] - ((d_left + d_right + 2) / 4);
        }

        // Undo predict step
        for i in 0..(width / 2) {
            let left = data[offset + 2 * i];
            let right = if 2 * i + 2 < width {
                data[offset + 2 * i + 2]
            } else {
                data[offset + 2 * i]
            };
            data[offset + 2 * i + 1] = temp[width / 2 + i] + ((left + right) / 2);
        }
    }

    /// Inverse 1D transform on a column
    fn inverse_1d_col(&self, data: &mut [i32], x: usize, height: usize, stride: usize) {
        if height < 2 {
            return;
        }

        let mut temp = vec![0i32; height];

        // Copy to temp
        for i in 0..height {
            temp[i] = data[i * stride + x];
        }

        // Undo update step
        for i in 0..(height / 2) {
            let d_top = if i > 0 {
                temp[height / 2 + i - 1]
            } else {
                temp[height / 2]
            };
            let d_bottom = if i < height / 2 {
                temp[height / 2 + i]
            } else {
                temp[height / 2 + i - 1]
            };
            data[2 * i * stride + x] = temp[i] - ((d_top + d_bottom + 2) / 4);
        }

        // Undo predict step
        for i in 0..(height / 2) {
            let top = data[2 * i * stride + x];
            let bottom = if 2 * i + 2 < height {
                data[(2 * i + 2) * stride + x]
            } else {
                data[2 * i * stride + x]
            };
            data[(2 * i + 1) * stride + x] = temp[height / 2 + i] + ((top + bottom) / 2);
        }
    }
}

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

    #[test]
    fn test_perfect_reconstruction() {
        // Create a small test image (8x8)
        let width = 8;
        let height = 8;
        let image: Vec<i32> = (0..(width * height))
            .map(|i| ((i % 256) as i32 - 128))
            .collect();

        // Forward transform
        let transform = Cdf53Transform::new(2);
        let coeffs = transform.forward(&image, width, height).unwrap();

        // Inverse transform
        let reconstructed = transform.inverse(&coeffs, width, height).unwrap();

        // Check perfect reconstruction
        for (orig, recon) in image.iter().zip(reconstructed.iter()) {
            assert_eq!(orig, recon, "Perfect reconstruction failed");
        }
    }
}