use std::cmp::Ordering;
use std::collections::{BinaryHeap, VecDeque};

use crate::maps::VertexMap;
use crate::traits::{GraphTopology, IncidenceCursor};

#[derive(PartialEq, Eq)]
struct DistanceOrderedVertex<V> {
    distance: u32,
    vertex: V,
}

impl<V: Eq> PartialOrd for DistanceOrderedVertex<V> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl<V: Eq> Ord for DistanceOrderedVertex<V> {
    fn cmp(&self, other: &Self) -> Ordering {
        other.distance.cmp(&self.distance)
    }
}

pub struct DijkstraResult<V: Copy> {
    pub distances: VertexMap<V, Option<u32>>,
    pub predecessors: VertexMap<V, Option<V>>,
}

// TODO: Generalize the return type of the weight function.
pub fn dijkstra<G, W>(graph: &G, source: G::Vertex, weight: W) -> DijkstraResult<G::Vertex>
where
    G: GraphTopology,
    W: Fn(G::Edge) -> u32,
{
    let mut predecessors = graph.vertex_map(None);
    let distances = dijkstra_impl(graph, source, weight, |adjacent, predecessor| {
        predecessors[adjacent] = Some(predecessor);
    });
    DijkstraResult {
        distances,
        predecessors,
    }
}

pub fn dijkstra_unweighted<G>(graph: &G, source: G::Vertex) -> DijkstraResult<G::Vertex>
where
    G: GraphTopology,
{
    dijkstra(graph, source, |_| 1)
}

pub fn dijkstra_distances<G, W>(
    graph: &G,
    source: G::Vertex,
    weight: W,
) -> VertexMap<G::Vertex, Option<u32>>
where
    G: GraphTopology,
    W: Fn(G::Edge) -> u32,
{
    dijkstra_impl(graph, source, weight, |_, _| {})
}

pub fn dijkstra_distances_unweighted<G>(
    graph: &G,
    source: G::Vertex,
) -> VertexMap<G::Vertex, Option<u32>>
where
    G: GraphTopology,
{
    dijkstra_distances(graph, source, |_| 1)
}

fn dijkstra_impl<G, W, F>(
    graph: &G,
    source: G::Vertex,
    weight: W,
    mut on_relax: F,
) -> VertexMap<G::Vertex, Option<u32>>
where
    G: GraphTopology,
    W: Fn(G::Edge) -> u32,
    F: FnMut(G::Vertex, G::Vertex),
{
    let mut distances = graph.vertex_map(None);
    let mut heap = BinaryHeap::new();

    distances[source] = Some(0);
    heap.push(DistanceOrderedVertex {
        vertex: source,
        distance: 0,
    });

    while let Some(v) = heap.pop() {
        for incidence in graph.incidences(v.vertex) {
            let new_distance = distances[v.vertex].unwrap() + weight(incidence.1);
            if match distances[incidence.0] {
                None => true,
                Some(old_distance) if old_distance > new_distance => true,
                _ => false,
            } {
                distances[incidence.0] = Some(new_distance);
                on_relax(incidence.0, v.vertex);
                heap.push(DistanceOrderedVertex {
                    vertex: incidence.0,
                    distance: new_distance,
                });
            }
        }
    }
    distances
}

pub struct BfsResult<V: Copy> {
    pub distances: VertexMap<V, Option<u32>>,
    pub predecessors: VertexMap<V, Option<V>>,
}

pub fn bfs<G>(graph: &G, source: G::Vertex) -> BfsResult<G::Vertex>
where
    G: GraphTopology,
{
    let mut predecessors = graph.vertex_map(None);
    let result = bfs_impl(graph, source, |neighbor, predecessor| {
        predecessors[neighbor] = Some(predecessor);
        true
    });
    BfsResult {
        distances: result.distances,
        predecessors,
    }
}

pub fn bfs_distances<G>(graph: &G, source: G::Vertex) -> VertexMap<G::Vertex, Option<u32>>
where
    G: GraphTopology,
{
    bfs_impl(graph, source, |_, _| true).distances
}

pub fn bfs_find<G>(graph: &G, source: G::Vertex, target: G::Vertex) -> Option<u32>
where
    G: GraphTopology,
{
    bfs_find_where(graph, source, |v| v == target).map(|(_, distance)| distance)
}

pub fn bfs_find_where<G, P>(graph: &G, source: G::Vertex, predicate: P) -> Option<(G::Vertex, u32)>
where
    G: GraphTopology,
    P: Fn(G::Vertex) -> bool,
{
    if predicate(source) {
        return Some((source, 0));
    }
    bfs_impl(graph, source, |neighbor, _| !predicate(neighbor)).found
}

struct BfsImplResult<V: Copy> {
    distances: VertexMap<V, Option<u32>>,
    found: Option<(V, u32)>,
}

fn bfs_impl<G, F>(graph: &G, source: G::Vertex, mut on_discover: F) -> BfsImplResult<G::Vertex>
where
    G: GraphTopology,
    F: FnMut(G::Vertex, G::Vertex) -> bool,
{
    let mut distances = graph.vertex_map(None);
    let mut queue = VecDeque::new();

    distances[source] = Some(0);
    queue.push_back(source);

    while let Some(v) = queue.pop_front() {
        for neighbor in graph.adjacent_vertices(v) {
            if distances[neighbor].is_none() {
                let distance = distances[v].unwrap() + 1;
                distances[neighbor] = Some(distance);
                if !on_discover(neighbor, v) {
                    return BfsImplResult {
                        distances,
                        found: Some((neighbor, distance)),
                    };
                }
                queue.push_back(neighbor);
            }
        }
    }
    BfsImplResult {
        distances,
        found: None,
    }
}

pub struct DfsResult<V: Copy> {
    pub visited: VertexMap<V, bool>,
    pub predecessors: VertexMap<V, Option<V>>,
}

pub fn dfs<G>(graph: &G, source: G::Vertex) -> DfsResult<G::Vertex>
where
    G: GraphTopology,
{
    let mut predecessors = graph.vertex_map(None);
    let result = dfs_impl(graph, source, |neighbor, predecessor| {
        predecessors[neighbor] = Some(predecessor);
        true
    });
    DfsResult {
        visited: result.visited,
        predecessors,
    }
}

pub fn dfs_visited<G>(graph: &G, source: G::Vertex) -> VertexMap<G::Vertex, bool>
where
    G: GraphTopology,
{
    dfs_impl(graph, source, |_, _| true).visited
}

pub fn dfs_find<G>(graph: &G, source: G::Vertex, target: G::Vertex) -> bool
where
    G: GraphTopology,
{
    dfs_find_where(graph, source, |v| v == target).is_some()
}

pub fn dfs_find_where<G, P>(graph: &G, source: G::Vertex, predicate: P) -> Option<G::Vertex>
where
    G: GraphTopology,
    P: Fn(G::Vertex) -> bool,
{
    if predicate(source) {
        return Some(source);
    }
    dfs_impl(graph, source, |neighbor, _| !predicate(neighbor)).found
}

struct DfsImplResult<V: Copy> {
    visited: VertexMap<V, bool>,
    found: Option<V>,
}

fn dfs_impl<G, F>(graph: &G, source: G::Vertex, mut on_discover: F) -> DfsImplResult<G::Vertex>
where
    G: GraphTopology,
    F: FnMut(G::Vertex, G::Vertex) -> bool,
{
    let mut visited = graph.vertex_map(false);
    visited[source] = true;
    let mut stack = vec![(source, None::<G::Vertex>)];

    while let Some((v, predecessor)) = stack.pop() {
        if let Some(p) = predecessor {
            if !on_discover(v, p) {
                return DfsImplResult {
                    visited,
                    found: Some(v),
                };
            }
        }
        for neighbor in graph.adjacent_vertices(v) {
            if !visited[neighbor] {
                visited[neighbor] = true;
                stack.push((neighbor, Some(v)));
            }
        }
    }
    DfsImplResult {
        visited,
        found: None,
    }
}

pub fn dfs_find_path<G>(graph: &G, source: G::Vertex, target: G::Vertex) -> Option<Vec<G::Edge>>
where
    G: GraphTopology,
{
    dfs_find_path_where(graph, source, |v| v == target)
}

pub fn dfs_find_path_where<G, P>(graph: &G, source: G::Vertex, predicate: P) -> Option<Vec<G::Edge>>
where
    G: GraphTopology,
    P: Fn(G::Vertex) -> bool,
{
    if predicate(source) {
        return Some(vec![]);
    }

    let mut visited = graph.vertex_map(false);
    visited[source] = true;

    struct Frame<G: GraphTopology> {
        arrival_edge: Option<G::Edge>,
        cursor: G::IncidenceCursor,
    }

    let mut stack: Vec<Frame<G>> = vec![Frame {
        arrival_edge: None,
        cursor: graph.incidence_cursor(source),
    }];

    while let Some(frame) = stack.last_mut() {
        match frame.cursor.next(graph) {
            None => {
                stack.pop();
            }
            Some((neighbor, edge)) => {
                if predicate(neighbor) {
                    let mut path: Vec<G::Edge> =
                        stack.iter().filter_map(|f| f.arrival_edge).collect();
                    path.push(edge);
                    return Some(path);
                }
                if !visited[neighbor] {
                    visited[neighbor] = true;
                    stack.push(Frame {
                        arrival_edge: Some(edge),
                        cursor: graph.incidence_cursor(neighbor),
                    });
                }
            }
        }
    }
    None
}