use std::fmt::Debug;
use std::fmt::Display;
use std::{collections::HashSet, ops::Add};

type Coordinate = (usize, usize);

const START: Coordinate = (0, 0);

pub trait Matrix<T>
where
    T: Add + PartialOrd + Debug + Copy,
{
    fn height(&self) -> usize;
    fn width(&self) -> usize;
    fn get(&self, coordinate: Coordinate) -> Option<&dyn Node<T>>;
    fn get_mut(&mut self, coordinate: Coordinate) -> Option<&mut dyn Node<T>>;
}

impl<T> Debug for dyn Matrix<T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        todo!()
    }
}

pub struct Dijkstra<'a, T> {
    matrix: Box<&'a mut dyn Matrix<T>>,
    has_distance: HashSet<Coordinate>,
    unvisited: HashSet<Coordinate>,
}

impl<'a, T> Debug for Dijkstra<'a, T> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Dijkstra")
            .field("unvisited", &self.unvisited)
            .finish()
    }
}

impl<'a, T> Dijkstra<'a, T>
where
    T: PartialOrd + Add<Output = T> + Display + Copy + Debug,
{
    pub fn new(matrix: Box<&'a mut dyn Matrix<T>>) -> Self {
        let has_distance = HashSet::from([START]);
        let unvisited = Self::initialize_unvisited_set(*matrix, START);
        Self {
            matrix,
            has_distance,
            unvisited,
        }
    }

    pub fn solve(&mut self) {
        while !(self.unvisited.is_empty() && self.has_distance.is_empty()) {
            let coord = self.get_unvisited_coord_with_smallest_distance_from_start();
            let neighbours = self.get_neighbours(&coord);
            let previous_distance = *self.matrix.get(coord).unwrap().get_distance().unwrap();
            for neighbour in neighbours {
                self.update_neighbour(neighbour, coord, previous_distance);
            }
        }
    }

    fn get_unvisited_coord_with_smallest_distance_from_start(&mut self) -> Coordinate {
        let coord = self
            .has_distance
            .iter()
            .filter_map(|c| {
                let node = self.matrix.get(*c).unwrap();
                node.get_distance().map(|value| (c, value))
            })
            .reduce(|(c0, v0), (c1, v1)| if v0 < v1 { (c0, v0) } else { (c1, v1) })
            .map(|(c, _)| *c)
            .unwrap_or((0, 0));
        self.has_distance.remove(&coord);
        coord
    }

    fn get_neighbours(&self, coordinate: &Coordinate) -> Vec<Coordinate> {
        let (x0, y0) = coordinate;
        let mut result = vec![];
        if *x0 < self.matrix.width() - 1 {
            result.push((*x0 + 1, *y0));
        }
        if *y0 < self.matrix.height() - 1 {
            result.push((*x0, *y0 + 1));
        }
        result
    }

    fn update_neighbour(&mut self, next: Coordinate, previous: Coordinate, previous_distance: T) {
        let next_node = self.matrix.get_mut(next).unwrap();
        let new_distance = previous_distance + *next_node.get_value();
        if let Some(old_distance) = next_node.get_distance() {
            if new_distance < *old_distance {
                Self::update_node(next_node, new_distance, previous);
                self.update_distances_hashset(next);
            }
        } else {
            Self::update_node(next_node, new_distance, previous);
            self.update_distances_hashset(next);
        }
    }

    fn initialize_unvisited_set(matrix: &dyn Matrix<T>, start: Coordinate) -> HashSet<Coordinate> {
        let unvisited: HashSet<_> = (0..matrix.height())
            .flat_map(|x| (0..matrix.width()).zip([x].into_iter().cycle()))
            .filter(|value| *value != start)
            .collect();
        unvisited
    }

    fn update_node(next_node: &mut dyn Node<T>, new_distance: T, previous: (usize, usize)) {
        next_node.set_distance(new_distance);
        next_node.set_previous(previous);
    }

    fn update_distances_hashset(&mut self, next: Coordinate) {
        self.unvisited.remove(&next);
        self.has_distance.insert(next);
    }
}

pub trait Node<T>
where
    T: Add + PartialOrd,
{
    fn get_value(&self) -> &T;
    fn get_distance(&self) -> Option<&T>;
    fn get_previous(&self) -> Option<Coordinate>;

    fn set_distance(&mut self, distance: T);
    fn set_previous(&mut self, coord: Coordinate);
}