tx2-iff 0.1.0

//! Color space conversions and subsampling
//!
//! Implements YCoCg-R (Reversible) color space and 4:2:0 chroma subsampling.
//!
//! ## YCoCg-R
//! Reversible transformation between RGB and YCoCg.
//!
//! Forward:
//! Co = R - B
//! tmp = B + Co/2
//! Cg = G - tmp
//! Y = tmp + Cg/2
//!
//! Inverse:
//! tmp = Y - Cg/2
//! G = Cg + tmp
//! B = tmp - Co/2
//! R = B + Co

use crate::error::{IffError, Result};

/// Image channel data
#[derive(Debug, Clone)]
pub struct Channel {
    pub width: usize,
    pub height: usize,
    pub data: Vec<i32>,
}

impl Channel {
    pub fn new(width: usize, height: usize) -> Self {
        Channel {
            width,
            height,
            data: vec![0; width * height],
        }
    }
}

/// YCoCg image
#[derive(Debug, Clone)]
pub struct YCoCgImage {
    pub width: usize,
    pub height: usize,
    pub y: Channel,
    pub co: Channel,
    pub cg: Channel,
}

/// Convert RGB to YCoCg-R
pub fn rgb_to_ycocg(rgb: &[[u8; 3]], width: usize, height: usize) -> YCoCgImage {
    let mut y_channel = Channel::new(width, height);
    let mut co_channel = Channel::new(width, height);
    let mut cg_channel = Channel::new(width, height);

    for i in 0..rgb.len() {
        let r = rgb[i][0] as i32;
        let g = rgb[i][1] as i32;
        let b = rgb[i][2] as i32;

        let co = r - b;
        let tmp = b + (co >> 1);
        let cg = g - tmp;
        let y = tmp + (cg >> 1);

        y_channel.data[i] = y;
        co_channel.data[i] = co;
        cg_channel.data[i] = cg;
    }

    YCoCgImage {
        width,
        height,
        y: y_channel,
        co: co_channel,
        cg: cg_channel,
    }
}

/// Convert YCoCg-R to RGB
pub fn ycocg_to_rgb(image: &YCoCgImage) -> Result<Vec<[u8; 3]>> {
    let width = image.width;
    let height = image.height;
    let len = width * height;

    if image.y.data.len() != len || image.co.data.len() != len || image.cg.data.len() != len {
        return Err(IffError::Other("Channel dimensions mismatch".to_string()));
    }

    let mut rgb = Vec::with_capacity(len);

    for i in 0..len {
        let y = image.y.data[i];
        let co = image.co.data[i];
        let cg = image.cg.data[i];

        let tmp = y - (cg >> 1);
        let g = cg + tmp;
        let b = tmp - (co >> 1);
        let r = b + co;

        rgb.push([
            r.clamp(0, 255) as u8,
            g.clamp(0, 255) as u8,
            b.clamp(0, 255) as u8,
        ]);
    }

    Ok(rgb)
}

/// Subsample channel (2x2 average)
pub fn subsample_420(channel: &Channel) -> Channel {
    let new_width = channel.width / 2;
    let new_height = channel.height / 2;
    let mut new_data = Vec::with_capacity(new_width * new_height);

    for y in 0..new_height {
        for x in 0..new_width {
            // Average 2x2 block
            let idx = (y * 2) * channel.width + (x * 2);
            let sum = channel.data[idx]
                + channel.data[idx + 1]
                + channel.data[idx + channel.width]
                + channel.data[idx + channel.width + 1];
            new_data.push(sum / 4);
        }
    }

    Channel {
        width: new_width,
        height: new_height,
        data: new_data,
    }
}

/// Upsample channel (nearest neighbor or bilinear)
/// Using simple duplication for now to match subsample
pub fn upsample_420(channel: &Channel, target_width: usize, target_height: usize) -> Channel {
    let mut new_data = vec![0; target_width * target_height];

    for y in 0..target_height {
        for x in 0..target_width {
            let src_x = (x / 2).min(channel.width - 1);
            let src_y = (y / 2).min(channel.height - 1);
            let val = channel.data[src_y * channel.width + src_x];
            new_data[y * target_width + x] = val;
        }
    }

    Channel {
        width: target_width,
        height: target_height,
        data: new_data,
    }
}

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

    #[test]
    fn test_ycocg_reversible() {
        let rgb = vec![[100, 150, 200]];
        let ycocg = rgb_to_ycocg(&rgb, 1, 1);
        let restored = ycocg_to_rgb(&ycocg).unwrap();
        assert_eq!(rgb, restored);
    }

    #[test]
    fn test_subsample() {
        let data = vec![10, 10, 20, 20];
        let channel = Channel {
            width: 2,
            height: 2,
            data,
        };
        let sub = subsample_420(&channel);
        assert_eq!(sub.width, 1);
        assert_eq!(sub.height, 1);
        assert_eq!(sub.data[0], 15);
    }
}