aboutsummaryrefslogtreecommitdiff
path: root/src/main.rs
blob: 5cfc10cbd53c157d034b5cdf45cc77afd6ed02e2 (plain)
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
#![feature(iterator_try_collect)]

use clap::Parser;
use priority_queue::PriorityQueue;
use std::cmp::Ordering;
use std::collections::HashSet;
use std::path::PathBuf;

use embedders::*;

mod embedders;

#[derive(Debug, Clone, Copy, clap::ValueEnum)]
enum Embedder {
    Brightness,
    Hue,
}

#[derive(Debug, Parser)]
struct Args {
    #[arg(short, long, default_value = "hue")]
    embedder: Embedder,

    images: Vec<PathBuf>,
}

//#[derive(Debug)]
//struct Config {
//    base_dir: xdg::BaseDirectories,
//}
//
//fn get_config() -> Result<Config, String> {
//    let dirs = xdg::BaseDirectories::with_prefix("embeddings-sort")
//        .map_err(|_| "oh no")?;
//
//    Ok(Config{base_dir: dirs})
//}

fn process_embedder<E>(mut e: E, args: Args) -> Result<(), String>
    where E: EmbedderT
{
    // wrapper struct to
    // - reverse the ordering
    // - implement Ord, even though the type is backed by an f64
    #[repr(transparent)]
    #[derive(Debug, PartialEq)]
    struct DownOrd (f64);
    impl Eq for DownOrd {}
    impl PartialOrd for DownOrd {
        fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
            self.0.partial_cmp(&other.0).map(|e| e.reverse())
        }
    }
    impl Ord for DownOrd {
        fn cmp(&self, other: &Self) -> Ordering {self.partial_cmp(other).unwrap()}
    }

    let num_ims = args.images.len();
    if num_ims == 0 {
        return Ok(());
    }

    let embeds: Vec<_> = e
        .embed(&args.images)?
        .into_iter()
        .collect();

    let mut possibleEdges =
        PriorityQueue::with_capacity((num_ims * num_ims - num_ims) / 2);
    let mut visitedImages = HashSet::with_capacity(num_ims);
    let mut takenEdges = HashSet::with_capacity(num_ims - 1); // TODO das muss anders

    // here, we start at 0.
    // we might get a better result in the end if we started with a vertex next
    // to the lowest-cost edge, but we don't know which one that is (though we
    // could compute that without changing our asymptotic complexity)
    visitedImages.insert(0);
    for i in 1..num_ims {
        possibleEdges.push((0, i), DownOrd(embeds[0].dist(&embeds[i])));
    }

    // prims algorithm or something like that
    while visitedImages.len() < num_ims {
        // find the edge with the least cost that connects us to a new vertex
        let (new, old) = loop {
            let ((a, b), _) = possibleEdges.pop().unwrap();
            if !visitedImages.contains(&a) {
                break (a, b);
            }
            else if !visitedImages.contains(&b) {
                break (b, a);
            }
        };
        visitedImages.insert(new);

        println!("next edge: {:?} -> {:?}", old, new);

        // insert all the new edges we could take
        takenEdges.insert((new, old));
        for i in 0..num_ims {
            // don't consider edges taking us to nodes we already visited
            if visitedImages.contains(&i) {
                continue;
            }

            possibleEdges.push((new, i), DownOrd(embeds[new].dist(&embeds[i])));
        }
    }

    Ok(())
}

fn main() -> Result<(), String> {
    //let cfg = get_config()?;
    let args = Args::parse();

    match args.embedder {
        Embedder::Brightness => process_embedder(BrightnessEmbedder, args),
        Embedder::Hue => process_embedder(HueEmbedder, args),
    }
}