#![feature(iterator_try_collect)]

use clap::Parser;
use priority_queue::PriorityQueue;
use std::cmp::Ordering;
use std::collections::HashMap;
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<Vec<PathBuf>, 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(Vec::new());
    }

    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 = HashMap::with_capacity(num_ims);

    // 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, Vec::new());
    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_key(&a) {
                break (a, b);
            }
            else if !visitedImages.contains_key(&b) {
                break (b, a);
            }
        };
        visitedImages.insert(new, Vec::new());

        // insert all the new edges we could take
        visitedImages.entry(old).and_modify(|v|v.push(new));
        for i in 0..num_ims {
            // don't consider edges taking us to nodes we already visited
            if visitedImages.contains_key(&i) {
                continue;
            }

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

    // find TSP approximation via DFS through the MST
    fn dfs(cur: usize, t: &HashMap<usize, Vec<usize>>, into: &mut Vec<usize>) {
        into.push(cur);
        t.get(&cur).unwrap().iter().for_each(|c| dfs(*c, t, into));
    }
    let mut tsp_path = Vec::with_capacity(num_ims);
    dfs(0, &visitedImages, &mut tsp_path);

    Ok(tsp_path.iter().map(|i| args.images[*i].clone()).collect())
}

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

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

    for p in tsp_path {
        println!("{:?}", p);
    }

    Ok(())
}