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

type Coordinate = (usize, usize);

pub enum AllowedMovements {
    DownRight,
    DownRightUp,
    DownRightUpLeft,
}

impl AllowedMovements {
    fn get_neighbours(&self, c: &Coordinate) -> Vec<Coordinate> {
        match self {
            Self::DownRight => self.get_neighbours_down_right(c),
            Self::DownRightUp => self.get_neighbours_down_right_up(c),
            Self::DownRightUpLeft => self.get_neighbours_down_right_up_left(c),
        }
    }

    fn get_neighbours_down_right(&self, (x, y): &Coordinate) -> Vec<Coordinate> {
        vec![(*x + 1, *y), (*x, *y + 1)]
    }

    fn get_neighbours_down_right_up(&self, (x, y): &Coordinate) -> Vec<Coordinate> {
        let mut result = self.get_neighbours_down_right(&(*x, *y));
        if *y > 0 {
            result.push((*x, *y - 1));
        }
        result
    }

    fn get_neighbours_down_right_up_left(&self, (x, y): &Coordinate) -> Vec<Coordinate> {
        let mut result = self.get_neighbours_down_right_up(&(*x, *y));
        if *x > 0 {
            result.push((*x - 1, *y));
        }
        result
    }
}

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: &'a mut dyn Matrix<T>,
    unvisited_with_distance: HashSet<Coordinate>,
    unvisited_without_distance: HashSet<Coordinate>,
    allowed_movements: AllowedMovements,
    start: 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_without_distance)
            .finish()
    }
}

impl<'a, T> Dijkstra<'a, T>
where
    T: PartialOrd + Add<Output = T> + Display + Copy + Debug,
{
    pub fn new(
        matrix: &'a mut dyn Matrix<T>,
        allowed_movements: AllowedMovements,
        start: Coordinate,
    ) -> Self {
        let unvisited_with_distance = HashSet::new();
        let unvisited_without_distance = Self::initialize_unvisited_set(&*matrix, start);
        Self {
            matrix,
            unvisited_with_distance,
            unvisited_without_distance,
            allowed_movements,
            start,
        }
    }

    pub fn solve(&mut self) {
        while !(self.unvisited_without_distance.is_empty()
            && self.unvisited_with_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
            .unvisited_with_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(self.start);
        self.unvisited_with_distance.remove(&coord);
        coord
    }

    fn get_neighbours(&self, coordinate: &Coordinate) -> Vec<Coordinate> {
        let possible_neighbours = self.allowed_movements.get_neighbours(coordinate);
        self.remove_out_of_bounds_neighbours(possible_neighbours)
    }

    fn remove_out_of_bounds_neighbours(&self, neighbours: Vec<Coordinate>) -> Vec<Coordinate> {
        neighbours
            .into_iter()
            .filter(|c| self.is_in_bounds(c))
            .collect()
    }

    fn is_in_bounds(&self, (x, y): &Coordinate) -> bool {
        *x < self.matrix.width() && *y < self.matrix.height()
    }

    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_hashsets(next);
            }
        } else {
            Self::update_node(next_node, new_distance, previous);
            self.update_distances_hashsets(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_hashsets(&mut self, next: Coordinate) {
        self.unvisited_without_distance.remove(&next);
        self.unvisited_with_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);
}