-use atools::prelude::*;

-

-pub struct KD {

- median: [f32; 4],

- normal: [f32; 4],

- index: u32,

- left: Option<Box<KD>>,

- right: Option<Box<KD>>,

- depth: u32,

-}

-

-impl std::fmt::Debug for KD {

- fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {

- write!(f, "#{} {:?}", self.index, self.median)?;

- let d = self.depth as usize;

- if let Some(ref left) = self.left {

- write!(f, "\n{} ⬅️ {left:?}", " ".repeat(d))?;

- }

- if let Some(ref right) = self.right {

- write!(f, "\n{} ➡️ {right:?}", " ".repeat(d))?;

- }

-

- Ok(())

- }

-}

-

-struct KDNearest {

- index: u32,

- distance: f32,

-}

-

-trait Dot {

- fn dot(self, other: Self) -> f32;

-}

-

-impl Dot for [f32; 4] {

- fn dot(self, other: Self) -> f32 {

- self.amul(other).sum()

- }

-}

-

-impl KD {

- pub fn new(colors: &[[f32; 4]]) -> Self {

- assert!(colors.len() < u32::MAX as usize);

- let colors = colors.iter().copied().zip(0u32..).collect::<Vec<_>>();

- Self::_new(&mut { colors }, 0).unwrap()

- }

-

- fn _new(colors: &mut [([f32; 4], u32)], depth: u32) -> Option<KD> {

- if colors.len() == 0 {

- return None;

- };

-

- let middle = colors.len() / 2;

-

- let mut sum = [0.; 4];

- let mut sum2 = [0.; 4];

- for &(c, _) in colors.iter() {

- sum = sum.aadd(c);

- sum2 = sum2.aadd(c.amul(c));

- }

- let [r, g, b, a] = { sum2.asub(sum.amul(sum).mul(1.0 / colors.len() as f32)) };

-

- let normal = [

- (r > g && r > b && r > a) as u8 as f32,

- (g > b && g > a) as u8 as f32,

- (b > a) as u8 as f32,

- 1.,

- ];

- // colors.sort_by(|(a, _), (b, _)| );

- let (before, median, after) = colors.select_nth_unstable_by(middle, |(a, _), (b, _)| {

- a.dot(normal).partial_cmp(&b.dot(normal)).unwrap()

- });

- Some(KD {

- median: median.0,

- index: median.1,

- left: KD::_new(before, depth + 1).map(Box::new),

- right: KD::_new(after, depth + 1).map(Box::new),

- normal,

- depth,

- })

- }

-

- fn _find_nearest(&self, needle: [f32; 4], mut limit: f32) -> Option<KDNearest> {

- let mut result = None;

-

- let diff = needle.asub(self.median);

- let distance = diff.dot(diff).sqrt();

-

- if distance < limit {

- limit = distance;

- result = Some(KDNearest {

- index: self.index,

- distance,

- })

- }

-

- let dot = diff.dot(self.normal);

- if dot <= 0.0 {

- if let Some(ref left) = self.left

- && let Some(nearest) = left._find_nearest(needle, limit)

- {

- limit = nearest.distance;

- result = Some(nearest);

- }

-

- if -dot < limit {

- if let Some(ref right) = self.right

- && let Some(nearest) = right._find_nearest(needle, limit)

- {

- result = Some(nearest);

- }

- }

- } else {

- if let Some(ref right) = self.right

- && let Some(nearest) = right._find_nearest(needle, limit)

- {

- limit = nearest.distance;

- result = Some(nearest);

- }

-

- if dot < limit {

- if let Some(ref left) = self.left

- && let Some(nearest) = left._find_nearest(needle, limit)

- {

- result = Some(nearest);

- }

- }

- }

-

- result

- }

-

- fn _find_nearest_excepting(

- &self,

- needle: [f32; 4],

- mut limit: f32,

- ignore_index: u32,

- ) -> Option<KDNearest> {

- let mut result = None;

-

- let diff = needle.asub(self.median);

- let distance = diff.dot(diff).sqrt();

-

- if distance < limit && self.index != ignore_index {

- limit = distance;

- result = Some(KDNearest {

- index: self.index,

- distance,

- })

- }

-

- let dot = diff.dot(self.normal);

- if dot <= 0.0 {

- if let Some(ref left) = self.left

- && let Some(nearest) = left._find_nearest_excepting(needle, limit, ignore_index)

- {

- limit = nearest.distance;

- result = Some(nearest);

- }

-

- if -dot < limit {

- if let Some(ref right) = self.right

- && let Some(nearest) =

- right._find_nearest_excepting(needle, limit, ignore_index)

- {

- result = Some(nearest);

- }

- }

- } else {

- if let Some(ref right) = self.right

- && let Some(nearest) = right._find_nearest_excepting(needle, limit, ignore_index)

- {

- limit = nearest.distance;

- result = Some(nearest);

- }

-

- if dot < limit {

- if let Some(ref left) = self.left

- && let Some(nearest) = left._find_nearest_excepting(needle, limit, ignore_index)

- {

- result = Some(nearest);

- }

- }

- }

-

- result

- }

-

- pub fn find_nearest(&self, color: [f32; 4]) -> u32 {

- self._find_nearest(color, f32::MAX)

- .map(|x| x.index)

- .unwrap_or(0)

- }

-

- pub fn space(&self, colors: &[[f32; 4]]) -> f32 {

- let mut i = colors

- .iter()

- .zip(0..)

- .map(|(c, i)| {

- self._find_nearest_excepting(*c, f32::MAX, i)

- .unwrap()

- .distance

- })

- .array_chunks::<256>();

-

- let (c, sum) = i.by_ref().fold((0.0, 0.0), |(c, sum), chunk| {

- let y = sum_block(chunk) - c;

- let t = sum + y;

- ((t - sum) - y, t)

- });

- (sum + (i.into_remainder().map(sum_block).unwrap_or(0.) - c)) / colors.len() as f32

- }

-}

-

-use std::intrinsics::fadd_algebraic;

-fn sum_block(arr: impl IntoIterator<Item = f32>) -> f32 {

- arr.into_iter().fold(0.0, |x, y| fadd_algebraic(x, y))

-}