blad-cfa 0.0.2

//! Colour filter array decorrelation.
//!
//! A Bayer mosaic interleaves four colour channels on a 2×2 grid, so horizontally
//! adjacent samples measure *different* colours. Predictive image coders assume
//! neighbouring samples are correlated, and that assumption is violated on every step —
//! which is why compressing a mosaic as if it were a grayscale image leaves a large
//! amount on the table.
//!
//! Splitting the mosaic into four half-resolution planes, one per CFA position, restores
//! the assumption: within a plane, every neighbour measures the same colour.
//!
//! Measured on a Hasselblad X1D 3FR (8384×6304, 16-bit, ISO 400), lossless JXL:
//!
//! | approach                    | ratio |
//! |-----------------------------|-------|
//! | mosaic as one grayscale     | 0.613 |
//! | four CFA sub-planes         | 0.534 |
//!
//! **12.8% smaller, and faster** — the planes are quarter-size and encode in parallel.
//! General-purpose compressors cannot do this because they do not know the data is a
//! mosaic.
//!
//! # Why this operates on bytes
//!
//! The transform is a pure reordering: samples are moved, never read. So it needs
//! neither the CFA pattern (RGGB vs BGGR vs …) nor the byte order — which keeps it
//! exactly reversible for any sensor, and avoids materialising the image as `u16`.

#[derive(Debug, thiserror::Error, PartialEq, Eq)]
pub enum Error {
    #[error("CFA split needs even dimensions, got {width}×{height}")]
    OddDimensions { width: u32, height: u32 },
    #[error("expected {expected} bytes of mosaic data, got {actual}")]
    WrongLength { expected: usize, actual: usize },
}

pub type Result<T> = std::result::Result<T, Error>;

/// The four CFA positions, in raster order: (0,0), (1,0), (0,1), (1,1).
pub const PLANE_COUNT: usize = 4;

/// Four half-resolution planes decomposed from a mosaic, as raw sample bytes in the
/// source's own byte order.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Planes {
    /// Width of each plane (half the mosaic width).
    pub width: u32,
    /// Height of each plane (half the mosaic height).
    pub height: u32,
    pub bytes_per_sample: usize,
    pub planes: [Vec<u8>; PLANE_COUNT],
}

/// Split a mosaic into four sub-planes.
pub fn split(mosaic: &[u8], width: u32, height: u32, bytes_per_sample: usize) -> Result<Planes> {
    if !width.is_multiple_of(2) || !height.is_multiple_of(2) {
        return Err(Error::OddDimensions { width, height });
    }
    let bps = bytes_per_sample;
    let expected = width as usize * height as usize * bps;
    if mosaic.len() != expected {
        return Err(Error::WrongLength {
            expected,
            actual: mosaic.len(),
        });
    }

    let (w, h) = (width as usize, height as usize);
    let (hw, hh) = (w / 2, h / 2);
    let mut planes: [Vec<u8>; PLANE_COUNT] =
        std::array::from_fn(|_| Vec::with_capacity(hw * hh * bps));

    for y in 0..h {
        let row = &mosaic[y * w * bps..(y + 1) * w * bps];
        // Rows alternate between plane pair {0,1} (even y) and {2,3} (odd y).
        let pair = (y % 2) * 2;
        for x in 0..w {
            planes[pair + (x % 2)].extend_from_slice(&row[x * bps..(x + 1) * bps]);
        }
    }

    Ok(Planes {
        width: hw as u32,
        height: hh as u32,
        bytes_per_sample: bps,
        planes,
    })
}

/// Re-interleave four sub-planes back into a mosaic.
pub fn merge(planes: &Planes) -> Vec<u8> {
    let bps = planes.bytes_per_sample;
    let (hw, hh) = (planes.width as usize, planes.height as usize);
    let (w, h) = (hw * 2, hh * 2);
    let mut out = vec![0u8; w * h * bps];

    for y in 0..h {
        let plane_row = (y / 2) * hw * bps;
        let pair = (y % 2) * 2;
        for x in 0..w {
            let src = &planes.planes[pair + (x % 2)];
            let s = plane_row + (x / 2) * bps;
            let d = (y * w + x) * bps;
            out[d..d + bps].copy_from_slice(&src[s..s + bps]);
        }
    }
    out
}

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

    /// Deterministic pseudo-random bytes; no dev-dependency needed.
    fn noise(n: usize, seed: u64) -> Vec<u8> {
        let mut s = seed | 1;
        (0..n)
            .map(|_| {
                s ^= s << 13;
                s ^= s >> 7;
                s ^= s << 17;
                (s >> 24) as u8
            })
            .collect()
    }

    #[test]
    fn round_trip_is_byte_exact() {
        for (w, h) in [(8u32, 6u32), (64, 32), (2, 2), (1024, 8)] {
            let src = noise(w as usize * h as usize * 2, u64::from(w * h));
            let p = split(&src, w, h, 2).unwrap();
            assert_eq!(src, merge(&p), "round trip failed for {w}×{h}");
        }
    }

    #[test]
    fn planes_are_quarter_size_and_partition_the_mosaic() {
        let (w, h) = (8u32, 6u32);
        let src = noise(w as usize * h as usize * 2, 7);
        let p = split(&src, w, h, 2).unwrap();
        assert_eq!((p.width, p.height), (4, 3));
        let total: usize = p.planes.iter().map(|v| v.len()).sum();
        assert_eq!(total, (w * h) as usize * 2);
    }

    #[test]
    fn split_routes_positions_correctly() {
        // 2×2 mosaic of 16-bit samples, raster order 1,2,3,4 (little-endian bytes).
        let src: Vec<u8> = [1u16, 2, 3, 4]
            .iter()
            .flat_map(|v| v.to_le_bytes())
            .collect();
        let p = split(&src, 2, 2, 2).unwrap();
        assert_eq!(p.planes[0], vec![1, 0]); // (0,0)
        assert_eq!(p.planes[1], vec![2, 0]); // (1,0)
        assert_eq!(p.planes[2], vec![3, 0]); // (0,1)
        assert_eq!(p.planes[3], vec![4, 0]); // (1,1)
    }

    /// The transform must not depend on byte order: the same bytes in, the same bytes
    /// out, whatever they happen to mean.
    #[test]
    fn byte_order_is_irrelevant() {
        let src = noise(16 * 16 * 2, 3);
        let a = split(&src, 16, 16, 2).unwrap();
        let mut swapped = src.clone();
        for c in swapped.chunks_exact_mut(2) {
            c.swap(0, 1);
        }
        let b = split(&swapped, 16, 16, 2).unwrap();
        for i in 0..PLANE_COUNT {
            let unswapped: Vec<u8> = b.planes[i]
                .chunks_exact(2)
                .flat_map(|c| [c[1], c[0]])
                .collect();
            assert_eq!(a.planes[i], unswapped);
        }
    }

    #[test]
    fn odd_dimensions_rejected() {
        let src = vec![0u8; 3 * 4 * 2];
        assert_eq!(
            split(&src, 3, 4, 2),
            Err(Error::OddDimensions {
                width: 3,
                height: 4
            })
        );
    }

    #[test]
    fn wrong_length_rejected() {
        assert!(matches!(
            split(&[0u8; 10], 4, 4, 2),
            Err(Error::WrongLength { .. })
        ));
    }
}