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

use crate::union_find::UnionFind;

pub type Distance = usize;

pub trait Graph {
    type Node;
    type Edge: Edge;

    /// All nodes.
    fn nodes(&self) -> impl Iterator<Item = &Self::Node>;
    /// All edges.
    fn edges(&self) -> impl Iterator<Item = &Self::Edge>;

    /// Indices of neighbors reachable from node.
    fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>>;
    /// Indices of neighbors that can reach node.
    /// For directed graphs, returns same result as neighbor_idxs
    fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>>;
    /// All edges going to or from node.
    /// For directed graphs, mind the from and to distinction.
    fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>>;

    /// Look if there is an edge between nodes from and to.
    /// It is expected this function will be overriden with more performant version.
    fn has_edge(&self, from: usize, to: usize) -> bool {
        self.edges()
            .any(|edge| edge.from_node() == from && edge.to_node() == to)
    }

    /// Find an edge that connects nodes from and to, if it exists.
    /// If it doesn't, return none.
    /// It is expected this function will be overriden with more performant version.
    fn get_edge_between(&self, from: usize, to: usize) -> Option<usize> {
        self.edges_of_idxs(from)
            .map(|edges| edges
                 .filter(|&edge| self.edge(edge).unwrap().to_node() == to)
                 .next())
            .flatten()
    }

    /// Get a single edge at the given index.
    fn edge(&self, id: usize) -> Option<&Self::Edge> {
        self.edges().skip(id).next()
    }

    /// Get a single node at the given index.
    fn node(&self, id: usize) -> Option<&Self::Node> {
        self.nodes().skip(id).next()
    }
}

pub trait MutGraph: Graph {
    fn nodes_mut(&mut self) -> impl Iterator<Item = &mut Self::Node>;
    fn edges_mut(&mut self) -> impl Iterator<Item = &mut Self::Edge>;

    fn node_mut(&mut self, id: usize) -> Option<&mut Self::Node> {
        self.nodes_mut().skip(id).next()
    }

    fn edge_mut(&mut self, id: usize) -> Option<&mut Self::Edge> {
        self.edges_mut().skip(id).next()
    }
}

pub trait WeightedEdge {
    type Cost: PartialOrd;
    fn cost(&self) -> Self::Cost;
}

/// An edge.
pub trait Edge {
    /// Index of a node this edge goes from.
    /// For undirected graphs, from_node and to_node have the same meaning.
    fn from_node(&self) -> usize;
    /// Index of a node this edge goes to.
    /// For undirected graphs, from_node and to_node have the same meaning.
    fn to_node(&self) -> usize;
}

/// An edge that might be reversed easily.
pub trait ReversibleEdge: Edge {
    fn reverse(self) -> Self;
}

#[derive(Debug, Clone, PartialEq)]
pub struct GenericEdge {
    from: usize,
    to: usize,
}

impl GenericEdge {
    pub fn new(from: usize, to: usize) -> Self {
        Self {
            from,
            to
        }
    }
}

impl From<(usize, usize)> for GenericEdge {
    fn from(value: (usize, usize)) -> Self {
        Self {
            from: value.0,
            to: value.1
        }
    }
}

impl Edge for GenericEdge {
    fn from_node(&self) -> usize {
        self.from
    }

    fn to_node(&self) -> usize {
        self.to
    }

}

impl ReversibleEdge for GenericEdge {
    fn reverse(mut self) -> Self {
        (self.from, self.to) = (self.to, self.from);
        self
    }
}


/// A directed graph that owns nodes and edges of
/// specific types given by generics.
#[derive(Debug, Clone, PartialEq)]
pub struct GenericDirectedGraph<T, TEdge: Edge>
{
    nodes: Vec<T>,
    edges: Vec<TEdge>,
    node_edges: Vec<Vec<usize>>,
    reverse_node_edges: Vec<Vec<usize>>,
}

impl<T, TEdge: Edge> GenericDirectedGraph<T, TEdge> {
    pub fn new(nodes: Vec<T>) -> Self {
        let nodes_len = nodes.len();
        Self {
            nodes,
            edges: vec![],
            node_edges: vec![vec![]; nodes_len],
            reverse_node_edges: vec![vec![]; nodes_len],
        }
    }

    pub fn add_generic_edge(&mut self, edge: TEdge) -> usize {
        let idx = self.edges.len();

        self.node_edges[edge.from_node()].push(idx);
        self.reverse_node_edges[edge.to_node()].push(idx);
        self.edges.push(edge);

        idx
    }

    pub fn decompose(self) -> Vec<T> {
        self.nodes
    }
}

impl<T, TEdge: ReversibleEdge> GenericDirectedGraph<T, TEdge> {
    pub fn reverse(mut self) -> Self {
        (self.node_edges, self.reverse_node_edges) =
            (self.reverse_node_edges, self.node_edges);

        self.edges = self.edges
            .into_iter()
            .map(|edge| edge.reverse())
            .collect();

        self
    }
}

impl<T, TEdge: Edge + From<(usize, usize)>> GenericDirectedGraph<T, TEdge> {
    pub fn add_edge(&mut self, from: usize, to: usize) -> usize {
        let edge = (from, to).into();
        self.add_generic_edge(edge)
    }
}

impl<T, TEdge: Edge> Graph for GenericDirectedGraph<T, TEdge> {
    type Node = T;
    type Edge = TEdge;

    fn nodes(&self) -> impl Iterator<Item = &T> {
        self.nodes.iter()
    }

    fn edges(&self) -> impl Iterator<Item = &TEdge> {
        self.edges.iter()
    }

    fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.node_edges.get(node)
            .map(|edges| edges.iter()
                 .map(|i| self.edges[*i].to_node()))
    }

    fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.reverse_node_edges.get(node)
            .map(|edges| edges.iter()
                 .map(|i| self.edges[*i].from_node()))
    }

    fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        let reverse_edges = self.reverse_node_edges.get(node)?;

        self.node_edges.get(node)
            .map(|edges|
                 edges.iter()
                 .chain(reverse_edges.iter())
                 .map(|i| *i))
    }
}

impl<T, TEdge: Edge> MutGraph for GenericDirectedGraph<T, TEdge> {
    fn nodes_mut(&mut self) -> impl Iterator<Item = &mut Self::Node> {
        self.nodes.iter_mut()
    }

    fn edges_mut(&mut self) -> impl Iterator<Item = &mut Self::Edge> {
        self.edges.iter_mut()
    }
}

/// An undirected graph that owns nodes and edges of
/// specific types given by generics.
pub struct GenericGraph<T, TEdge: Edge> {
    nodes: Vec<T>,
    edges: Vec<TEdge>,
    node_neighbors: Vec<Vec<usize>>,
    node_edges: Vec<Vec<usize>>,
    adjacency_matrix: Option<Vec<Vec<Option<usize>>>>
}

impl<T, TEdge: Edge> GenericGraph<T, TEdge>
{
    pub fn new(nodes: Vec<T>, adjacency_matrix: bool) -> Self {
        let nodes_count = nodes.len();
        Self {
            nodes,
            edges: vec![],
            node_neighbors: vec![vec![]; nodes_count],
            node_edges: vec![vec![]; nodes_count],
            adjacency_matrix: if adjacency_matrix {
                Some(vec![vec![None; nodes_count]; nodes_count])
            } else {
                None
            }
        }
    }

    pub fn add_generic_edge(&mut self, edge: TEdge) -> usize {
        let idx = self.edges.len();

        if let Some(adjacency_matrix) = self.adjacency_matrix.as_deref_mut() {
            adjacency_matrix[edge.from_node()][edge.to_node()] = Some(idx);
            adjacency_matrix[edge.to_node()][edge.from_node()] = Some(idx);
        }

        self.node_edges[edge.from_node()].push(idx);
        self.node_edges[edge.to_node()].push(idx);
        self.node_neighbors[edge.from_node()].push(edge.to_node());
        self.node_neighbors[edge.to_node()].push(edge.from_node());
        self.edges.push(edge);

        idx
    }

    pub fn filter_edges(self, filter: impl Fn(usize, &TEdge) -> bool) -> Self {
        let keep_edges = self.edges
            .into_iter()
            .enumerate()
            .filter(|(i, e)| filter(*i, e))
            .map(|(_, e)| e);

        let mut new_graph = Self::new(self.nodes, self.adjacency_matrix.is_some());

        for edge in keep_edges {
            new_graph.add_generic_edge(edge);
        }

        new_graph
    }

    // NOTE: it's expected the edges will not reconnect, only the type will change.
    // from_node() and to_node() should stay the same!
    pub fn map_edges<TNewEdge: Edge>(self, map: impl Fn(TEdge) -> TNewEdge) -> GenericGraph<T, TNewEdge> {
        GenericGraph::<T, TNewEdge> {
            nodes: self.nodes,
            edges: self.edges.into_iter().map(|edge| map(edge)).collect(),
            node_neighbors: self.node_neighbors,
            node_edges: self.node_edges,
            adjacency_matrix: self.adjacency_matrix
        }
    }

    pub fn map_nodes<TNewNode>(self, map: impl Fn(T) -> TNewNode) -> GenericGraph<TNewNode, TEdge> {
        GenericGraph::<TNewNode, TEdge> {
            nodes: self.nodes.into_iter().map(|node| map(node)).collect(),
            edges: self.edges,
            node_neighbors: self.node_neighbors,
            node_edges: self.node_edges,
            adjacency_matrix: self.adjacency_matrix
        }
    }
}

impl<T, TEdge: Edge> Graph for GenericGraph<T, TEdge> {
    type Node = T;
    type Edge = TEdge;

    fn nodes(&self) -> impl Iterator<Item = &T> {
        self.nodes.iter()
    }

    fn edges(&self) -> impl Iterator<Item = &TEdge> {
        self.edges.iter()
    }

    fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.node_neighbors.get(node).map(|neighbors| neighbors.iter().map(|&node| node))
    }

    fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.node_neighbors.get(node).map(|neighbors| neighbors.iter().map(|&node| node))
    }

    fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.node_edges.get(node).map(|edges| edges.iter().map(|&edge| edge))
    }

    fn has_edge(&self, from: usize, to: usize) -> bool {
        if let Some(adjacency_matrix) = &self.adjacency_matrix {
            adjacency_matrix[from][to].is_some()
        } else {
            self.edges()
                .any(|edge| edge.from_node() == from && edge.to_node() == to)
        }
    }

    fn get_edge_between(&self, from: usize, to: usize) -> Option<usize> {
        if let Some(adjacency_matrix) = &self.adjacency_matrix {
            adjacency_matrix[from][to]
        } else {
            self.edges_of_idxs(from)
                .map(|edges| edges
                     .filter(|&edge| self.edge(edge).unwrap().to_node() == to)
                     .next())
                .flatten()
        }
    }
}

impl<T, TEdge: Edge> MutGraph for GenericGraph<T, TEdge> {
    fn nodes_mut(&mut self) -> impl Iterator<Item = &mut Self::Node> {
        self.nodes.iter_mut()
    }

    fn edges_mut(&mut self) -> impl Iterator<Item = &mut Self::Edge> {
        self.edges.iter_mut()
    }
}

pub struct ReversedEdge<'a, T: Edge> {
    edge: &'a T
}

impl<'a, T: Edge> ReversedEdge<'a, T> {
    pub fn new(edge: &'a T) -> Self {
        Self {
            edge
        }
    }
}

impl<'a, T: Edge> Edge for ReversedEdge<'a, T> {
    fn from_node(&self) -> usize {
        self.edge.to_node()
    }

    fn to_node(&self) -> usize {
        self.edge.from_node()
    }
}

/// A view on a graph that reverses all its
/// edges.
pub struct ReversedGraph<'a, T: Graph> {
    graph: &'a T,
    edges: Vec<ReversedEdge<'a, T::Edge>>
}

impl<'a, T: Graph> ReversedGraph<'a, T> {
    pub fn new(graph: &'a T) -> Self {
        Self {
            graph,
            edges: graph.edges()
                .map(|edge| ReversedEdge::new(edge))
                .collect()
        }
    }
}

impl<'a, T: Graph> Graph for ReversedGraph<'a, T> {
    type Node = T::Node;
    type Edge = ReversedEdge<'a, T::Edge>;

    fn nodes(&self) -> impl Iterator<Item = &Self::Node> {
        self.graph.nodes()
    }

    fn edges(&self) -> impl Iterator<Item = &Self::Edge> {
        self.edges.iter()
    }

    fn neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.graph.reverse_neighbor_idxs(node)
    }

    fn reverse_neighbor_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.graph.neighbor_idxs(node)
    }

    fn edges_of_idxs(&self, node: usize) -> Option<impl Iterator<Item = usize>> {
        self.graph.edges_of_idxs(node)
    }
}

/// Make a search out of the starting node, visiting every
/// node that can be visited from it. Denoting the
/// distances between the nodes and the parents that visited the
/// given node for reconstructing the shortest path.
pub fn breadth_first_search<TGraph>(
    graph: &TGraph,
    starting_node: usize,
) -> (Vec<Option<Distance>>, Vec<usize>)
where
    TGraph: Graph
{
    let nodes_len = graph.nodes().count();
    let mut distances = vec![None; nodes_len];
    let mut visited = vec![false; nodes_len];
    let mut parents = vec![0; nodes_len];

    let mut node_queue = VecDeque::with_capacity(nodes_len / 4);

    node_queue.push_back(starting_node);
    distances[starting_node] = Some(0);
    visited[starting_node] = true;

    while let Some(current) = node_queue.pop_front() {
        for neighbor in graph.neighbor_idxs(current).unwrap() {
            if visited[neighbor] {
                continue
            }

            visited[neighbor] = true;
            distances[neighbor] = Some(distances[current].unwrap() + 1);
            parents[neighbor] = current;

            node_queue.push_back(neighbor);
        }
    }

    (distances, parents)
}

/// Like breadth first search, but only look if a node
/// is reachable, do not figure out the distance, do not
/// figure out the shortest path.
pub fn breadth_first_reachable<TGraph>(
    graph: &TGraph,
    starting_node: usize,
) -> Vec<bool>
where
    TGraph: Graph
{
    let nodes_len = graph.nodes().count();
    let mut visited = vec![false; nodes_len];

    let mut node_queue = VecDeque::with_capacity(nodes_len / 4);

    node_queue.push_back(starting_node);
    visited[starting_node] = true;

    while let Some(current) = node_queue.pop_front() {
        for neighbor in graph.neighbor_idxs(current).unwrap() {
            if visited[neighbor] {
                continue
            }

            visited[neighbor] = true;

            node_queue.push_back(neighbor);
        }
    }

    visited
}

/// Figure out distance from each node to each node.
/// In case the node is unreachable, Distance::MAX is
/// in the matrix.
pub fn floyd_warshall<TNode, TEdge, TGraph>(
    graph: &TGraph
) -> Vec<Vec<Distance>>
where
    TEdge: Edge,
    TGraph: Graph<Node = TNode, Edge = TEdge>
{
    let nodes = graph.nodes().count();
    let mut distances = vec![vec![Distance::MAX; nodes]; nodes];

    for edge in graph.edges() {
        distances[edge.from_node()][edge.to_node()] = 1;
    }

    for v in 0..nodes {
        distances[v][v] = 0;
    }

    for k in 0..nodes {
        for i in 0..nodes {
            for j in 00..nodes {
                if distances[i][k] == Distance::MAX || distances[k][j] == Distance::MAX {
                    continue;
                }

                if distances[i][j] > distances[i][k] + distances[k][j] {
                    distances[i][j] = distances[i][k] + distances[k][j];
                }
            }
        }
    }

    distances
}


/// A generic representation of a minimum spanning
/// tree for cases where the graph might not be
/// fully connected, and thus it is possible the minimum
/// spanning tree will not be a single component.
#[derive(Clone, Debug, PartialEq)]
pub struct MinimumSpanningTree {
    // What component does node i belong to?
    pub components: UnionFind,
    // Edge indices in the tree
    pub edges: Vec<usize>,
}

impl MinimumSpanningTree {
    pub fn nodes_count(&self) -> usize {
        self.components.len()
    }

    pub fn components_count(&self) -> usize {
        self.nodes_count() - self.edges.len()
    }
}

/// Use the kruskal algorithm for finding the minimum spanning tree.
/// Take only edges filtered by selector as candidates for the spanning tree.
/// Initial minimum spanning tree can be passed, that should usually be a result
/// of prior run of this function with a different selector.
pub fn minimal_spanning_tree_kruskal<'a, TNode, TWeight, TEdge, TGraph>(
    graph: &TGraph,
    initial: Option<MinimumSpanningTree>,
    selector: impl Fn(&TEdge) -> bool
) -> MinimumSpanningTree
where
    TWeight: PartialOrd,
    TEdge: Edge + WeightedEdge<Cost = TWeight>,
    TGraph: Graph<Node = TNode, Edge = TEdge>
{
    // let separate_new_edges = initial.is_some();
    let nodes = graph.nodes().count();
    let mut initial_edge_selected = vec![false; graph.edges().count()];
    let mut current = if let Some(initial) = initial {
        for edge in &initial.edges {
            initial_edge_selected[*edge] = true;
        }

        initial
    } else {
        MinimumSpanningTree {
            components: UnionFind::make_set(nodes),
            edges: Vec::with_capacity(nodes - 1),
        }
    };

    let mut remaining_edges = graph.edges()
        .enumerate()
        .filter(|(i, e)| !initial_edge_selected[*i] && selector(e))
        .collect::<Vec<_>>();

    remaining_edges.sort_by(
        |a, b| a.1.cost().partial_cmp(&b.1.cost()).unwrap_or(Ordering::Less)
    );

    for (i, edge) in remaining_edges {
        // 1. does the edge connect two components?
        let (root_a, root_b) = {
            if current.components_count() == 1 {
                break;
            }

            let root_a = current.components.find(edge.from_node());
            let root_b = current.components.find(edge.to_node());

            if root_a == root_b {
                continue;
            }

            (root_a, root_b)
        };
        // 2. if so, use it and connect them
        current.components.union(root_a, root_b);
        current.edges.push(i);
    }

    current
}