//! Colour reduction and noise removal.

use crate::{Mode, Raster, TraceParams};

/// Label of a transparent (untraced) pixel.
pub const TRANSPARENT: u16 = u16::MAX;

/// A palette-indexed image.
#[derive(Clone, Debug, PartialEq)]
pub struct Quantized {
    /// One palette index per pixel (row-major), or [`TRANSPARENT`].
    pub labels: Vec<u16>,
    pub palette: Vec<[u8; 3]>,
}

fn luma(p: [u8; 4]) -> u8 {
    ((p[0] as u32 * 299 + p[1] as u32 * 587 + p[2] as u32 * 114 + 500) / 1000) as u8
}

/// Reduce `img` to a palette according to `params.mode`.
pub fn quantize(img: &Raster, params: &TraceParams) -> Quantized {
    let px = img.rgba.as_chunks::<4>().0.iter().map(|c| [c[0], c[1], c[2], c[3]]);
    match params.mode {
        Mode::BlackAndWhite => {
            let t = params.threshold;
            let labels = px
                .map(|p| {
                    if p[3] < 128 {
                        TRANSPARENT
                    } else if luma(p) < t {
                        1
                    } else {
                        0
                    }
                })
                .collect();
            Quantized { labels, palette: vec![[255, 255, 255], [0, 0, 0]] }
        }
        Mode::Grayscale => {
            let mut hist = [0u64; 256];
            for p in img.rgba.as_chunks::<4>().0 {
                if p[3] >= 128 {
                    hist[luma([p[0], p[1], p[2], p[3]]) as usize] += 1;
                }
            }
            let levels = kmeans_1d(&hist, params.colors.clamp(2, 256) as usize);
            let mut lut = [0u16; 256];
            for (v, slot) in lut.iter_mut().enumerate() {
                *slot = nearest_1d(&levels, v as f64) as u16;
            }
            let labels = px.map(|p| if p[3] < 128 { TRANSPARENT } else { lut[luma(p) as usize] }).collect();
            let palette = levels.iter().map(|&l| {
                let g = l.round().clamp(0.0, 255.0) as u8;
                [g, g, g]
            });
            compact(labels, palette.collect())
        }
        Mode::Color => {
            // 15-bit colour histogram with per-bin sums (exact bin means).
            let mut bins: Vec<Bin> = vec![Bin::default(); 1 << 15];
            for p in img.rgba.as_chunks::<4>().0 {
                if p[3] >= 128 {
                    let b = &mut bins[key(p[0], p[1], p[2])];
                    b.n += 1;
                    b.sum[0] += p[0] as u64;
                    b.sum[1] += p[1] as u64;
                    b.sum[2] += p[2] as u64;
                }
            }
            let used: Vec<(usize, [f64; 3], f64)> =
                bins.iter().enumerate().filter(|(_, b)| b.n > 0).map(|(i, b)| (i, b.mean(), b.n as f64)).collect();
            let k = params.colors.clamp(2, 256) as usize;
            let centers = kmeans_rgb(&used, k);
            let mut lut = vec![0u16; 1 << 15];
            for (i, c, _) in &used {
                lut[*i] = nearest_rgb(&centers, *c) as u16;
            }
            let labels = px.map(|p| if p[3] < 128 { TRANSPARENT } else { lut[key(p[0], p[1], p[2])] }).collect();
            let palette = centers.iter().map(|c| [c[0].round() as u8, c[1].round() as u8, c[2].round() as u8]).collect();
            compact(labels, palette)
        }
    }
}

/// Drop unused palette entries and merge duplicates.
fn compact(mut labels: Vec<u16>, palette: Vec<[u8; 3]>) -> Quantized {
    let mut used = vec![false; palette.len()];
    for &l in &labels {
        if l != TRANSPARENT {
            used[l as usize] = true;
        }
    }
    let mut remap = vec![0u16; palette.len()];
    let mut out: Vec<[u8; 3]> = vec![];
    for (i, c) in palette.iter().enumerate() {
        if !used[i] {
            continue;
        }
        remap[i] = match out.iter().position(|o| o == c) {
            Some(j) => j as u16,
            None => {
                out.push(*c);
                (out.len() - 1) as u16
            }
        };
    }
    for l in &mut labels {
        if *l != TRANSPARENT {
            *l = remap[*l as usize];
        }
    }
    Quantized { labels, palette: out }
}

#[derive(Clone, Copy, Default)]
struct Bin {
    n: u64,
    sum: [u64; 3],
}

impl Bin {
    fn mean(&self) -> [f64; 3] {
        let n = self.n.max(1) as f64;
        [self.sum[0] as f64 / n, self.sum[1] as f64 / n, self.sum[2] as f64 / n]
    }
}

fn key(r: u8, g: u8, b: u8) -> usize {
    ((r as usize >> 3) << 10) | ((g as usize >> 3) << 5) | (b as usize >> 3)
}

fn nearest_1d(levels: &[f64], v: f64) -> usize {
    let mut best = (f64::INFINITY, 0);
    for (i, &l) in levels.iter().enumerate() {
        let d = (l - v).abs();
        if d < best.0 {
            best = (d, i);
        }
    }
    best.1
}

/// 1-D k-means over a 256-bin histogram (Lloyd iterations, evenly spaced start).
fn kmeans_1d(hist: &[u64; 256], k: usize) -> Vec<f64> {
    let lo = hist.iter().position(|&n| n > 0).unwrap_or(0) as f64;
    let hi = hist.iter().rposition(|&n| n > 0).unwrap_or(255) as f64;
    let mut levels: Vec<f64> = (0..k).map(|i| if k == 1 { lo } else { lo + (hi - lo) * i as f64 / (k - 1) as f64 }).collect();
    for _ in 0..20 {
        let mut sum = vec![0.0; k];
        let mut cnt = vec![0.0; k];
        for (v, &n) in hist.iter().enumerate() {
            if n > 0 {
                let j = nearest_1d(&levels, v as f64);
                sum[j] += v as f64 * n as f64;
                cnt[j] += n as f64;
            }
        }
        let mut moved = false;
        for j in 0..k {
            if cnt[j] > 0.0 {
                let m = sum[j] / cnt[j];
                moved |= (m - levels[j]).abs() > 0.01;
                levels[j] = m;
            }
        }
        if !moved {
            break;
        }
    }
    levels
}

fn dist2(a: [f64; 3], b: [f64; 3]) -> f64 {
    // Slightly perceptual weights.
    let (dr, dg, db) = (a[0] - b[0], a[1] - b[1], a[2] - b[2]);
    2.0 * dr * dr + 4.0 * dg * dg + 3.0 * db * db
}

fn nearest_rgb(centers: &[[f64; 3]], c: [f64; 3]) -> usize {
    let mut best = (f64::INFINITY, 0);
    for (i, &m) in centers.iter().enumerate() {
        let d = dist2(m, c);
        if d < best.0 {
            best = (d, i);
        }
    }
    best.1
}

/// Median cut for the initial palette, then weighted k-means refinement over histogram bins.
fn kmeans_rgb(used: &[(usize, [f64; 3], f64)], k: usize) -> Vec<[f64; 3]> {
    if used.is_empty() {
        return vec![[255.0; 3]];
    }
    // Median cut: repeatedly split the box with the largest weighted extent.
    let mut boxes: Vec<Vec<usize>> = vec![(0..used.len()).collect()];
    while boxes.len() < k {
        let mut best: Option<(f64, usize, usize)> = None;
        for (bi, b) in boxes.iter().enumerate() {
            if b.len() < 2 {
                continue;
            }
            let (mut lo, mut hi) = ([f64::INFINITY; 3], [f64::NEG_INFINITY; 3]);
            let mut wsum = 0.0;
            for &i in b {
                for ch in 0..3 {
                    lo[ch] = lo[ch].min(used[i].1[ch]);
                    hi[ch] = hi[ch].max(used[i].1[ch]);
                }
                wsum += used[i].2;
            }
            let (ch, range) = (0..3).map(|c| (c, hi[c] - lo[c])).fold((0, -1.0), |a, b| if b.1 > a.1 { b } else { a });
            if range <= 0.0 {
                continue;
            }
            let score = range * wsum.sqrt();
            if best.is_none_or(|x| score > x.0) {
                best = Some((score, bi, ch));
            }
        }
        let Some((_, bi, ch)) = best else { break };
        let mut b = boxes.swap_remove(bi);
        b.sort_by(|x, y| used[*x].1[ch].total_cmp(&used[*y].1[ch]));
        let total: f64 = b.iter().map(|&i| used[i].2).sum();
        let mut acc = 0.0;
        let mut cut = 1;
        for (j, &i) in b.iter().enumerate() {
            acc += used[i].2;
            if acc >= total / 2.0 {
                cut = (j + 1).clamp(1, b.len() - 1);
                break;
            }
        }
        let tail = b.split_off(cut);
        boxes.push(b);
        boxes.push(tail);
    }
    let mut centers: Vec<[f64; 3]> = boxes
        .iter()
        .map(|b| {
            let mut s = [0.0; 3];
            let mut w = 0.0;
            for &i in b {
                for (acc, v) in s.iter_mut().zip(used[i].1) {
                    *acc += v * used[i].2;
                }
                w += used[i].2;
            }
            [s[0] / w, s[1] / w, s[2] / w]
        })
        .collect();
    // Lloyd refinement.
    for _ in 0..12 {
        let mut sum = vec![[0.0; 3]; centers.len()];
        let mut cnt = vec![0.0; centers.len()];
        for (_, c, n) in used {
            let j = nearest_rgb(&centers, *c);
            for ch in 0..3 {
                sum[j][ch] += c[ch] * n;
            }
            cnt[j] += n;
        }
        let mut moved = 0.0f64;
        for j in 0..centers.len() {
            if cnt[j] > 0.0 {
                let m = [sum[j][0] / cnt[j], sum[j][1] / cnt[j], sum[j][2] / cnt[j]];
                moved = moved.max(dist2(m, centers[j]));
                centers[j] = m;
            }
        }
        if moved < 0.25 {
            break;
        }
    }
    centers
}

/// Merge 4-connected same-label components smaller than `min_area` pixels into their most common
/// neighbouring label. Transparent pixels are never relabelled.
pub fn denoise(labels: &mut [u16], w: usize, h: usize, min_area: usize) {
    if min_area <= 1 || w == 0 || h == 0 {
        return;
    }
    let mut seen = vec![false; w * h];
    let mut stack: Vec<usize> = Vec::new();
    let mut comp: Vec<usize> = Vec::new();
    let mut neigh: Vec<(u16, u32)> = Vec::new();
    for start in 0..w * h {
        if seen[start] || labels[start] == TRANSPARENT {
            continue;
        }
        let l = labels[start];
        comp.clear();
        neigh.clear();
        stack.push(start);
        seen[start] = true;
        let mut small = true;
        while let Some(i) = stack.pop() {
            if small {
                comp.push(i);
                if comp.len() >= min_area {
                    small = false;
                }
            }
            let (x, y) = (i % w, i / w);
            let mut visit = |j: usize| {
                let lj = labels[j];
                if lj == l {
                    if !seen[j] {
                        seen[j] = true;
                        stack.push(j);
                    }
                } else if lj != TRANSPARENT {
                    match neigh.iter_mut().find(|(n, _)| *n == lj) {
                        Some(e) => e.1 += 1,
                        None => neigh.push((lj, 1)),
                    }
                }
            };
            if x > 0 {
                visit(i - 1);
            }
            if x + 1 < w {
                visit(i + 1);
            }
            if y > 0 {
                visit(i - w);
            }
            if y + 1 < h {
                visit(i + w);
            }
        }
        if small && let Some(&(to, _)) = neigh.iter().max_by_key(|(n, c)| (*c, u16::MAX - *n)) {
            for &i in &comp {
                labels[i] = to;
            }
        }
    }
}