# test_disjoint_paths.py - Tests for flow based node and edge disjoint paths.
#
# Copyright 2016 NetworkX developers.
#
# This file is part of NetworkX.
#
# NetworkX is distributed under a BSD license; see LICENSE.txt for more
# information.
import pytest

import networkx as nx
from networkx.algorithms import flow
from networkx.utils import pairwise

flow_funcs = [
    flow.boykov_kolmogorov,
    flow.edmonds_karp,
    flow.dinitz,
    flow.preflow_push,
    flow.shortest_augmenting_path,
]

msg = "Assertion failed in function: {0}"


def is_path(G, path):
    return all(v in G[u] for u, v in pairwise(path))


def are_edge_disjoint_paths(G, paths):
    if not paths:
        return False
    for path in paths:
        assert is_path(G, path)
    paths_edges = [list(pairwise(p)) for p in paths]
    num_of_edges = sum(len(e) for e in paths_edges)
    num_unique_edges = len(set.union(*[set(es) for es in paths_edges]))
    if num_of_edges == num_unique_edges:
        return True
    return False


def are_node_disjoint_paths(G, paths):
    if not paths:
        return False
    for path in paths:
        assert is_path(G, path)
    # first and last nodes are source and target
    st = {paths[0][0], paths[0][-1]}
    num_of_nodes = len([n for path in paths for n in path if n not in st])
    num_unique_nodes = len({n for path in paths for n in path if n not in st})
    if num_of_nodes == num_unique_nodes:
        return True
    return False


def test_graph_from_pr_2053():
    G = nx.Graph()
    G.add_edges_from([
        ('A', 'B'), ('A', 'D'), ('A', 'F'), ('A', 'G'),
        ('B', 'C'), ('B', 'D'), ('B', 'G'), ('C', 'D'),
        ('C', 'E'), ('C', 'Z'), ('D', 'E'), ('D', 'F'),
        ('E', 'F'), ('E', 'Z'), ('F', 'Z'), ('G', 'Z')])
    for flow_func in flow_funcs:
        kwargs = dict(flow_func=flow_func)
        # edge disjoint paths
        edge_paths = list(nx.edge_disjoint_paths(G, 'A', 'Z', **kwargs))
        assert are_edge_disjoint_paths(G, edge_paths), msg.format(flow_func.__name__)
        assert nx.edge_connectivity(G, 'A', 'Z') == len(edge_paths), msg.format(flow_func.__name__)
        # node disjoint paths
        node_paths = list(nx.node_disjoint_paths(G, 'A', 'Z', **kwargs))
        assert are_node_disjoint_paths(G, node_paths), msg.format(flow_func.__name__)
        assert nx.node_connectivity(G, 'A', 'Z') == len(node_paths), msg.format(flow_func.__name__)


def test_florentine_families():
    G = nx.florentine_families_graph()
    for flow_func in flow_funcs:
        kwargs = dict(flow_func=flow_func)
        # edge disjoint paths
        edge_dpaths = list(nx.edge_disjoint_paths(G, 'Medici', 'Strozzi', **kwargs))
        assert are_edge_disjoint_paths(G, edge_dpaths), msg.format(flow_func.__name__)
        assert nx.edge_connectivity(G, 'Medici', 'Strozzi') == len(edge_dpaths), msg.format(flow_func.__name__)
        # node disjoint paths
        node_dpaths = list(nx.node_disjoint_paths(G, 'Medici', 'Strozzi', **kwargs))
        assert are_node_disjoint_paths(G, node_dpaths), msg.format(flow_func.__name__)
        assert nx.node_connectivity(G, 'Medici', 'Strozzi') == len(node_dpaths), msg.format(flow_func.__name__)


def test_karate():
    G = nx.karate_club_graph()
    for flow_func in flow_funcs:
        kwargs = dict(flow_func=flow_func)
        # edge disjoint paths
        edge_dpaths = list(nx.edge_disjoint_paths(G, 0, 33, **kwargs))
        assert are_edge_disjoint_paths(G, edge_dpaths), msg.format(flow_func.__name__)
        assert nx.edge_connectivity(G, 0, 33) == len(edge_dpaths), msg.format(flow_func.__name__)
        # node disjoint paths
        node_dpaths = list(nx.node_disjoint_paths(G, 0, 33, **kwargs))
        assert are_node_disjoint_paths(G, node_dpaths), msg.format(flow_func.__name__)
        assert nx.node_connectivity(G, 0, 33) == len(node_dpaths), msg.format(flow_func.__name__)


def test_petersen_disjoint_paths():
    G = nx.petersen_graph()
    for flow_func in flow_funcs:
        kwargs = dict(flow_func=flow_func)
        # edge disjoint paths
        edge_dpaths = list(nx.edge_disjoint_paths(G, 0, 6, **kwargs))
        assert are_edge_disjoint_paths(G, edge_dpaths), msg.format(flow_func.__name__)
        assert 3 == len(edge_dpaths), msg.format(flow_func.__name__)
        # node disjoint paths
        node_dpaths = list(nx.node_disjoint_paths(G, 0, 6, **kwargs))
        assert are_node_disjoint_paths(G, node_dpaths), msg.format(flow_func.__name__)
        assert 3 == len(node_dpaths), msg.format(flow_func.__name__)


def test_octahedral_disjoint_paths():
    G = nx.octahedral_graph()
    for flow_func in flow_funcs:
        kwargs = dict(flow_func=flow_func)
        # edge disjoint paths
        edge_dpaths = list(nx.edge_disjoint_paths(G, 0, 5, **kwargs))
        assert are_edge_disjoint_paths(G, edge_dpaths), msg.format(flow_func.__name__)
        assert 4 == len(edge_dpaths), msg.format(flow_func.__name__)
        # node disjoint paths
        node_dpaths = list(nx.node_disjoint_paths(G, 0, 5, **kwargs))
        assert are_node_disjoint_paths(G, node_dpaths), msg.format(flow_func.__name__)
        assert 4 == len(node_dpaths), msg.format(flow_func.__name__)


def test_icosahedral_disjoint_paths():
    G = nx.icosahedral_graph()
    for flow_func in flow_funcs:
        kwargs = dict(flow_func=flow_func)
        # edge disjoint paths
        edge_dpaths = list(nx.edge_disjoint_paths(G, 0, 6, **kwargs))
        assert are_edge_disjoint_paths(G, edge_dpaths), msg.format(flow_func.__name__)
        assert 5 == len(edge_dpaths), msg.format(flow_func.__name__)
        # node disjoint paths
        node_dpaths = list(nx.node_disjoint_paths(G, 0, 6, **kwargs))
        assert are_node_disjoint_paths(G, node_dpaths), msg.format(flow_func.__name__)
        assert 5 == len(node_dpaths), msg.format(flow_func.__name__)


def test_cutoff_disjoint_paths():
    G = nx.icosahedral_graph()
    for flow_func in flow_funcs:
        kwargs = dict(flow_func=flow_func)
        for cutoff in [2, 4]:
            kwargs['cutoff'] = cutoff
            # edge disjoint paths
            edge_dpaths = list(nx.edge_disjoint_paths(G, 0, 6, **kwargs))
            assert are_edge_disjoint_paths(G, edge_dpaths), msg.format(flow_func.__name__)
            assert cutoff == len(edge_dpaths), msg.format(flow_func.__name__)
            # node disjoint paths
            node_dpaths = list(nx.node_disjoint_paths(G, 0, 6, **kwargs))
            assert are_node_disjoint_paths(G, node_dpaths), msg.format(flow_func.__name__)
            assert cutoff == len(node_dpaths), msg.format(flow_func.__name__)


def test_missing_source_edge_paths():
    with pytest.raises(nx.NetworkXError):
        G = nx.path_graph(4)
        list(nx.edge_disjoint_paths(G, 10, 1))


def test_missing_source_node_paths():
    with pytest.raises(nx.NetworkXError):
        G = nx.path_graph(4)
        list(nx.node_disjoint_paths(G, 10, 1))


def test_missing_target_edge_paths():
    with pytest.raises(nx.NetworkXError):
        G = nx.path_graph(4)
        list(nx.edge_disjoint_paths(G, 1, 10))


def test_missing_target_node_paths():
    with pytest.raises(nx.NetworkXError):
        G = nx.path_graph(4)
        list(nx.node_disjoint_paths(G, 1, 10))


def test_not_weakly_connected_edges():
    with pytest.raises(nx.NetworkXNoPath):
        G = nx.DiGraph()
        nx.add_path(G, [1, 2, 3])
        nx.add_path(G, [4, 5])
        list(nx.edge_disjoint_paths(G, 1, 5))


def test_not_weakly_connected_nodes():
    with pytest.raises(nx.NetworkXNoPath):
        G = nx.DiGraph()
        nx.add_path(G, [1, 2, 3])
        nx.add_path(G, [4, 5])
        list(nx.node_disjoint_paths(G, 1, 5))


def test_not_connected_edges():
    with pytest.raises(nx.NetworkXNoPath):
        G = nx.Graph()
        nx.add_path(G, [1, 2, 3])
        nx.add_path(G, [4, 5])
        list(nx.edge_disjoint_paths(G, 1, 5))


def test_not_connected_nodes():
    with pytest.raises(nx.NetworkXNoPath):
        G = nx.Graph()
        nx.add_path(G, [1, 2, 3])
        nx.add_path(G, [4, 5])
        list(nx.node_disjoint_paths(G, 1, 5))


def test_isolated_edges():
    with pytest.raises(nx.NetworkXNoPath):
        G = nx.Graph()
        G.add_node(1)
        nx.add_path(G, [4, 5])
        list(nx.edge_disjoint_paths(G, 1, 5))


def test_isolated_nodes():
    with pytest.raises(nx.NetworkXNoPath):
        G = nx.Graph()
        G.add_node(1)
        nx.add_path(G, [4, 5])
        list(nx.node_disjoint_paths(G, 1, 5))


def test_invalid_auxiliary():
    with pytest.raises(nx.NetworkXError):
        G = nx.complete_graph(5)
        list(nx.node_disjoint_paths(G, 0, 3, auxiliary=G))