use std::collections::BTreeMap;
use std::ops::Add;

// Граф - дерево деревьев
type Graph<V, E> = BTreeMap<V, BTreeMap<V, E>>;

pub fn dijkstra<V: Ord + Copy, E: Ord + Copy + Add<Output = E>>(
    graph: &Graph<V, E>,
    start: V,
) -> BTreeMap<V, Option<(V, E)>> {
    let mut res = BTreeMap::new();
    let mut queue = BTreeMap::new();

    // Путь от начала до начала равен 0
    res.insert(start, None);

    for (new, weight) in &graph[&start] {
        // Сохраняем путь
        res.insert(*new, Some((start, *weight)));
        // Сохраняем в очередь
        queue.insert(*new, *weight);
    }

    // Проходим по всем элементам очереди
    while let Some((vertex, path_weight)) = queue.pop_first() {
        // Проходим по всем следующим связанным элементам
        for (next, weight) in &graph[&vertex] {
            // Считаем итоговую сумму на данный момент
            let new_weight = path_weight + *weight;
            // Ищем следующий элемент в матрице
            match res.get(next) {
                // Если следующий элемент есть в матрице и имеет расстояние, меньшее найденного на этой итерации - не делаем ничего
                Some(Some((_, weight_next))) if new_weight >= *weight_next => {}
                // Если следующий элемент начальный - не делаем ничего
                Some(None) => {}
                // Если следующего элемента нет - сохраняем в матрицу и очередь с новым весом
                _ => {
                    res.insert(*next, Some((vertex, new_weight)));
                    queue.insert(*next, new_weight);
                }
            }
        }
    }

    res
}

fn insert_edge<V: Ord + Copy, E: Ord + Copy>(graph: &mut Graph<V, E>, v1: V, v2: V, c: E) {
    // Добавляем А->Б
    graph.entry(v1).or_default().insert(v2, c);
    // Добавляем Б-А
    graph.entry(v2).or_default().insert(v1, c);
}

fn main() {
    let args: Vec<String> = std::env::args().skip(1).collect();

    // Если аргументы пусты или -h/--help
    if args.is_empty() || args[0] == "--help" || args[0] == "-h" {
        println!("Usage: {} <start> <edges>.", std::env::args().nth(0).unwrap_or("dijkstra".to_string()));
        println!("Each vertex is specified as a single latin character");
        println!("Edges are separated by a whitespace and formatted as A-B=C where A, B are vertex names and C is the length of the edge");
        return;
    }

    // Стартовая вершина - первый аргумент
    let start = args[0].chars().nth(0).expect("No start specified");

    // Парсим аргументы в список рёбер
    let edges: Vec<(char, char, u128)> = args
        .iter()
        .skip(1) // Пропускаем первый аргумент - стартовую вершину
        .map(|s| {
            // Разбиваем аргумент по символам - и =
            let parts: Vec<&str> = s.split(&['-', '='][..]).collect();
            (
                parts[0].chars().next().unwrap(), // Первая вершина
                parts[1].chars().next().unwrap(), // Вторая вершина
                parts[2].parse().unwrap(),        // Длина
            )
        })
        .collect();

    // Создаем граф
    let mut graph = Graph::new();

    // Для каждого найденного ребра из вершины добавляем в граф
    for edge in &edges {
        insert_edge(&mut graph, edge.0.clone(), edge.1.clone(), edge.2);
    }

    println!(
        "Total edges: {}\nStarting vertex: {}\nFinding paths...",
        edges.len(),
        start
    );

    // Считаем длины путей
    let paths = dijkstra(&graph, start);

    // Выводим длину для каждого пути
    for path in &paths {
        let to = path.0;
        if let Some(second) = path.1 {
            println!("Length of path from {} to {} is {}", start, to, second.1);
        } else {
            println!("Length of path from {} to {} is 0", start, to);
        }
    }
}
