# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.

import unittest

import numpy

import retworkx


class TestFloydWarshall(unittest.TestCase):

    def test_floyd_warshall(self):
        """Test the algorithm on a 5q x 4 depth circuit."""
        dag = retworkx.PyDAG()
        # inputs
        qr_0 = dag.add_node('qr[0]')
        qr_1 = dag.add_node('qr[1]')
        qr_2 = dag.add_node('qr[2]')
        cr_0 = dag.add_node('cr[0]')
        cr_1 = dag.add_node('cr[1]')
        # wires
        cx_1 = dag.add_node('cx_1')
        dag.add_edge(qr_0, cx_1, 'qr[0]')
        dag.add_edge(qr_1, cx_1, 'qr[1]')
        h_1 = dag.add_node('h_1')
        dag.add_edge(cx_1, h_1, 'qr[0]')
        cx_2 = dag.add_node('cx_2')
        dag.add_edge(cx_1, cx_2, 'qr[1]')
        dag.add_edge(qr_2, cx_2, 'qr[2]')
        cx_3 = dag.add_node('cx_3')
        dag.add_edge(h_1, cx_3, 'qr[0]')
        dag.add_edge(cx_2, cx_3, 'qr[2]')
        h_2 = dag.add_node('h_2')
        dag.add_edge(cx_3, h_2, 'qr[2]')
        # # outputs
        qr_0_out = dag.add_node('qr[0]_out')
        dag.add_edge(cx_3, qr_0_out, 'qr[0]')
        qr_1_out = dag.add_node('qr[1]_out')
        dag.add_edge(cx_2, qr_1_out, 'qr[1]')
        qr_2_out = dag.add_node('qr[2]_out')
        dag.add_edge(h_2, qr_2_out, 'qr[2]')
        cr_0_out = dag.add_node('cr[0]_out')
        dag.add_edge(cr_0, cr_0_out, 'qr[2]')
        cr_1_out = dag.add_node('cr[1]_out')
        dag.add_edge(cr_1, cr_1_out, 'cr[1]')

        result = retworkx.floyd_warshall(dag)
        expected = {
            0: {0: 0, 5: 1, 6: 2, 7: 2, 8: 3, 9: 4, 10: 4, 11: 3, 12: 5},
            1: {1: 0, 5: 1, 6: 2, 7: 2, 8: 3, 9: 4, 10: 4, 11: 3, 12: 5},
            2: {2: 0, 7: 1, 8: 2, 9: 3, 10: 3, 11: 2, 12: 4},
            3: {3: 0, 13: 1},
            4: {4: 0, 14: 1},
            5: {5: 0, 6: 1, 7: 1, 8: 2, 9: 3, 10: 3, 11: 2, 12: 4},
            6: {6: 0, 8: 1, 9: 2, 10: 2, 12: 3},
            7: {7: 0, 8: 1, 9: 2, 10: 2, 11: 1, 12: 3},
            8: {8: 0, 9: 1, 10: 1, 12: 2},
            9: {9: 0, 12: 1},
            10: {10: 0},
            11: {11: 0},
            12: {12: 0},
            13: {13: 0},
            14: {14: 0},
        }
        self.assertDictEqual(result, expected)

    def test_floyd_warshall_numpy_three_edges(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(6)))
        weights = [2, 12, 1, 5, 1]
        graph.add_edges_from([(i, i + 1, weights[i]) for i in range(5)])
        graph.add_edge(5, 0, 10)
        dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
        self.assertEqual(dist[0, 3], 15)
        self.assertEqual(dist[3, 0], 15)

    def test_weighted_numpy_two_edges(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(8)))
        graph.add_edges_from([
            (0, 1, 2),
            (1, 2, 2),
            (2, 3, 1),
            (3, 4, 1),
            (4, 5, 1),
            (5, 6, 1),
            (6, 7, 1),
            (7, 0, 1),
        ])
        dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
        self.assertEqual(dist[0, 2], 4)
        self.assertEqual(dist[2, 0], 4)

    def test_weighted_numpy_negative_cycle(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(4)))
        graph.add_edges_from([
            (0, 1, 1),
            (1, 2, -1),
            (2, 3, -1),
            (3, 0, -1),
        ])
        dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
        expected = numpy.array(
            [[-6, -7, -8, -13],
             [-7, -8, -9, -14],
             [-8, -9, -10, -15],
             [-13, -14, -15, -20]], dtype=numpy.float64)
        self.assertTrue(numpy.array_equal(dist, expected))

    def test_floyd_warshall_numpy_cycle(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(7)))
        graph.add_edges_from_no_data(
            [(0, 1), (0, 6), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)])
        dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: 1)
        self.assertEqual(dist[0, 3], 3)
        self.assertEqual(dist[0, 4], 3)

    def test_directed_floyd_warshall_numpy_cycle_as_undirected(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(7)))
        graph.add_edges_from_no_data(
            [(0, 1), (0, 6), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)])
        dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: 1,
                                                     as_undirected=True)
        expected = numpy.array([[0., 1., 2., 3., 3., 2., 1.],
                                [1., 0., 1., 2., 3., 3., 2.],
                                [2., 1., 0., 1., 2., 3., 3.],
                                [3., 2., 1., 0., 1., 2., 3.],
                                [3., 3., 2., 1., 0., 1., 2.],
                                [2., 3., 3., 2., 1., 0., 1.],
                                [1., 2., 3., 3., 2., 1., 0.]])
        self.assertTrue(numpy.array_equal(dist, expected))

    def test_floyd_warshall_numpy_digraph_three_edges(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(6)))
        weights = [2, 12, 1, 5, 1]
        graph.add_edges_from([(i, i + 1, weights[i]) for i in range(5)])
        graph.add_edge(5, 0, 10)
        dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
        self.assertEqual(dist[0, 3], 15)
        self.assertEqual(dist[3, 0], 16)

    def test_weighted_numpy_digraph_two_edges(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(8)))
        graph.add_edges_from([
            (0, 1, 2),
            (1, 2, 2),
            (2, 3, 1),
            (3, 4, 1),
            (4, 5, 1),
            (5, 6, 1),
            (6, 7, 1),
            (7, 0, 1),
        ])
        dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
        self.assertEqual(dist[0, 2], 4)
        self.assertEqual(dist[2, 0], 6)

    def test_floyd_warshall_numpy_digraph_cycle(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(7)))
        graph.add_edges_from_no_data(
            [(0, 1), (0, 6), (1, 2), (2, 3), (3, 4), (4, 5), (5, 6)])
        dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: 1)
        self.assertEqual(dist[0, 3], 3)
        self.assertEqual(dist[0, 4], 4)

    def test_weighted_numpy_directed_negative_cycle(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(4)))
        graph.add_edges_from([
            (0, 1, 1),
            (1, 2, -1),
            (2, 3, -1),
            (3, 0, -1),
        ])
        dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
        expected = numpy.array(
            [[-2, -1, -2, -3],
             [-3, -2, -3, -4],
             [-2, -1, -2, -3],
             [-3, -2, -3, -4]], dtype=numpy.float64)
        self.assertTrue(numpy.array_equal(dist, expected))

    def test_numpy_directed_no_edges(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(4)))
        dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: x)
        expected = numpy.full((4, 4), numpy.inf)
        numpy.fill_diagonal(expected, 0)
        self.assertTrue(numpy.array_equal(dist, expected))

    def test_numpy_no_edges(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(4)))
        dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: x)
        expected = numpy.full((4, 4), numpy.inf)
        numpy.fill_diagonal(expected, 0)
        self.assertTrue(numpy.array_equal(dist, expected))

    def test_floyd_warshall_numpy_digraph_cycle_with_removals(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(8)))
        graph.remove_node(0)
        graph.add_edges_from_no_data(
            [(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
        dist = retworkx.digraph_floyd_warshall_numpy(graph, lambda x: 1)
        self.assertEqual(dist[0, 3], 3)
        self.assertEqual(dist[0, 4], 4)

    def test_floyd_warshall_numpy_graph_cycle_with_removals(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(8)))
        graph.remove_node(0)
        graph.add_edges_from_no_data(
            [(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
        dist = retworkx.graph_floyd_warshall_numpy(graph, lambda x: 1)
        self.assertEqual(dist[0, 3], 3)
        self.assertEqual(dist[0, 4], 3)

    def test_floyd_warshall_numpy_digraph_cycle_no_weight_fn(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(8)))
        graph.remove_node(0)
        graph.add_edges_from_no_data(
            [(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
        dist = retworkx.digraph_floyd_warshall_numpy(graph)
        self.assertEqual(dist[0, 3], 3)
        self.assertEqual(dist[0, 4], 4)

    def test_floyd_warshall_numpy_graph_cycle_no_weight_fn(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(8)))
        graph.remove_node(0)
        graph.add_edges_from_no_data(
            [(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
        dist = retworkx.graph_floyd_warshall_numpy(graph)
        self.assertEqual(dist[0, 3], 3)
        self.assertEqual(dist[0, 4], 3)

    def test_floyd_warshall_numpy_digraph_cycle_default_weight(self):
        graph = retworkx.PyDiGraph()
        graph.add_nodes_from(list(range(8)))
        graph.remove_node(0)
        graph.add_edges_from_no_data(
            [(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
        dist = retworkx.digraph_floyd_warshall_numpy(graph, default_weight=2)
        self.assertEqual(dist[0, 3], 6)
        self.assertEqual(dist[0, 4], 8)

    def test_floyd_warshall_numpy_graph_cycle_default_weight(self):
        graph = retworkx.PyGraph()
        graph.add_nodes_from(list(range(8)))
        graph.remove_node(0)
        graph.add_edges_from_no_data(
            [(1, 2), (1, 7), (2, 3), (3, 4), (4, 5), (5, 6), (6, 7)])
        dist = retworkx.graph_floyd_warshall_numpy(graph, default_weight=2)
        self.assertEqual(dist[0, 3], 6)
        self.assertEqual(dist[0, 4], 6)