From fc7dcd441077fd2d97ed11f2d249f6a7ec8a0c42 Mon Sep 17 00:00:00 2001
From: DarthSim <darthsim@gmail.com>
Date: Tue, 12 Apr 2022 13:22:57 +0600
Subject: [PATCH] Better quantization and dithering

---
 Cargo.lock       |  27 ++++++
 Cargo.toml       |   1 +
 src/cluster.rs   | 139 +++++++++++--------------------
 src/colormap.rs  | 208 +++++++++++++++++++++++++----------------------
 src/histogram.rs |  58 ++++++++++---
 src/quantize.rs  | 154 +++++++++++++++++------------------
 6 files changed, 312 insertions(+), 275 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index b167858..1858c65 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -2,6 +2,33 @@
 # It is not intended for manual editing.
 version = 3
 
+[[package]]
+name = "autocfg"
+version = "1.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "d468802bab17cbc0cc575e9b053f41e72aa36bfa6b7f55e3529ffa43161b97fa"
+
+[[package]]
+name = "num-traits"
+version = "0.2.14"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "9a64b1ec5cda2586e284722486d802acf1f7dbdc623e2bfc57e65ca1cd099290"
+dependencies = [
+ "autocfg",
+]
+
 [[package]]
 name = "quantizr"
 version = "1.1.0"
+dependencies = [
+ "vpsearch",
+]
+
+[[package]]
+name = "vpsearch"
+version = "2.0.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "98657e55d93988f2a6780e2548ef4fde753c90b0ac1f91149e30eef68abf4260"
+dependencies = [
+ "num-traits",
+]
diff --git a/Cargo.toml b/Cargo.toml
index 0db20dd..58853a2 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -17,5 +17,6 @@ debug-assertions = false
 "capi" = []
 
 [dependencies]
+vpsearch = "2.0.1"
 
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
diff --git a/src/cluster.rs b/src/cluster.rs
index 599936e..0f11f17 100644
--- a/src/cluster.rs
+++ b/src/cluster.rs
@@ -1,44 +1,28 @@
-use std::cmp::{Ord,Ordering};
+use std::cmp::Ordering;
 
 use crate::histogram::{Histogram, HistogramEntry};
 
 pub struct Cluster {
     pub entries: Vec<HistogramEntry>,
-    pub mean: [u8; 4],
-    pub weight: usize,
-    pub priority: u64,
+    pub mean: [f64; 4],
+    pub weight: f64,
+    pub chan_diff: f64,
     widest_chan: u8,
 }
 
-impl Ord for Cluster {
-    fn cmp(&self, other: &Self) -> Ordering {
-        self.priority.cmp(&other.priority)
-    }
-}
-
-impl PartialOrd for Cluster {
-    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
-        Some(self.cmp(other))
-    }
-}
-
-impl Eq for Cluster {}
-
-impl PartialEq for Cluster {
-    fn eq(&self, other: &Self) -> bool {
-        self.priority == other.priority
-    }
-}
-
 impl Cluster {
     pub fn new(entries: Vec<HistogramEntry>) -> Self {
-        Self{
+        let mut cluster = Self{
             entries: entries,
-            mean: [0; 4],
-            priority: 0,
-            weight: 0,
+            mean: [0.0; 4],
+            weight: 0.0,
+            chan_diff: 0.0,
             widest_chan: 0,
-        }
+        };
+
+        cluster.calc_stats();
+
+        cluster
     }
 
     pub fn from_histogram(hist: &Histogram) -> Self {
@@ -51,33 +35,47 @@ impl Cluster {
         Self::new(entries)
     }
 
-    pub fn calc_widest_and_priority(&mut self) {
-        let mut diff_sum: [usize; 4] = [0; 4];
+    fn calc_stats(&mut self) {
+        self.mean = [0.0; 4];
+        self.weight = 0.0;
+
+        if self.entries.is_empty() {
+            self.chan_diff = 0.0;
+            return
+        }
 
         for e in self.entries.iter() {
-            let weight = e.weight as usize;
+            let weight = e.weight as f64;
+
+            self.mean[0] += e.color[0] as f64 * weight;
+            self.mean[1] += e.color[1] as f64 * weight;
+            self.mean[2] += e.color[2] as f64 * weight;
+            self.mean[3] += e.color[3] as f64 * weight;
 
-            diff_sum[0] += diff(e.color[0], self.mean[0]) as usize * weight;
-            diff_sum[1] += diff(e.color[1], self.mean[1]) as usize * weight;
-            diff_sum[2] += diff(e.color[2], self.mean[2]) as usize * weight;
-            diff_sum[3] += diff(e.color[3], self.mean[3]) as usize * weight;
+            self.weight += weight;
         }
 
-        let mut chan = 0;
-        let mut max_diff_sum = 0;
+        self.mean[0] /= self.weight;
+        self.mean[1] /= self.weight;
+        self.mean[2] /= self.weight;
+        self.mean[3] /= self.weight;
 
-        for ch in 0..=3usize {
-            let d = diff_sum[ch];
+        let mut diff_sum: [f64; 4] = [0f64; 4];
 
-            if d > max_diff_sum {
-                chan = ch;
-                max_diff_sum = d;
-            }
+        for e in self.entries.iter() {
+            let weight = e.weight as f64;
+
+            diff_sum[0] += (e.color[0] as f64 - self.mean[0]).abs() * weight;
+            diff_sum[1] += (e.color[1] as f64 - self.mean[1]).abs() * weight;
+            diff_sum[2] += (e.color[2] as f64 - self.mean[2]).abs() * weight;
+            diff_sum[3] += (e.color[3] as f64 - self.mean[3]).abs() * weight;
         }
 
-        let chan_diff = max_diff_sum as f64 / self.weight as f64;
+        let (chan, max_diff_sum) = diff_sum.iter().enumerate()
+            .max_by(|&(_, a), &(_, b)| a.partial_cmp(&b).unwrap_or(Ordering::Equal))
+            .unwrap();
 
-        self.priority = (chan_diff * (self.weight as f64).sqrt()) as u64;
+        self.chan_diff = max_diff_sum / self.weight;
         self.widest_chan = chan as u8;
     }
 
@@ -93,26 +91,27 @@ impl Cluster {
         let mut gt_weight: usize = 0;
 
         while i <= gt {
-            let val = self.entries[i].color[widest_chan];
+            let entry = &self.entries[i];
+            let val = entry.color[widest_chan] as f64;
 
             if val < widest_chan_mean {
+                lt_weight += entry.weight as usize;
                 if lt != i {
                     self.entries.swap(lt, i);
                 }
                 lt += 1;
                 i += 1;
-                lt_weight += self.entries[i].weight as usize;
             } else if val > widest_chan_mean {
+                gt_weight += entry.weight as usize;
                 self.entries.swap(gt, i);
                 gt -= 1;
-                gt_weight += self.entries[i].weight as usize;
             } else {
                 i += 1;
             }
         }
 
         let mut split_pos = i;
-        if lt_weight < gt_weight {
+        if lt_weight > gt_weight {
             split_pos = lt;
         }
 
@@ -120,44 +119,4 @@ impl Cluster {
 
         (Self::new(sp1.to_vec()), Self::new(sp2.to_vec()))
     }
-
-    pub fn calc_mean_and_weight(&mut self) {
-        self.weight = 0;
-
-        if self.entries.is_empty() {
-            self.mean = [0; 4];
-            return
-        }
-
-        let mut rsum: usize = 0;
-        let mut gsum: usize = 0;
-        let mut bsum: usize = 0;
-        let mut asum: usize = 0;
-
-        for e in self.entries.iter() {
-            let weight = e.weight as usize;
-
-            rsum += e.color[0] as usize * weight;
-            gsum += e.color[1] as usize * weight;
-            bsum += e.color[2] as usize * weight;
-            asum += e.color[3] as usize * weight;
-
-            self.weight += weight;
-        }
-
-        self.mean = [
-            (rsum / self.weight) as u8,
-            (gsum / self.weight) as u8,
-            (bsum / self.weight) as u8,
-            (asum / self.weight) as u8,
-        ]
-    }
-}
-
-fn diff(a: u8, b: u8) -> u8 {
-    if a > b {
-        return a - b
-    }
-
-    b - a
 }
diff --git a/src/colormap.rs b/src/colormap.rs
index a5f33a6..a8ea6a9 100644
--- a/src/colormap.rs
+++ b/src/colormap.rs
@@ -1,73 +1,59 @@
 use std::cmp::Ordering;
 
+use vpsearch;
+
 use crate::palette::Palette;
 use crate::cluster::Cluster;
 
-macro_rules! color_dist {
-    ($c1: expr, $c2: expr) => {
-        ($c1[0] - $c2[0]).powi(2) +
-        ($c1[1] - $c2[1]).powi(2) +
-        ($c1[2] - $c2[2]).powi(2) +
-        ($c1[3] - $c2[3]).powi(2)
-    };
-}
+struct ColormapEntry;
+impl vpsearch::MetricSpace<ColormapEntry> for [f32; 4] {
+    type UserData = ();
+    type Distance = f32;
 
-struct ColormapEntry {
-    pub color: [f32; 4],
-    popularity: usize,
-    radius: f32,
+    fn distance(&self, other: &Self, _: &Self::UserData) -> Self::Distance {
+        color_dist(self, other).sqrt()
+    }
 }
 
-pub struct Colormap(Vec<ColormapEntry>);
+pub struct Colormap {
+    entries: Vec<[f32; 4]>,
+    tree: vpsearch::Tree::<[f32; 4], ColormapEntry>,
+    pub error: f32,
+    total_weight: f32,
+}
 
 impl Colormap {
     pub fn new(clusters: &Vec::<Cluster>) -> Self {
-        let mut entries = vec![];
-
-        let mut ind = 0;
-        entries.resize_with(clusters.len() as usize, ||{
-            let c = &clusters[ind];
-            let e = ColormapEntry{
-                color: [c.mean[0] as f32, c.mean[1] as f32, c.mean[2] as f32, c.mean[3] as f32],
-                popularity: c.weight,
-                radius: 0.0,
-            };
-            ind += 1;
-            e
-        });
+        let mut total_weight = 0f32;
 
-        let mut res = Self{0: entries};
-        res.calc_radiuses();
-        res.kmeans(clusters);
-        res
-    }
+        let entries: Vec::<[f32; 4]> = clusters.iter().map(|c|{
+            total_weight += c.weight as f32;
+            [c.mean[0] as f32, c.mean[1] as f32, c.mean[2] as f32, c.mean[3] as f32]
+        }).collect();
+
+        let tree = vpsearch::Tree::new(&entries);
 
-    fn calc_radiuses(&mut self) {
-        let count = self.0.len();
-        let mut nearest = [f32::MAX; 256];
+        let mut res = Self{
+            entries: entries,
+            tree: tree,
+            error: 0f32,
+            total_weight: total_weight,
+        };
 
-        assert!(count <= 256);
+        res.kmeans(clusters);
+        res.tree = vpsearch::Tree::new(&res.entries);
 
-        for i in 0..count-1 {
-            for j in i+1..count {
-                let dist = color_dist!(self.0[i].color, self.0[j].color);
+        res.kmeans(clusters);
+        res.sort();
 
-                nearest[i] = nearest[i].min(dist);
-                nearest[j] = nearest[j].min(dist);
-            }
-        }
+        res.tree = vpsearch::Tree::new(&res.entries);
 
-        for i in 0..count {
-            self.0[i].radius = nearest[i] / 2.0;
-        }
+        res
     }
 
     fn sort(&mut self) {
-        self.0.sort_by(|e1, e2| {
-            match e1.color[3].partial_cmp(&e2.color[3]).unwrap_or(Ordering::Equal) {
-                Ordering::Equal => e2.popularity.cmp(&e1.popularity),
-                o => o,
-            }
+        self.entries.sort_unstable_by(|e1, e2| {
+            e1[3].partial_cmp(&e2[3]).unwrap_or(Ordering::Equal)
         });
     }
 
@@ -75,80 +61,112 @@ impl Colormap {
         let mut colors = [[0f32; 4]; 256];
         let mut weights = [0f32; 256];
 
+        let mut total_err = 0f32;
+
         for cluster in clusters.iter() {
             for entry in cluster.entries.iter() {
-                let r = entry.color[0] as f32;
-                let g = entry.color[1] as f32;
-                let b = entry.color[2] as f32;
-                let a = entry.color[3] as f32;
-
+                let hist_color = [
+                    entry.color[0] as f32,
+                    entry.color[1] as f32,
+                    entry.color[2] as f32,
+                    entry.color[3] as f32,
+                ];
                 let weight = entry.weight as f32;
 
-                let ind = self.nearest_ind(&[r, g, b, a], 0);
+                let (ind, err) = self.nearest_ind(&hist_color);
 
                 let color = &mut colors[ind];
-                color[0] += r * weight;
-                color[1] += g * weight;
-                color[2] += b * weight;
-                color[3] += a * weight;
+                add_color(color, &hist_color, weight);
 
                 weights[ind] += weight;
+                total_err += err*err;
             }
         }
 
-        for (i, c) in colors.iter().enumerate() {
-            let weight = weights[i];
-
+        for ((pal_c, c), weight) in self.entries.iter_mut().zip(colors).zip(weights) {
             if weight > 0.0 {
-                let pal_c = &mut self.0[i];
-                pal_c.color[0] = c[0] / weight;
-                pal_c.color[1] = c[1] / weight;
-                pal_c.color[2] = c[2] / weight;
-                pal_c.color[3] = c[3] / weight;
+                pal_c[0] = c[0] / weight;
+                pal_c[1] = c[1] / weight;
+                pal_c[2] = c[2] / weight;
+                pal_c[3] = c[3] / weight;
             }
         }
 
-        self.sort();
-        self.calc_radiuses()
+        self.error = total_err / self.total_weight;
     }
 
     pub fn generate_palette(&self, palette: &mut Palette) {
-        palette.count = self.0.len() as u32;
+        palette.count = self.entries.len() as u32;
 
-        for (i, e) in self.0.iter().enumerate() {
+        for (i, e) in self.entries.iter().enumerate() {
             let c = &mut palette.entries[i];
-            c.r = e.color[0].round().clamp(0.0, 255.0) as u8;
-            c.g = e.color[1].round().clamp(0.0, 255.0) as u8;
-            c.b = e.color[2].round().clamp(0.0, 255.0) as u8;
-            c.a = e.color[3].round().clamp(0.0, 255.0) as u8;
+            c.r = e[0].round().clamp(0.0, 255.0) as u8;
+            c.g = e[1].round().clamp(0.0, 255.0) as u8;
+            c.b = e[2].round().clamp(0.0, 255.0) as u8;
+            c.a = e[3].round().clamp(0.0, 255.0) as u8;
         }
     }
 
-    pub fn nearest_ind(&self, color: &[f32; 4], last_ind: usize) -> usize {
-        let mut best_ind = last_ind;
+    #[inline(always)]
+    pub fn nearest_ind(&self, color: &[f32; 4]) -> (usize, f32) {
+        self.tree.find_nearest(color)
+    }
 
-        let mut best_dist = color_dist!(self.0[last_ind].color, color);
-        if best_dist <= self.0[last_ind].radius {
-            return last_ind;
-        }
+    pub fn color(&self, ind: usize) -> &[f32; 4] {
+        &self.entries[ind]
+    }
+}
 
-        for (i, e) in self.0.iter().enumerate() {
-            let dist = color_dist!(e.color, color);
+#[cfg(target_arch = "x86_64")]
+#[inline(always)]
+fn add_color(dst: &mut [f32; 4], src: &[f32; 4], weight: f32) {
+    unsafe {
+        use std::arch::x86_64::*;
 
-            if dist <= e.radius {
-                return i
-            }
+        let mut psrc = _mm_loadu_ps(src.as_ptr());
+        let mut pdst = _mm_loadu_ps(dst.as_ptr());
+        let pweights = _mm_set1_ps(weight);
 
-            if dist < best_dist {
-                best_dist = dist;
-                best_ind = i;
-            }
-        }
+        psrc = _mm_mul_ps(psrc, pweights);
+        pdst = _mm_add_ps(pdst, psrc);
 
-        best_ind
+        _mm_storeu_ps(dst.as_mut_ptr(), pdst);
     }
+}
 
-    pub fn color(&self, ind: usize) -> [f32; 4] {
-        self.0[ind].color
+#[cfg(not(target_arch = "x86_64"))]
+#[inline(always)]
+fn add_color(dst: &mut [f32; 4], src: &[f32; 4], weight: f32) {
+    dst[0] += src[0] * weight;
+    dst[1] += src[1] * weight;
+    dst[2] += src[2] * weight;
+    dst[3] += src[3] * weight;
+}
+
+#[cfg(target_arch = "x86_64")]
+#[inline(always)]
+fn color_dist(c1: &[f32; 4], c2: &[f32; 4]) -> f32 {
+    unsafe {
+        use std::arch::x86_64::*;
+
+        let pc1 = _mm_loadu_ps(c1.as_ptr());
+        let pc2 = _mm_loadu_ps(c2.as_ptr());
+
+        let mut dist = _mm_sub_ps(pc1, pc2);
+        dist = _mm_mul_ps(dist, dist);
+
+        let mut tmp = [0f32; 4];
+        _mm_storeu_ps(tmp.as_mut_ptr(), dist);
+
+        tmp[0] + tmp[1] + tmp[2] + tmp[3]
     }
 }
+
+#[cfg(not(target_arch = "x86_64"))]
+#[inline(always)]
+fn color_dist(c1: &[f32; 4], c2: &[f32; 4]) -> f32 {
+    (c1[0] - c2[0]).powi(2) +
+    (c1[1] - c2[1]).powi(2) +
+    (c1[2] - c2[2]).powi(2) +
+    (c1[3] - c2[3]).powi(2)
+}
diff --git a/src/histogram.rs b/src/histogram.rs
index 3978403..ac21cdb 100644
--- a/src/histogram.rs
+++ b/src/histogram.rs
@@ -1,13 +1,14 @@
 use std::collections::HashMap;
+use std::hash::{BuildHasher,Hasher};
 
 use crate::image::Image;
 
 macro_rules! hist_key {
     ($c: expr) => {
-        $c[0] as usize +
-        $c[1] as usize * 256 +
-        $c[2] as usize * 256 * 256 +
-        $c[3] as usize * 256 * 256 * 256
+        (($c[0] as u64) << 24) +
+        (($c[1] as u64) << 16) +
+        (($c[2] as u64) << 8) +
+        ($c[3] as u64)
     };
 }
 
@@ -17,28 +18,61 @@ pub struct HistogramEntry {
     pub weight: u32,
 }
 
-pub struct Histogram(pub HashMap<usize, HistogramEntry>);
+pub struct Histogram(pub HashMap<u64, HistogramEntry, ColorHasher>);
 
 impl Histogram {
     pub fn new() -> Self {
-        Self{0: HashMap::new()}
+        Self{0: HashMap::with_hasher(ColorHasher(0))}
     }
 
     pub fn add_image(&mut self, image: &Image) {
         for ind in (0..image.width*image.height*4).step_by(4) {
             let pix = &image.data[ind..ind+4];
-            let color = [pix[0], pix[1], pix[2], pix[3]];
+            let color = if pix[3] != 0 {
+                [pix[0], pix[1], pix[2], pix[3]]
+            } else {
+                [0, 0, 0, 0]
+            };
 
             let key = hist_key!(color);
 
-            if let Some(e) = self.0.get_mut(&key) {
-                e.weight += 1;
-            } else {
-                self.0.insert(key, HistogramEntry{
+            self.0.entry(key)
+                .and_modify(|e| e.weight += 1)
+                .or_insert(HistogramEntry{
                     color: color,
                     weight: 1,
                 });
-            }
         }
     }
 }
+
+pub struct ColorHasher(u64);
+impl BuildHasher for ColorHasher {
+    type Hasher = Self;
+    #[inline(always)]
+    fn build_hasher(&self) -> Self {
+        Self(0)
+    }
+}
+
+impl Hasher for ColorHasher {
+    // Magick number from https://github.com/cbreeden/fxhash/blob/master/lib.rs
+    #[inline(always)]
+    fn finish(&self) -> u64 { self.0.wrapping_mul(0x517cc1b727220a95) }
+
+    #[inline(always)]
+    fn write_u64(&mut self, i: u64) { self.0 = i; }
+
+    fn write(&mut self, _bytes: &[u8]) { unimplemented!() }
+    fn write_u8(&mut self, _i: u8) { unimplemented!() }
+    fn write_u16(&mut self, _i: u16) { unimplemented!() }
+    fn write_u32(&mut self, _i: u32) { unimplemented!() }
+    fn write_u128(&mut self, _i: u128) { unimplemented!() }
+    fn write_usize(&mut self, _i: usize) { unimplemented!() }
+    fn write_i8(&mut self, _i: i8) { unimplemented!() }
+    fn write_i16(&mut self, _i: i16) { unimplemented!() }
+    fn write_i32(&mut self, _i: i32) { unimplemented!() }
+    fn write_i64(&mut self, _i: i64) { unimplemented!() }
+    fn write_i128(&mut self, _i: i128) { unimplemented!() }
+    fn write_isize(&mut self, _i: isize) { unimplemented!() }
+}
diff --git a/src/quantize.rs b/src/quantize.rs
index df226a3..c5700f7 100644
--- a/src/quantize.rs
+++ b/src/quantize.rs
@@ -1,4 +1,4 @@
-use std::collections::BinaryHeap;
+use std::cmp::Ordering;
 
 use crate::histogram::Histogram;
 use crate::cluster::Cluster;
@@ -23,68 +23,52 @@ impl QuantizeResult {
     }
 
     pub fn quantize_histogram(hist: &Histogram, attr: &Options) -> Self {
-        let mut heap = BinaryHeap::new();
         let mut clusters = Vec::<Cluster>::with_capacity(attr.max_colors as usize);
 
-        let mut root = Cluster::from_histogram(&hist);
-        root.calc_mean_and_weight();
-        root.calc_widest_and_priority();
+        let root = Cluster::from_histogram(&hist);
 
-        // If priority is zero, then all colors in cluster are the same
-        if root.priority > 0 {
-            heap.push(root);
-        } else {
-            clusters.push(root);
-        }
+        clusters.push(root);
+
+        let max_colors_f64 = attr.max_colors as f64;
+        let max_colors_usize = attr.max_colors as usize;
+
+        while clusters.len() < max_colors_usize {
+            // We want to split bigger clusters in the beginning,
+            // and clusters with bigger chan_diff in the end
+            let weight_ratio = 0.75 - (clusters.len() as f64 + 1.0) / max_colors_f64 / 2.0;
+
+            // Get the best cluster to split
+            let to_split_opt = clusters.iter().enumerate()
+                .filter(|(_, c)| c.chan_diff > 0.0)
+                .map(|(i, c)|{
+                    let priority = c.chan_diff * c.weight.powf(weight_ratio);
+                    (i, priority)
+                })
+                .max_by(|&(_, a), &(_, b)| a.partial_cmp(&b).unwrap_or(Ordering::Equal))
+                .map(|(i, _)| clusters.swap_remove(i) );
 
-        loop {
-            // Try to pop cluster from queue
             // If nothing there, this means everything is ready
-            let mut to_split = match heap.pop() {
+            let mut to_split = match to_split_opt {
                 Some(c) => c,
                 None => break,
             };
 
             let (mut c1, mut c2) = to_split.split();
 
-            c1.calc_mean_and_weight();
-            c2.calc_mean_and_weight();
-
             if c1.entries.is_empty() {
+                c2.chan_diff = 0.0;
                 clusters.push(c2);
                 continue;
             }
 
             if c2.entries.is_empty() {
+                c1.chan_diff = 0.0;
                 clusters.push(c1);
                 continue;
             }
 
-            let colors = clusters.len() + heap.len() + 2;
-
-            // Looks like we reached the maximum of colors
-            // Add new clusters to ready and flush queue
-            if attr.max_colors == colors as i32 {
-                clusters.push(c1);
-                clusters.push(c2);
-                clusters.append(&mut heap.into_vec());
-                break
-            }
-
-            c1.calc_widest_and_priority();
-            c2.calc_widest_and_priority();
-
-            if c1.priority > 0 {
-                heap.push(c1);
-            } else {
-                clusters.push(c1)
-            }
-
-            if c2.priority > 0 {
-                heap.push(c2);
-            } else {
-                clusters.push(c2)
-            }
+            clusters.push(c1);
+            clusters.push(c2);
         }
 
         let mut palette = Palette::default();
@@ -124,9 +108,7 @@ impl QuantizeResult {
         Error::Ok
     }
 
-    fn remap_image_no_dither(&self, image: &Image, buf: &mut [u8]) {
-        let mut last_ind = 0;
-
+    fn remap_image_no_dither(&mut self, image: &Image, buf: &mut [u8]) {
         for point in 0..image.width*image.height {
             let data_point = point*4;
 
@@ -136,18 +118,20 @@ impl QuantizeResult {
             let b = pix[2] as f32;
             let a = pix[3] as f32;
 
-            last_ind = self.colormap.nearest_ind(&[r, g, b, a], last_ind);
+            let (ind, _) = self.colormap.nearest_ind(&[r, g, b, a]);
 
-            buf[point] = last_ind as u8;
+            buf[point] = ind as u8;
         }
     }
 
-    fn remap_image_dither(&self, image: &Image, buf: &mut [u8]) {
-        let mut last_ind = 0;
-
+    fn remap_image_dither(&mut self, image: &Image, buf: &mut [u8]) {
         let error_size = image.width+2;
-        let mut error_curr = vec![[0f32; 4]; error_size].into_boxed_slice();
-        let mut error_next = vec![[0f32; 4]; error_size].into_boxed_slice();
+        let mut error_curr = vec![[0f32; 4]; error_size];
+        let mut error_next = vec![[0f32; 4]; error_size];
+
+        let dithering_coeff = self.dithering_level * 15.0 / 16.0 / 16.0;
+        let err_threshold = self.colormap.error;
+        // println!("Err threshold {}", self.colormap.error);
 
         let mut x_reverse = true;
 
@@ -169,26 +153,45 @@ impl QuantizeResult {
                     true => [err_ind + 1, err_ind, err_ind - 1],
                 };
 
+                let err_pix = &mut error_curr[err_ind];
+
+                let err_total = err_pix[0].powi(2) + err_pix[1].powi(2) + err_pix[2].powi(2) + err_pix[3].powi(2);
+                if err_total > err_threshold {
+                    err_pix[0] = err_pix[0] * 0.8;
+                    err_pix[1] = err_pix[1] * 0.8;
+                    err_pix[2] = err_pix[2] * 0.8;
+                    err_pix[3] = err_pix[3] * 0.8;
+                }
+
                 let pix = &image.data[data_point..data_point+4];
-                let r = pix[0] as f32;
-                let g = pix[1] as f32;
-                let b = pix[2] as f32;
-                let a = pix[3] as f32;
-
-                let err_pix = &error_curr[err_ind];
-                let dr = (r + err_pix[0]).clamp(-0.1, 255.1);
-                let dg = (g + err_pix[1]).clamp(-0.1, 255.1);
-                let db = (b + err_pix[2]).clamp(-0.1, 255.1);
-                let da = (a + err_pix[3]).clamp(-0.1, 255.1);
-
-                last_ind = self.colormap.nearest_ind(&[dr, dg, db, da], last_ind);
-                buf[point] = last_ind as u8;
-
-                let pal_pix = self.colormap.color(last_ind);
-                let err_r = (r - pal_pix[0]) / 16.0 * self.dithering_level;
-                let err_g = (g - pal_pix[1]) / 16.0 * self.dithering_level;
-                let err_b = (b - pal_pix[2]) / 16.0 * self.dithering_level;
-                let err_a = (a - pal_pix[3]) / 16.0 * self.dithering_level;
+                let dith_pix = [
+                    pix[0] as f32 + err_pix[0],
+                    pix[1] as f32 + err_pix[1],
+                    pix[2] as f32 + err_pix[2],
+                    pix[3] as f32 + err_pix[3],
+                ];
+
+                let (ind, _) = self.colormap.nearest_ind(&dith_pix);
+                buf[point] = ind as u8;
+
+                let pal_pix = self.colormap.color(ind);
+                let mut err_r = dith_pix[0] - pal_pix[0];
+                let mut err_g = dith_pix[1] - pal_pix[1];
+                let mut err_b = dith_pix[2] - pal_pix[2];
+                let mut err_a = dith_pix[3] - pal_pix[3];
+
+                let err_total = err_r.powi(2) + err_g.powi(2) + err_b.powi(2) + err_a.powi(2);
+                if err_total > err_threshold {
+                    err_r = err_r * 0.75;
+                    err_g = err_g * 0.75;
+                    err_b = err_b * 0.75;
+                    err_a = err_a * 0.75;
+                }
+
+                err_r = err_r * dithering_coeff;
+                err_g = err_g * dithering_coeff;
+                err_b = err_b * dithering_coeff;
+                err_a = err_a * dithering_coeff;
 
                 let err = &mut error_next[err_inds[0]];
                 err[0] += err_r * 3.0;
@@ -227,13 +230,8 @@ impl QuantizeResult {
                 }
             }
 
-            let tmp_err = error_curr;
-            error_curr = error_next;
-            error_next = tmp_err;
-
-            for i in 0..error_size {
-                error_next[i] = [0.0, 0.0, 0.0, 0.0];
-            }
+            std::mem::swap(&mut error_curr, &mut error_next);
+            error_next.fill_with(|| [0f32; 4]);
         }
     }
 }