tpt-cv-core 0.1.0

// SPDX-License-Identifier: MIT OR Apache-2.0

//! Color space representations and conversions.
//!
//! All conversions are generic over [`crate::pixel::Sample`] and operate in
//! normalized `[0, 1]` space internally, so `u8`, `u16`, and `f32` images use
//! the same matrices. Same-channel-count conversions (RGB↔BGR, RGB↔HSV,
//! RGB↔YUV) are performed in place; conversions that change the channel count
//! (RGB↔grayscale) write into a caller-provided destination.

use crate::image::{Image, ImageMut};
use crate::pixel::Sample;

const C_LUMA_R: f32 = 0.299;
const C_LUMA_G: f32 = 0.587;
const C_LUMA_B: f32 = 0.114;

/// Luma (BT.601) from a unit `(r, g, b)` triplet.
#[inline]
pub fn luma_unit(r: f32, g: f32, b: f32) -> f32 {
    C_LUMA_R * r + C_LUMA_G * g + C_LUMA_B * b
}

/// HSV from a unit `(r, g, b)` triplet; `h` is returned in `[0, 1)`.
#[inline]
pub fn hsv_unit(r: f32, g: f32, b: f32) -> (f32, f32, f32) {
    let max = r.max(g).max(b);
    let min = r.min(g).min(b);
    let delta = max - min;

    // `(g - b) / delta` is in `[-1, 1]`; wrap into `[0, 6)` (no_std-safe
    // replacement for `rem_euclid(6.0)`).
    let h = if delta == 0.0 {
        0.0
    } else if max == r {
        let mut h = (g - b) / delta;
        if h < 0.0 {
            h += 6.0;
        }
        h
    } else if max == g {
        (b - r) / delta + 2.0
    } else {
        (r - g) / delta + 4.0
    };
    let s = if max == 0.0 { 0.0 } else { delta / max };
    (h / 6.0, s, max)
}

/// RGB from a unit `(h, s, v)` triplet.
#[inline]
pub fn rgb_unit(h: f32, s: f32, v: f32) -> (f32, f32, f32) {
    if s == 0.0 {
        return (v, v, v);
    }
    // Wrap into `[0, 1)` (no_std-safe replacement for `rem_euclid(1.0)`).
    let mut h = h % 1.0;
    if h < 0.0 {
        h += 1.0;
    }
    let h = h * 6.0;
    let sector = h as i32; // truncation == floor for h >= 0
    let f = h - sector as f32;
    let p = v * (1.0 - s);
    let q = v * (1.0 - s * f);
    let t = v * (1.0 - s * (1.0 - f));
    match sector {
        0 => (v, t, p),
        1 => (q, v, p),
        2 => (p, v, t),
        3 => (p, q, v),
        4 => (t, p, v),
        _ => (v, p, q),
    }
}

/// Full-range BT.601 YUV from a unit `(r, g, b)` triplet; all outputs in `[0, 1]`.
#[inline]
pub fn yuv_unit(r: f32, g: f32, b: f32) -> (f32, f32, f32) {
    let y = C_LUMA_R * r + C_LUMA_G * g + C_LUMA_B * b;
    let u = 0.5 - 0.169 * r - 0.331 * g + 0.5 * b;
    let v = 0.5 + 0.5 * r - 0.419 * g - 0.081 * b;
    (y, u, v)
}

/// RGB from a unit `(y, u, v)` triplet.
#[inline]
pub fn rgb_from_yuv_unit(y: f32, u: f32, v: f32) -> (f32, f32, f32) {
    let cu = u - 0.5;
    let cv = v - 0.5;
    let r = y + 1.402 * cv;
    let g = y - 0.344 * cu - 0.714 * cv;
    let b = y + 1.772 * cu;
    (r, g, b)
}

/// Swap channel 0 and channel 2 of every pixel (RGB↔BGR) in place.
#[inline]
pub fn swap_red_blue<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    for y in 0..img.height() {
        let row = img.row_mut(y);
        for px in row.chunks_exact_mut(3) {
            px.swap(0, 2);
        }
    }
}

/// RGB → BGR in place.
#[inline]
pub fn rgb_to_bgr<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    swap_red_blue(img);
}

/// BGR → RGB in place.
#[inline]
pub fn bgr_to_rgb<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    swap_red_blue(img);
}

/// RGB → grayscale, writing into a `width×height` `C = 1` destination.
///
/// Returns `false` if `dst` is not `src.width × src.height`.
pub fn rgb_to_gray<T: Sample>(src: &Image<'_, T, 3>, dst: &mut ImageMut<'_, T, 1>) -> bool {
    if src.width() != dst.width() || src.height() != dst.height() {
        return false;
    }
    let h = src.height();
    for y in 0..h {
        let srow = src.row(y);
        let drow = dst.row_mut(y);
        for (dst_px, src_px) in drow.iter_mut().zip(srow.chunks_exact(3)) {
            let (r, g, b) = (
                src_px[0].to_unit(),
                src_px[1].to_unit(),
                src_px[2].to_unit(),
            );
            *dst_px = T::from_unit(luma_unit(r, g, b));
        }
    }
    true
}

/// Grayscale → RGB, replicating the luma channel. Returns `false` on size
/// mismatch.
pub fn gray_to_rgb<T: Sample>(src: &Image<'_, T, 1>, dst: &mut ImageMut<'_, T, 3>) -> bool {
    if src.width() != dst.width() || src.height() != dst.height() {
        return false;
    }
    let h = src.height();
    for y in 0..h {
        let srow = src.row(y);
        let drow = dst.row_mut(y);
        for (dst_px, v) in drow.chunks_exact_mut(3).zip(srow.iter()) {
            dst_px[0] = *v;
            dst_px[1] = *v;
            dst_px[2] = *v;
        }
    }
    true
}

/// RGB → HSV in place.
pub fn rgb_to_hsv<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    let w = img.width();
    let h = img.height();
    for y in 0..h {
        let row = img.row_mut(y);
        for x in 0..w {
            let off = x * 3;
            let (r, g, b) = (
                row[off].to_unit(),
                row[off + 1].to_unit(),
                row[off + 2].to_unit(),
            );
            let (h_, s, v) = hsv_unit(r, g, b);
            row[off] = T::from_unit(h_);
            row[off + 1] = T::from_unit(s);
            row[off + 2] = T::from_unit(v);
        }
    }
}

/// HSV → RGB in place.
pub fn hsv_to_rgb<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    let w = img.width();
    let h = img.height();
    for y in 0..h {
        let row = img.row_mut(y);
        for x in 0..w {
            let off = x * 3;
            let (hh, s, v) = (
                row[off].to_unit(),
                row[off + 1].to_unit(),
                row[off + 2].to_unit(),
            );
            let (r, g, b) = rgb_unit(hh, s, v);
            row[off] = T::from_unit(r);
            row[off + 1] = T::from_unit(g);
            row[off + 2] = T::from_unit(b);
        }
    }
}

/// RGB → YUV in place (full-range BT.601).
pub fn rgb_to_yuv<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    let w = img.width();
    let h = img.height();
    for y in 0..h {
        let row = img.row_mut(y);
        for x in 0..w {
            let off = x * 3;
            let (r, g, b) = (
                row[off].to_unit(),
                row[off + 1].to_unit(),
                row[off + 2].to_unit(),
            );
            let (yy, u, v) = yuv_unit(r, g, b);
            row[off] = T::from_unit(yy);
            row[off + 1] = T::from_unit(u);
            row[off + 2] = T::from_unit(v);
        }
    }
}

/// YUV → RGB in place (full-range BT.601).
pub fn yuv_to_rgb<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    let w = img.width();
    let h = img.height();
    for y in 0..h {
        let row = img.row_mut(y);
        for x in 0..w {
            let off = x * 3;
            let (yy, u, v) = (
                row[off].to_unit(),
                row[off + 1].to_unit(),
                row[off + 2].to_unit(),
            );
            let (r, g, b) = rgb_from_yuv_unit(yy, u, v);
            row[off] = T::from_unit(r);
            row[off + 1] = T::from_unit(g);
            row[off + 2] = T::from_unit(b);
        }
    }
}

/// YCbCr (BT.601, studio-swing normalized) from a unit `(r, g, b)` triplet;
/// `y` is in `[0, 1]`, `cb`/`cr` are centered in `[0, 1]` around `0.5`.
#[inline]
pub fn ycbcr_unit(r: f32, g: f32, b: f32) -> (f32, f32, f32) {
    let y = 0.299 * r + 0.587 * g + 0.114 * b;
    let cb = 0.5 - 0.168_736 * r - 0.331_264 * g + 0.5 * b;
    let cr = 0.5 + 0.5 * r - 0.418_688 * g - 0.081_312 * b;
    (y, cb, cr)
}

/// RGB from a unit `(y, cb, cr)` triplet.
#[inline]
pub fn rgb_from_ycbcr_unit(y: f32, cb: f32, cr: f32) -> (f32, f32, f32) {
    let cb = cb - 0.5;
    let cr = cr - 0.5;
    (
        y + 1.402 * cr,
        y - 0.344_136 * cb - 0.714_136 * cr,
        y + 1.772 * cb,
    )
}

/// RGB → YCbCr in place (BT.601).
pub fn rgb_to_ycbcr<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    let w = img.width();
    let h = img.height();
    for y in 0..h {
        let row = img.row_mut(y);
        for x in 0..w {
            let off = x * 3;
            let (yy, cb, cr) = ycbcr_unit(
                row[off].to_unit(),
                row[off + 1].to_unit(),
                row[off + 2].to_unit(),
            );
            row[off] = T::from_unit(yy);
            row[off + 1] = T::from_unit(cb);
            row[off + 2] = T::from_unit(cr);
        }
    }
}

/// YCbCr → RGB in place (BT.601).
pub fn ycbcr_to_rgb<T: Sample>(img: &mut ImageMut<'_, T, 3>) {
    let w = img.width();
    let h = img.height();
    for y in 0..h {
        let row = img.row_mut(y);
        for x in 0..w {
            let off = x * 3;
            let (r, g, b) = rgb_from_ycbcr_unit(
                row[off].to_unit(),
                row[off + 1].to_unit(),
                row[off + 2].to_unit(),
            );
            row[off] = T::from_unit(r);
            row[off + 1] = T::from_unit(g);
            row[off + 2] = T::from_unit(b);
        }
    }
}

/// CIE L\*a\*b\* (D65) from a unit `(r, g, b)` triplet.
///
/// Returns `(l, a, b)` scaled toward a unit cube: `l` is the standard CIE L\*
/// divided by 100 (`[0, 1]`), while `a*`/`b*` are divided by 127 (roughly
/// `[-1, 1]`). Use [`rgb_from_lab_unit`] to invert.
#[inline]
pub fn lab_unit(r: f32, g: f32, b: f32) -> (f32, f32, f32) {
    // sRGB → linear.
    let lin = |c: f32| {
        if c <= 0.040_45 {
            c / 12.92
        } else {
            libm::powf((c + 0.055) / 1.055, 2.4)
        }
    };
    let (r, g, b) = (lin(r), lin(g), lin(b));

    // Linear sRGB → XYZ (D65).
    let x = 0.412_456_4 * r + 0.357_576_1 * g + 0.180_437_5 * b;
    let y = 0.212_672_9 * r + 0.715_152_2 * g + 0.072_175_0 * b;
    let z = 0.019_333_9 * r + 0.119_192 * g + 0.950_304_1 * b;

    // Normalize by the D65 white point.
    const WX: f32 = 0.950_47;
    const WZ: f32 = 1.088_83;
    let f = |t: f32| {
        if t > 216.0 / 24_389.0 {
            libm::powf(t, 1.0 / 3.0)
        } else {
            (24_389.0 / 27.0 * t + 16.0) / 116.0
        }
    };
    let fx = f(x / WX);
    let fy = f(y);
    let fz = f(z / WZ);

    let l = 116.0 * fy - 16.0;
    let a = 500.0 * (fx - fy);
    let bb = 200.0 * (fy - fz);
    (l / 100.0, a / 127.0, bb / 127.0)
}

/// Unit RGB from [`lab_unit`] output `(l, a, b)`.
#[inline]
pub fn rgb_from_lab_unit(l: f32, a: f32, b: f32) -> (f32, f32, f32) {
    let l = l * 100.0;
    let a = a * 127.0;
    let b = b * 127.0;

    let fy = (l + 16.0) / 116.0;
    let fx = fy + a / 500.0;
    let fz = fy - b / 200.0;
    const EPS: f32 = 216.0 / 24_389.0;
    const KAPA: f32 = 24_389.0 / 27.0;
    let finv = |t: f32| {
        let t3 = t * t * t;
        if t3 > EPS {
            t3
        } else {
            (116.0 * t - 16.0) / KAPA
        }
    };

    let x = finv(fx) * 0.950_47;
    let y = finv(fy);
    let z = finv(fz) * 1.088_83;

    // XYZ → linear sRGB.
    let rl = 3.240_454_2 * x - 1.537_138_5 * y - 0.498_531_4 * z;
    let gl = -0.969_266 * x + 1.876_010_8 * y + 0.041_556 * z;

    let bl = 0.055_643_4 * x - 0.204_025_9 * y + 1.057_225_2 * z;

    // Linear → sRGB, clamped to [0, 1].
    let gam = |c: f32| {
        let c = c.clamp(0.0, 1.0);
        if c <= 0.003_130_8 {
            12.92 * c
        } else {
            1.055 * libm::powf(c, 1.0 / 2.4) - 0.055
        }
    };
    (gam(rl), gam(gl), gam(bl))
}

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

    fn rgb_buf<T: Sample, const N: usize>(rows: &[[T; 3]; N]) -> ImageBuf<T, 3> {
        let mut data = alloc::vec![T::ZERO; N * 3];
        for (i, px) in rows.iter().enumerate() {
            data[i * 3] = px[0];
            data[i * 3 + 1] = px[1];
            data[i * 3 + 2] = px[2];
        }
        ImageBuf::from_vec(data, N, 1).unwrap()
    }

    #[test]
    fn bgr_roundtrip() {
        let mut buf = rgb_buf(&[[10u8, 200, 30], [255, 0, 128]]);
        rgb_to_bgr(&mut buf.as_image_mut());
        assert_eq!(buf.as_image().pixel(0, 0).channels, [30, 200, 10]);
        bgr_to_rgb(&mut buf.as_image_mut());
        assert_eq!(buf.as_image().pixel(0, 0).channels, [10, 200, 30]);
        assert_eq!(buf.as_image().pixel(1, 0).channels, [255, 0, 128]);
    }

    #[test]
    fn ycbcr_roundtrip_unit() {
        for &(r, g, b) in &[
            (0.0f32, 0.0, 0.0),
            (1.0, 1.0, 1.0),
            (1.0, 0.0, 0.0),
            (0.2, 0.5, 0.8),
        ] {
            let (y, cb, cr) = ycbcr_unit(r, g, b);
            let (r2, g2, b2) = rgb_from_ycbcr_unit(y, cb, cr);
            assert!((r2 - r).abs() < 1e-4);
            assert!((g2 - g).abs() < 1e-4);
            assert!((b2 - b).abs() < 1e-4);
        }
    }

    #[test]
    fn ycbcr_in_place_roundtrip() {
        // Roundtrip restores the original pixels within u8 quantization error.
        let mut buf = rgb_buf(&[[200u8, 200, 200], [10, 200, 30]]);
        rgb_to_ycbcr(&mut buf.as_image_mut());
        ycbcr_to_rgb(&mut buf.as_image_mut());
        for x in 0..2usize {
            for c in 0..3usize {
                let orig = [[200u8, 200, 200], [10, 200, 30]][x][c];
                let got = buf.as_image().pixel(x, 0).channels[c];
                assert!(
                    (i16::from(got) - i16::from(orig)).abs() <= 2,
                    "pixel {x} ch {c}: {got} vs {orig}"
                );
            }
        }
    }

    #[test]
    fn lab_known_values_and_roundtrip() {
        // Pure white is L*=100 → l=1; a/b ≈ 0.
        let (l, a, b) = lab_unit(1.0, 1.0, 1.0);
        assert!((l - 1.0).abs() < 1e-3, "l={l}");
        assert!(a.abs() < 5.0 / 127.0 && b.abs() < 5.0 / 127.0);

        // Roundtrip over a spread of colors.
        for &(r, g, b) in &[
            (0.1f32, 0.6, 0.9),
            (0.9, 0.3, 0.2),
            (0.5, 0.5, 0.5),
            (0.05, 0.05, 0.4),
        ] {
            let (l, a, bb) = lab_unit(r, g, b);
            let (r2, g2, b2) = rgb_from_lab_unit(l, a, bb);
            assert!((r2 - r).abs() < 1e-3, "r {r} -> {r2}");
            assert!((g2 - g).abs() < 1e-3, "g {g} -> {g2}");
            assert!((b2 - b).abs() < 1e-3, "b {b} -> {b2}");
        }
    }

    #[test]
    fn gray_known_values() {
        let src = rgb_buf(&[[255u8, 255, 255], [0, 0, 0], [255, 0, 0], [128, 128, 128]]);
        let mut gray = ImageBuf::<u8, 1>::new(4, 1);
        assert!(rgb_to_gray(&src.as_image(), &mut gray.as_image_mut()));
        let px = gray.as_image();
        assert_eq!(px.pixel(0, 0).scalar(), 255); // white
        assert_eq!(px.pixel(1, 0).scalar(), 0); // black
        assert_eq!(px.pixel(2, 0).scalar(), 76); // pure red ≈ 0.299*255 ≈ 76
        assert_eq!(px.pixel(3, 0).scalar(), 128); // mid gray
    }

    #[test]
    fn gray_roundtrip_via_gray_to_rgb() {
        let src = rgb_buf(&[[120u8, 34, 200], [1, 2, 3], [240, 240, 240]]);
        let mut gray = ImageBuf::<u8, 1>::new(3, 1);
        rgb_to_gray(&src.as_image(), &mut gray.as_image_mut());
        let mut back = ImageBuf::<u8, 3>::new(3, 1);
        gray_to_rgb(&gray.as_image(), &mut back.as_image_mut());
        let g = gray.as_image();
        let b = back.as_image();
        for x in 0..3 {
            let v = g.pixel(x, 0).scalar();
            assert_eq!(b.pixel(x, 0).channels, [v, v, v]);
        }
    }

    #[test]
    fn hsv_roundtrip_u8_within_tolerance() {
        let buf = rgb_buf(&[
            [255, 0, 0],
            [0, 255, 0],
            [0, 0, 255],
            [128, 64, 192],
            [10, 200, 30],
            [0, 0, 0],
            [255, 255, 255],
        ]);
        {
            let original = buf.as_image();
            // sanity: pure red -> h=0, s=1, v=1
            let mut scratch = ImageBuf::<u8, 3>::new(7, 1);
            scratch.as_image_mut().copy_from(&original);
            let mut m = scratch.as_image_mut();
            rgb_to_hsv(&mut m);
            assert_eq!(scratch.as_image().pixel(0, 0).channels, [0, 255, 255]);
            assert_eq!(scratch.as_image().pixel(3, 0).channels[2], 192); // v = max
        }
        let original = buf.as_image();
        let mut copy = ImageBuf::<u8, 3>::new(7, 1);
        copy.as_image_mut().copy_from(&original);
        rgb_to_hsv(&mut copy.as_image_mut());
        hsv_to_rgb(&mut copy.as_image_mut());
        for x in 0..7 {
            let a = original.pixel(x, 0).channels;
            let b = copy.as_image().pixel(x, 0).channels;
            for c in 0..3 {
                let diff = (a[c] as i32 - b[c] as i32).unsigned_abs();
                assert!(diff <= 2, "channel {c} pixel {x}: {} vs {}", a[c], b[c]);
            }
        }
    }

    #[test]
    fn hsv_unit_known_pure_hues() {
        assert_eq!(hsv_unit(1.0, 0.0, 0.0), (0.0, 1.0, 1.0)); // red
        let (h, s, v) = hsv_unit(0.0, 1.0, 0.0); // green
        assert!((h - 1.0 / 3.0).abs() < 1e-6, "green hue {h}");
        assert_eq!((s, v), (1.0, 1.0));
        let (h, _, _) = hsv_unit(0.0, 0.0, 1.0); // blue
        assert!((h - 2.0 / 3.0).abs() < 1e-6, "blue hue {h}");
    }

    #[test]
    fn yuv_roundtrip_all_samples() {
        // u8
        let buf = rgb_buf(&[[255u8, 0, 0], [0, 255, 0], [0, 0, 255], [123, 45, 67]]);
        let original = buf.as_image();
        let mut copy = ImageBuf::<u8, 3>::new(4, 1);
        copy.as_image_mut().copy_from(&original);
        rgb_to_yuv(&mut copy.as_image_mut());
        yuv_to_rgb(&mut copy.as_image_mut());
        for x in 0..4 {
            let a = original.pixel(x, 0).channels;
            let b = copy.as_image().pixel(x, 0).channels;
            for c in 0..3 {
                let diff = (a[c] as i32 - b[c] as i32).unsigned_abs();
                assert!(diff <= 2, "u8 channel {c} pixel {x}: {} vs {}", a[c], b[c]);
            }
        }

        // f32
        let fbuf = rgb_buf(&[[0.9f32, 0.1, 0.7], [0.0, 0.0, 0.0], [1.0, 1.0, 1.0]]);
        let foriginal = fbuf.as_image();
        let mut fcopy = ImageBuf::<f32, 3>::new(3, 1);
        fcopy.as_image_mut().copy_from(&foriginal);
        rgb_to_yuv(&mut fcopy.as_image_mut());
        yuv_to_rgb(&mut fcopy.as_image_mut());
        for x in 0..3 {
            let a = foriginal.pixel(x, 0).channels;
            let b = fcopy.as_image().pixel(x, 0).channels;
            for c in 0..3 {
                // Tolerance accounts for the rounded 3-digit YCbCr coefficients.
                assert!(
                    (a[c] - b[c]).abs() < 2e-3,
                    "f32 channel {c} pixel {x}: {} vs {}",
                    a[c],
                    b[c]
                );
            }
        }
    }

    #[test]
    fn yuv_known_gray() {
        let mut buf = rgb_buf(&[[128u8, 128, 128], [0, 0, 0], [255, 255, 255]]);
        rgb_to_yuv(&mut buf.as_image_mut());
        let img = buf.as_image();
        let within = |a: [u8; 3], b: [u8; 3]| {
            a.iter()
                .zip(b.iter())
                .all(|(x, y)| (i32::from(*x) - i32::from(*y)).unsigned_abs() <= 1)
        };
        // Neutral gray maps to u ≈ v ≈ 0.5 (128 in 8-bit)
        assert!(within(img.pixel(0, 0).channels, [128, 128, 128]));
        assert!(within(img.pixel(1, 0).channels, [0, 128, 128]));
        assert!(within(img.pixel(2, 0).channels, [255, 128, 128]));
    }

    #[test]
    fn grayscale_mismatch_returns_false() {
        let src = rgb_buf(&[[1u8, 2, 3]]);
        let mut gray = ImageBuf::<u8, 1>::new(2, 1);
        assert!(!rgb_to_gray(&src.as_image(), &mut gray.as_image_mut()));
    }
}