#  ___________________________________________________________________________
#
#  Pyomo: Python Optimization Modeling Objects
#  Copyright (c) 2008-2022
#  National Technology and Engineering Solutions of Sandia, LLC
#  Under the terms of Contract DE-NA0003525 with National Technology and
#  Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
#  rights in this software.
#  This software is distributed under the 3-clause BSD License.
#  ___________________________________________________________________________

import random
from pyomo.contrib.incidence_analysis.matching import maximum_matching
from pyomo.contrib.incidence_analysis.triangularize import (
        block_triangularize,
        get_diagonal_blocks,
        )
from pyomo.common.dependencies import (
        scipy,
        scipy_available,
        networkx_available,
        )
if scipy_available:
    sps = scipy.sparse

import pyomo.common.unittest as unittest


@unittest.skipUnless(networkx_available, "networkx is not available")
@unittest.skipUnless(scipy_available, "scipy is not available")
class TestTriangularize(unittest.TestCase):
    def test_low_rank_exception(self):
        N = 5
        row = list(range(N-1))
        col = list(range(N-1))
        data = [1 for _ in range(N-1)]

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        with self.assertRaises(ValueError) as exc:
            row_block_map, col_block_map = block_triangularize(matrix)
        self.assertIn('perfect matching', str(exc.exception))

    def test_non_square_exception(self):
        N = 5
        row = list(range(N-1))
        col = list(range(N-1))
        data = [1 for _ in range(N-1)]

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N-1))

        with self.assertRaises(ValueError) as exc:
            row_block_map, col_block_map = block_triangularize(matrix)
        self.assertIn('non-square matrices', str(exc.exception))

    def test_identity(self):
        N = 5
        matrix = sps.identity(N).tocoo()
        row_block_map, col_block_map = block_triangularize(matrix)
        row_values = set(row_block_map.values())
        col_values = set(row_block_map.values())

        # For a (block) diagonal matrix, the order of diagonal
        # blocks is arbitary, so we can't perform any strong
        # checks here.
        #
        # Perfect matching is unique, but order of strongly
        # connected components is not.

        self.assertEqual(len(row_block_map), N)
        self.assertEqual(len(col_block_map), N)
        self.assertEqual(len(row_values), N)
        self.assertEqual(len(col_values), N)

        for i in range(N):
            self.assertIn(i, row_block_map)
            self.assertIn(i, col_block_map)
            self.assertIn(i, row_values)
            self.assertIn(i, col_values)

    def test_lower_tri(self):
        """
        This matrix has a unique maximal matching and SCC
        order, making it a good test for a "fully decomposable"
        matrix.
        |x        |
        |x x      |
        |  x x    |
        |    x x  |
        |      x x|
        """
        N = 5
        row = []
        col = []
        data = []
        # Diagonal
        row.extend(range(N))
        col.extend(range(N))
        data.extend(1 for _ in range(N))

        # Below diagonal
        row.extend(range(1, N))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        row_block_map, col_block_map = block_triangularize(matrix)
        row_values = set(row_block_map.values())
        col_values = set(row_block_map.values())
        
        self.assertEqual(len(row_values), N)
        self.assertEqual(len(col_values), N)

        for i in range(N):
            self.assertEqual(row_block_map[i], i)
            self.assertEqual(col_block_map[i], i)

    def test_upper_tri(self):
        """
        This matrix has a unique maximal matching and SCC
        order, making it a good test for a "fully decomposable"
        matrix.
        |x x      |
        |  x x    |
        |    x x  |
        |      x x|
        |        x|
        """
        N = 5
        row = []
        col = []
        data = []
        # Diagonal
        row.extend(range(N))
        col.extend(range(N))
        data.extend(1 for _ in range(N))

        # Below diagonal
        row.extend(range(N-1))
        col.extend(range(1, N))
        data.extend(1 for _ in range(N-1))

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        row_block_map, col_block_map = block_triangularize(matrix)
        row_values = set(row_block_map.values())
        col_values = set(row_block_map.values())
        
        self.assertEqual(len(row_values), N)
        self.assertEqual(len(col_values), N)

        for i in range(N):
            # The block_triangularize function permutes
            # to lower triangular form, so rows and
            # columns are transposed to assemble the blocks.
            self.assertEqual(row_block_map[i], N-1-i)
            self.assertEqual(col_block_map[i], N-1-i)

    def test_bordered(self):
        """
        This matrix is non-decomposable
        |x       x|
        |  x     x|
        |    x   x|
        |      x x|
        |x x x x  |
        """
        N = 5
        row = []
        col = []
        data = []
        # Diagonal
        row.extend(range(N-1))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        # Bottom row
        row.extend(N-1 for _ in range(N-1))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        # Right column
        row.extend(range(N-1))
        col.extend(N-1 for _ in range(N-1))
        data.extend(1 for _ in range(N-1))

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        row_block_map, col_block_map = block_triangularize(matrix)
        row_values = set(row_block_map.values())
        col_values = set(row_block_map.values())

        self.assertEqual(len(row_values), 1)
        self.assertEqual(len(col_values), 1)

        for i in range(N):
            self.assertEqual(row_block_map[i], 0)
            self.assertEqual(col_block_map[i], 0)

    def test_decomposable_bordered(self):
        """
        This matrix decomposes
        |x        |
        |  x      |
        |    x   x|
        |      x x|
        |x x x x  |
        """
        N = 5
        half = N//2
        row = []
        col = []
        data = []

        # Diagonal
        row.extend(range(N-1))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        # Bottom row
        row.extend(N-1 for _ in range(N-1))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        # Right column
        row.extend(range(half, N-1))
        col.extend(N-1 for _ in range(half, N-1))
        data.extend(1 for _ in range(half, N-1))

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        row_block_map, col_block_map = block_triangularize(matrix)
        row_values = set(row_block_map.values())
        col_values = set(row_block_map.values())

        self.assertEqual(len(row_values), half+1)
        self.assertEqual(len(col_values), half+1)

        first_half_set = set(range(half))
        for i in range(N):
            if i < half:
                # The first N//2 diagonal blocks are unordered
                self.assertIn(row_block_map[i], first_half_set)
                self.assertIn(col_block_map[i], first_half_set)
            else:
                self.assertEqual(row_block_map[i], half)
                self.assertEqual(col_block_map[i], half)

    def test_decomposable_tridiagonal(self):
        """
        This matrix decomposes into 2x2 blocks
        |x x      |
        |x x      |
        |  x x x  |
        |    x x  |
        |      x x|
        """
        N = 5
        row = []
        col = []
        data = []

        # Diagonal
        row.extend(range(N))
        col.extend(range(N))
        data.extend(1 for _ in range(N))

        # Below diagonal
        row.extend(range(1, N))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        # Above diagonal
        row.extend(i for i in range(N-1) if not i%2)
        col.extend(i+1 for i in range(N-1) if not i%2)
        data.extend(1 for i in range(N-1) if not i%2)

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        row_block_map, col_block_map = block_triangularize(matrix)
        row_values = set(row_block_map.values())
        col_values = set(row_block_map.values())

        self.assertEqual(len(row_values), (N+1)//2)
        self.assertEqual(len(col_values), (N+1)//2)

        for i in range((N+1)//2):
            self.assertEqual(row_block_map[2*i], i)
            self.assertEqual(col_block_map[2*i], i)

            if 2*i+1 < N:
                self.assertEqual(row_block_map[2*i+1], i)
                self.assertEqual(col_block_map[2*i+1], i)

    def test_decomposable_tridiagonal_shuffled(self):
        """
        This matrix decomposes into 2x2 blocks
        |x x      |
        |x x      |
        |  x x x  |
        |    x x  |
        |      x x|
        """
        N = 5
        row = []
        col = []
        data = []

        # Diagonal
        row.extend(range(N))
        col.extend(range(N))
        data.extend(1 for _ in range(N))

        # Below diagonal
        row.extend(range(1, N))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        # Above diagonal
        row.extend(i for i in range(N-1) if not i%2)
        col.extend(i+1 for i in range(N-1) if not i%2)
        data.extend(1 for i in range(N-1) if not i%2)

        # Same results hold after applying a random permutation.
        row_perm = list(range(N))
        col_perm = list(range(N))
        random.shuffle(row_perm)
        random.shuffle(col_perm)

        row = [row_perm[i] for i in row]
        col = [col_perm[j] for j in col]

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        row_block_map, col_block_map = block_triangularize(matrix)
        row_values = set(row_block_map.values())
        col_values = set(row_block_map.values())

        self.assertEqual(len(row_values), (N+1)//2)
        self.assertEqual(len(col_values), (N+1)//2)

        for i in range((N+1)//2):
            row_idx = row_perm[2*i]
            col_idx = col_perm[2*i]
            self.assertEqual(row_block_map[row_idx], i)
            self.assertEqual(col_block_map[col_idx], i)

            if 2*i+1 < N:
                row_idx = row_perm[2*i+1]
                col_idx = col_perm[2*i+1]
                self.assertEqual(row_block_map[row_idx], i)
                self.assertEqual(col_block_map[col_idx], i)

    def test_decomposable_tridiagonal_diagonal_blocks(self):
        """
        This matrix decomposes into 2x2 blocks
        |x x      |
        |x x      |
        |  x x x  |
        |    x x  |
        |      x x|
        """
        N = 5
        row = []
        col = []
        data = []

        # Diagonal
        row.extend(range(N))
        col.extend(range(N))
        data.extend(1 for _ in range(N))

        # Below diagonal
        row.extend(range(1, N))
        col.extend(range(N-1))
        data.extend(1 for _ in range(N-1))

        # Above diagonal
        row.extend(i for i in range(N-1) if not i%2)
        col.extend(i+1 for i in range(N-1) if not i%2)
        data.extend(1 for i in range(N-1) if not i%2)

        matrix = sps.coo_matrix((data, (row, col)), shape=(N, N))

        row_blocks, col_blocks = get_diagonal_blocks(matrix)

        self.assertEqual(len(row_blocks), (N+1)//2)
        self.assertEqual(len(col_blocks), (N+1)//2)

        for i in range((N+1)//2):
            rows = row_blocks[i]
            cols = col_blocks[i]
            
            if 2*i+1 < N:
                self.assertEqual(set(rows), {2*i, 2*i+1})
                self.assertEqual(set(cols), {2*i, 2*i+1})
            else:
                self.assertEqual(set(rows), {2*i})
                self.assertEqual(set(cols), {2*i})


if __name__ == "__main__":
    unittest.main()
