1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
/// A data structure for fast nearest color lookups in a palette.
use atools::prelude::*;

pub struct KD {
    mid_point: [f32; 4],
    index: u32,
    normal: [f32; 4],
    left: Option<Box<KD>>,
    right: Option<Box<KD>>,
}

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 mut x = (0..colors.len() as u32).collect::<Vec<_>>();

        Self::_new(&mut x, colors)
    }

    fn _new(mut indices: &mut [u32], colors: &[[f32; 4]]) -> KD {
        assert!(indices.len() > 0);

        let middle = indices.len() / 2;

        let mut sum = [0.; 4];
        let mut sum2 = [0.; 4];
        for i in indices.iter() {
            let c = colors[*i as usize];
            sum = sum.aadd(c);
            sum2 = sum2.aadd(c.amul(c));
        }
        let [r, g, b, a] = { sum2.asub(sum.amul(sum).mul(1.0 / indices.len() as f32)) };
        let normal = if r > g && r > b && r > a {
            (1.0).join([0.; 3])
        } else if g > b && g > a {
            [0.0, 1.0].couple([0.; 2])
        } else if b > a {
            [0.; 2].couple([1., 0.])
        } else {
            [0.; 3].join(1.)
        };
        indices.sort_by(|a, b| {
            colors[*a as usize]
                .dot(normal)
                .partial_cmp(&colors[*b as usize].dot(normal))
                .unwrap()
        });
        let i = indices.len() / 2;
        let left = if i > 0 {
            Some(Box::new(KD::_new(&mut indices[0..i], colors)))
        } else {
            None
        };
        let right = if i + 1 < indices.len() {
            Some(Box::new(KD::_new(&mut indices[(i + 1)..], colors)))
        } else {
            None
        };
        KD {
            mid_point: colors[indices[i as usize] as usize],
            index: indices[i as usize],
            normal: normal,
            left: left,
            right: right,
        }
    }

    fn _find_nearest(
        &self,
        needle: [f32; 4],
        mut limit: f32,
        ignore_index: u32,
    ) -> Option<KDNearest> {
        let mut result = None;

        let diff = needle.asub(self.mid_point);
        let distance = diff.dot(diff).sqrt();

        if distance < limit && self.index != ignore_index {
            limit = distance;
            result = Some(KDNearest {
                index: self.index,
                distance: distance,
            })
        }

        let dot = diff.dot(self.normal);
        if dot <= 0.0 {
            if let Some(ref left) = self.left {
                if let Some(nearest) = left._find_nearest(needle, limit, ignore_index) {
                    limit = nearest.distance;
                    result = Some(nearest);
                }
            }

            if -dot < limit {
                if let Some(ref right) = self.right {
                    if let Some(nearest) = right._find_nearest(needle, limit, ignore_index) {
                        result = Some(nearest);
                    }
                }
            }
        } else {
            if let Some(ref right) = self.right {
                if let Some(nearest) = right._find_nearest(needle, limit, ignore_index) {
                    limit = nearest.distance;
                    result = Some(nearest);
                }
            }

            if dot < limit {
                if let Some(ref left) = self.left {
                    if let Some(nearest) = left._find_nearest(needle, limit, ignore_index) {
                        result = Some(nearest);
                    }
                }
            }
        }

        result
    }

    pub fn find_nearest(&self, color: [f32; 4]) -> u32 {
        self._find_nearest(color, f32::MAX, u32::MAX)
            .map(|x| x.index)
            .unwrap_or(0)
    }
}

fn occludes(origin: [f32; 4], occluder: [f32; 4], target: [f32; 4]) -> bool {
    let dir = occluder.asub(origin);
    dir.dot(dir) * 0.5 <= (target.asub(origin)).dot(dir)
}