// Copyright (c) 2016, 2017, 2018 Frank Fischer <frank-fischer@shadow-soft.de>
//
// This program is free software: you can redistribute it and/or
// modify it under the terms of the GNU General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful, but
// WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
// General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see  <http://www.gnu.org/licenses/>
//

//! Implementation of bidirectional Dijkstra's algorithm for shortest paths.

use shortestpath::binheap::BinHeap;
use shortestpath::heap::Heap;
use EdgeMap;
use {Digraph, Edge, IndexDigraph, IndexGraph, IndexNetwork, Node};

use num::traits::NumAssign;
use std::cmp::Ordering;
use std::collections::HashMap;
use std::hash::Hash;
use std::ops::Index;

/// Bidirectional Dijkstra's shortest path algorithm.
pub struct BiDijkstra<'a, G, W, N, E, H = BinHeap<NodeItem<N, E, W>, usize>>
where
    G: 'a + IndexGraph<'a, Node = N, Edge = E>,
    N: 'a + Node + Hash,
    E: 'a + Edge,
    W: NumAssign + Ord + Copy,
    H: Heap<Key = NodeItem<G::Node, G::Edge, W>, Item = usize>,
{
    g: &'a G,
    srcitms: HashMap<G::Node, NodeState<W>>,
    srcheap: H,
    pred: Vec<(G::Node, (G::Node, G::Edge))>,
    snkitms: HashMap<G::Node, NodeState<W>>,
    snkheap: H,
    succ: Vec<(G::Node, (G::Node, G::Edge))>,
}

impl<'a, G, W, N, E, H> BiDijkstra<'a, G, W, N, E, H>
where
    G: 'a + IndexGraph<'a, Node = N, Edge = E>,
    N: 'a + Node + Hash,
    E: 'a + Edge,
    W: NumAssign + Ord + Copy,
    H: Heap<Key = NodeItem<G::Node, G::Edge, W>, Item = usize>,
{
    /// Create new data structure for running Dijkstra on a graph.
    pub fn new(g: &'a G) -> Self {
        BiDijkstra {
            g,
            srcitms: HashMap::new(),
            srcheap: H::new(),
            pred: vec![],
            snkitms: HashMap::new(),
            snkheap: H::new(),
            succ: vec![],
        }
    }

    /// Compute a shortest path with Dijkstra's algorithm on an undirected graph.
    ///
    /// Each edge can be traversed in both directions with the same weight.
    ///
    /// - `weights` the (non-negative) edge weights
    /// - `src` the source node
    /// - `snk` the sink node
    ///
    /// Return the incoming edge for each node forming the shortest path
    /// tree.
    pub fn undirected<Ws>(&mut self, weights: Ws, src: G::Node, snk: G::Node) -> Vec<G::Edge>
    where
        Ws: EdgeMap<'a, G, W>,
    {
        self.generic(weights, src, snk, |g, u| g.neighs(u), |g, u| g.neighs(u))
    }

    /// Solve shortest path with Dijkstra as bidirected graph.
    ///
    /// The graph is considered as bidirected graph with different weights
    /// for each direction.
    ///
    /// - `weights` the (non-negative) arc weights
    /// - `src` the source node
    /// - `snk` the sink node
    ///
    /// Return the incoming arc for each node forming the shortest path
    /// tree.
    pub fn bidirected<Ws>(&mut self, weights: Ws, src: G::Node, snk: G::Node) -> Vec<G::Edge>
    where
        G: IndexNetwork<'a>,
        Ws: EdgeMap<'a, G, W>,
    {
        self.generic(weights, src, snk, |g, u| g.neighs(u), |g, u| g.neighs(u))
    }

    /// Solve shortest path with Dijkstra as directed graph.
    ///
    /// The graph is considered directed, travel is only allowed along
    /// outgoing edges.
    ///
    /// - `weights` the (non-negative) arc weights
    /// - `src` the source node
    /// - `snk` the sink node
    ///
    /// Return the incoming arc for each node forming the shortest path
    /// tree.
    pub fn directed<Ws>(&mut self, weights: Ws, src: G::Node, snk: G::Node) -> Vec<G::Edge>
    where
        G: Digraph<'a>,
        Ws: EdgeMap<'a, G, W>,
    {
        self.generic(weights, src, snk, |g, u| g.outedges(u), |g, u| g.inedges(u))
    }

    pub fn generic<Weights, Out, OutIt, In, InIt>(
        &mut self,
        weights: Weights,
        src: G::Node,
        snk: G::Node,
        outedges: Out,
        inedges: In,
    ) -> Vec<G::Edge>
    where
        Weights: Index<G::Edge, Output = W>,
        Out: Fn(&'a G, G::Node) -> OutIt,
        OutIt: Iterator<Item = (G::Edge, G::Node)>,
        In: Fn(&'a G, G::Node) -> InIt,
        InIt: Iterator<Item = (G::Edge, G::Node)>,
    {
        for e in self.g.edges() {
            assert!(weights[e] >= W::zero(), "Weights must be non-negative");
        }

        // source data
        self.srcheap.clear();
        self.srcitms.clear();
        self.pred.clear();

        self.srcitms.insert(
            src,
            NodeState::Seen(self.srcheap.insert(NodeItem {
                node: src,
                weight: W::zero(),
                pred: None,
            })),
        );

        // sink data
        self.snkheap.clear();
        self.snkitms.clear();
        self.succ.clear();

        self.snkitms.insert(
            snk,
            NodeState::Seen(self.snkheap.insert(NodeItem {
                node: snk,
                weight: W::zero(),
                pred: None,
            })),
        );

        let mut midnode = None;
        while !self.srcheap.is_empty() && !self.snkheap.is_empty() {
            // source Dijkstra
            let udata = self.srcheap.pop_min().unwrap();
            let u = udata.node;
            let d = udata.weight;

            // add incoming arc to result
            if let Some(e) = udata.pred {
                self.pred.push((u, e));
            }

            if let Some(&NodeState::Finished(dsnk)) = self.snkitms.get(&u) {
                let mut u = u;
                let mut d = d + dsnk;
                for (v, vst) in &self.srcitms {
                    if let NodeState::Finished(dsrc) = *vst {
                        if let Some(&NodeState::Finished(dsnk)) = self.snkitms.get(v) {
                            if dsrc + dsnk < d {
                                d = dsrc + dsnk;
                                u = *v;
                            }
                        }
                    }
                }
                midnode = Some(u);
                break;
            }

            self.srcitms.insert(u, NodeState::Finished(d));

            for (e, v) in outedges(self.g, u) {
                match self.srcitms.get(&v) {
                    // unknown node
                    None => {
                        let newd = d + weights[e];
                        self.srcitms.insert(
                            v,
                            NodeState::Seen(self.srcheap.insert(NodeItem {
                                node: v,
                                weight: newd,
                                pred: Some((u, e)),
                            })),
                        );
                    }
                    // seen but uncompleted node
                    Some(&NodeState::Seen(itm)) => {
                        let newd = d + weights[e];
                        if newd < self.srcheap.key(itm).weight {
                            {
                                let data = self.srcheap.key(itm);
                                data.weight = newd;
                                data.pred = Some((u, e));
                            }
                            self.srcheap.decrease(itm);
                        }
                    }
                    // completed node
                    _ => (),
                }
            }

            // sink Dijkstra
            let vdata = self.snkheap.pop_min().unwrap();
            let v = vdata.node;
            let d = vdata.weight;

            // add outgoing arc to result
            if let Some(e) = vdata.pred {
                self.succ.push((v, e));
            }

            if let Some(&NodeState::Finished(dsrc)) = self.srcitms.get(&v) {
                let mut v = v;
                let mut d = d + dsrc;
                for (w, wst) in &self.snkitms {
                    if let NodeState::Finished(dsnk) = *wst {
                        if let Some(&NodeState::Finished(dsrc)) = self.srcitms.get(w) {
                            if dsrc + dsnk < d {
                                d = dsrc + dsnk;
                                v = *w;
                            }
                        }
                    }
                }
                midnode = Some(v);
                break;
            }

            self.snkitms.insert(v, NodeState::Finished(d));

            for (e, u) in inedges(self.g, v) {
                match self.snkitms.get(&u) {
                    // unknown node
                    None => {
                        let newd = d + weights[e];
                        self.snkitms.insert(
                            u,
                            NodeState::Seen(self.snkheap.insert(NodeItem {
                                node: u,
                                weight: newd,
                                pred: Some((v, e)),
                            })),
                        );
                    }
                    // seen but uncompleted node
                    Some(&NodeState::Seen(itm)) => {
                        let newd = d + weights[e];
                        if newd < self.snkheap.key(itm).weight {
                            {
                                let data = self.snkheap.key(itm);
                                data.weight = newd;
                                data.pred = Some((v, e));
                            }
                            self.snkheap.decrease(itm);
                        }
                    }
                    // completed node
                    _ => (),
                }
            }
        }

        if let Some(midnode) = midnode {
            let mut path = vec![];
            let mut y = midnode;
            for &(v, (u, e)) in self.pred.iter().rev() {
                if v == y {
                    path.push(e);
                    y = u;
                }
            }
            path.reverse();

            let mut x = midnode;
            for &(u, (v, e)) in self.succ.iter().rev() {
                if u == x {
                    path.push(e);
                    x = v;
                }
            }

            path
        } else {
            vec![]
        }
    }
}

/// Compute a shortest path with bidirectional Dijkstra algorithm on an undirected graph.
///
/// Each edge can be traversed in both directions with the same weight.
///
/// - `g` the graph
/// - `weights` the (non-negative) edge weights
/// - `src` the source node
/// - `snk` the sink node
///
/// Return the edges forming the path.
///
/// # Example
///
/// ```
/// use rs_graph::*;
/// use rs_graph::shortestpath::bidijkstra;
///
/// let mut g = LinkedListGraph::<u32>::default();
/// let mut weights: Vec<usize> = vec![];
///
/// let nodes: Vec<_> = (0..6).map(|_| g.add_node()).collect();
/// for &(u,v,w) in [(0,1,7), (0,2,9), (0,3,14),
///                  (1,2,10), (1,3,15),
///                  (2,3,11), (2,5,2),
///                  (3,4,6), (5,4,9)].iter() {
///     g.add_edge(nodes[u], nodes[v]);
///     weights.push(w);
/// }
///
/// let path = bidijkstra::undirected(&g, EdgeVec::from(weights), nodes[0], nodes[4]);
/// assert_eq!(path.into_iter().map(|e| (g.src(e), g.snk(e))).collect::<Vec<_>>(),
///            vec![(nodes[0],nodes[2]), (nodes[2],nodes[5]), (nodes[5],nodes[4])]);
/// ```
pub fn undirected<'a, G, Ws, W>(g: &'a G, weights: Ws, src: G::Node, snk: G::Node) -> Vec<G::Edge>
where
    G: IndexGraph<'a>,
    G::Node: Hash,
    Ws: EdgeMap<'a, G, W>,
    W: NumAssign + Ord + Copy,
{
    let mut d = BiDijkstra::<_, _, _, _, BinHeap<NodeItem<_, _, W>, usize>>::new(g);
    d.undirected(weights, src, snk)
}

/// Solve shortest path with bidirectional Dijkstra as bidirected graph.
///
/// The graph is considered as bidirected graph with different weights
/// for each direction.
///
/// - `g` the graph
/// - `weights` the (non-negative) arc weights
/// - `src` the source node
/// - `snk` the sink node
///
/// Return the incoming arc for each node forming the shortest path
/// tree.
pub fn bidirected<'a, G, Ws, W>(g: &'a G, weights: Ws, src: G::Node, snk: G::Node) -> Vec<G::Edge>
where
    G: IndexNetwork<'a>,
    G::Node: Hash,
    Ws: EdgeMap<'a, G, W>,
    W: NumAssign + Ord + Copy,
{
    let mut d = BiDijkstra::<_, _, _, _, BinHeap<NodeItem<_, _, W>, usize>>::new(g);
    d.bidirected(weights, src, snk)
}

/// Solve shortest path with bidirectional Dijkstra as directed graph.
///
/// The graph is considered directed, travel is only allowed along
/// outgoing edges.
///
/// - `g` the graph
/// - `weights` the (non-negative) arc weights
/// - `src` the source node
/// - `snk` the sink node
///
/// Return the incoming arc for each node forming the shortest path
/// tree.
pub fn directed<'a, G, Ws, W>(g: &'a G, weights: Ws, src: G::Node, snk: G::Node) -> Vec<G::Edge>
where
    G: IndexDigraph<'a>,
    G::Node: Hash,
    Ws: EdgeMap<'a, G, W>,
    W: NumAssign + Ord + Copy,
{
    let mut d = BiDijkstra::<_, _, _, _, BinHeap<NodeItem<_, _, W>, usize>>::new(g);
    d.directed(weights, src, snk)
}

pub fn generic<'a, G, W, Weights, Out, OutIt, In, InIt, H>(
    g: &'a G,
    weights: Weights,
    src: G::Node,
    snk: G::Node,
    outedges: Out,
    inedges: In,
) -> Vec<G::Edge>
where
    G: IndexGraph<'a>,
    G::Node: Hash,
    W: NumAssign + Ord + Copy,
    Weights: Index<G::Edge, Output = W>,
    Out: Fn(&'a G, G::Node) -> OutIt,
    OutIt: Iterator<Item = (G::Edge, G::Node)>,
    In: Fn(&'a G, G::Node) -> InIt,
    InIt: Iterator<Item = (G::Edge, G::Node)>,
    H: Heap<Key = NodeItem<G::Node, G::Edge, W>, Item = usize>,
{
    let mut d = BiDijkstra::<_, _, _, _, BinHeap<NodeItem<_, _, W>, usize>>::new(g);
    d.generic(weights, src, snk, outedges, inedges)
}

#[derive(Clone, Copy)]
pub struct NodeItem<Node: Copy, Edge: Copy, W: Copy + PartialOrd> {
    node: Node,
    pred: Option<(Node, Edge)>,
    weight: W,
}

impl<Node: Copy, Edge: Copy, W: PartialOrd + Copy> PartialEq for NodeItem<Node, Edge, W> {
    fn eq(&self, other: &Self) -> bool {
        self.weight.eq(&other.weight)
    }
}

impl<Node: Copy, Edge: Copy, W: PartialOrd + Copy> PartialOrd for NodeItem<Node, Edge, W> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        self.weight.partial_cmp(&other.weight)
    }
}

#[derive(Clone, Copy)]
enum NodeState<T> {
    /// Node has been seen (heap item).
    Seen(usize),
    /// Node has been finished.
    Finished(T),
}