ocr 0.1.0

use image::{DynamicImage, GrayImage};

pub fn grayscale(img: &DynamicImage) -> DynamicImage {
    img.grayscale()
}

pub fn threshold(img: &DynamicImage, value: u8) -> DynamicImage {
    let gray = img.grayscale();
    let mut luma = gray.to_luma8();
    for pixel in luma.pixels_mut() {
        pixel[0] = if pixel[0] > value { 255 } else { 0 };
    }
    DynamicImage::ImageLuma8(luma)
}

pub fn preprocess(img: &DynamicImage) -> DynamicImage {
    let gray = grayscale(img);
    threshold(&gray, 128)
}

pub fn binarize(img: &DynamicImage) -> Vec<Vec<bool>> {
    let gray = img.grayscale();
    let luma = gray.to_luma8();
    let (width, height) = luma.dimensions();
    let thr = otsu_threshold(&luma);

    let below_count = luma.iter().filter(|&&p| p <= thr).count();
    let total = (width * height) as usize;
    let foreground_is_dark = below_count < total;

    let mut grid = vec![vec![false; width as usize]; height as usize];
    for y in 0..height as usize {
        for x in 0..width as usize {
            let dark = luma.get_pixel(x as u32, y as u32)[0] <= thr;
            grid[y][x] = if foreground_is_dark { dark } else { !dark };
        }
    }
    grid
}

fn otsu_threshold(gray: &GrayImage) -> u8 {
    let mut hist = [0u32; 256];
    for pixel in gray.iter() {
        hist[*pixel as usize] += 1;
    }
    let total = gray.width() * gray.height();

    let sum: f64 = (0..256).map(|i| i as f64 * hist[i] as f64).sum();
    let mut sum_b = 0.0_f64;
    let mut w_b = 0.0_f64;
    let mut best_var = 0.0_f64;
    let mut best_t = 0u8;

    for t in 0..256 {
        w_b += hist[t] as f64;
        if w_b == 0.0 {
            continue;
        }
        let w_f = total as f64 - w_b;
        if w_f == 0.0 {
            break;
        }
        sum_b += t as f64 * hist[t] as f64;
        let m_b = sum_b / w_b;
        let m_f = (sum - sum_b) / w_f;
        let var = w_b * w_f * (m_b - m_f).powi(2);
        if var > best_var {
            best_var = var;
            best_t = t as u8;
        }
    }
    best_t
}

pub fn clean_noise(grid: &mut [Vec<bool>], min_component_size: usize) {
    let h = grid.len();
    if h == 0 {
        return;
    }
    let w = grid[0].len();
    let mut visited = vec![vec![false; w]; h];
    let mut components: Vec<Vec<(usize, usize)>> = Vec::new();

    for y in 0..h {
        for x in 0..w {
            if grid[y][x] && !visited[y][x] {
                let mut comp = Vec::new();
                flood_fill(grid, &mut visited, x, y, &mut comp, w, h);
                components.push(comp);
            }
        }
    }

    for comp in &components {
        if comp.len() < min_component_size {
            for &(x, y) in comp {
                grid[y][x] = false;
            }
        }
    }
}

fn flood_fill(
    grid: &[Vec<bool>],
    visited: &mut [Vec<bool>],
    x: usize,
    y: usize,
    comp: &mut Vec<(usize, usize)>,
    w: usize,
    h: usize,
) {
    let mut stack = vec![(x, y)];
    while let Some((cx, cy)) = stack.pop() {
        if cx >= w || cy >= h || visited[cy][cx] || !grid[cy][cx] {
            continue;
        }
        visited[cy][cx] = true;
        comp.push((cx, cy));
        if cx > 0 { stack.push((cx - 1, cy)); }
        if cy > 0 { stack.push((cx, cy - 1)); }
        if cx + 1 < w { stack.push((cx + 1, cy)); }
        if cy + 1 < h { stack.push((cx, cy + 1)); }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use image::{DynamicImage, GrayImage, Luma};

    #[test]
    fn otsu_all_white() {
        let img = GrayImage::from_pixel(10, 10, Luma([255u8]));
        assert_eq!(otsu_threshold(&img), 0);
    }

    #[test]
    fn otsu_all_black() {
        let img = GrayImage::from_pixel(10, 10, Luma([0u8]));
        assert_eq!(otsu_threshold(&img), 0);
    }

    #[test]
    fn otsu_half_half() {
        let mut img = GrayImage::from_pixel(10, 10, Luma([255u8]));
        for y in 0..10 {
            for x in 0..5 {
                img.put_pixel(x, y, Luma([0u8]));
            }
        }
        let t = otsu_threshold(&img);
        assert!(t < 255);
    }

    #[test]
    fn binarize_half_half() {
        let mut luma = GrayImage::from_pixel(10, 10, Luma([255u8]));
        for y in 0..10u32 {
            for x in 0..5u32 {
                luma.put_pixel(x, y, image::Luma([0u8]));
            }
        }
        let img = DynamicImage::ImageLuma8(luma);
        let grid = binarize(&img);
        let left_black: usize = grid.iter().map(|r| r[..5].iter().filter(|&&b| b).count()).sum();
        assert!(left_black > 0);
    }

    #[test]
    fn clean_noise_removes_small_components() {
        let mut grid = vec![
            vec![false, false, false, false, false, false],
            vec![false, true,  true,  false, true,  false],
            vec![false, false, false, false, false, false],
            vec![false, false, false, false, false, false],
        ];
        clean_noise(&mut grid, 2);
        assert!(grid[1][1]);
        assert!(grid[1][2]);
        assert!(!grid[1][4]);
    }
}