import numpy as np
from numpy import linalg as la
from scipy.linalg import block_diag
from scipy.stats import ortho_group

############### PURPOSE #################
# SBD is a simple and efficient algorithm for finding a (often finest) simultaneous block diagonalization of an arbitrary number of symmetric matrices.
# The algorithm works by finding the eigendecomposition of a single matrix, which is a random linear combination of all the matrices to be simultaneously block diagonalized.
# The current algorithm is inspired by the results presented in K. Murota et al. Jpn. J. Ind. Appl. Math 27, 125–160 (2010).
# For an efficient algorithm that also guarantees the optimal SBD, see https://github.com/y-z-zhang/net-sync-sym.
# There, you can also find concrete examples on how the SBD algorithm can be applied to analyze cluster synchronization patterns in complex networks.

############### USEAGE #################
# [P,BlockSizes] = sbd(A,threshold)
# A --- list containing the set of matrices to be simultaneously block diagonalized
# threshold --- real number, matrix entries below threshold are considered zero
# P --- orthogonal transformation matrix that performs SBD on A
# BlockSizes --- array listing the size of each common block

############### REFERENCE #################
# Y. Zhang, V. Latora, and A. E. Motter,
# Unified treatment of synchronization patterns in generalized networks with higher-order, multilayer, and temporal interactions,
# Commun. Phys. 4, 195 (2021).

def sbd(A,threshold):
	n = len(A[0])	# size of the matrices to be simultaneously block diagonalized
	m = len(A)		# number of matrices to be simultaneously block diagonalized
	BlockSizes = []	# initialize the array that lists the size of each common block

    # B is a random self-adjoint matrix generated by matrices from A (and their conjugate transposes)
	B = np.zeros((n,n))
	for p in range(m):
		B = B + np.random.normal()*(A[p]+A[p].transpose())

    # find the eigenvalues and eigenvectors of B
	D, V = la.eigh(B)

    # C is a matrix used to sort the column vectors of V (i.e., the base vectors)
    # such that the base vectors corresponding to the same common block are next to each other
	C = np.zeros((n,n))
	for p in range(m):
		C = C + np.random.normal()*(A[p]+A[p].transpose())
	C = V.transpose()@C@V
	#for p in range(m):
	#	C = C + np.random.normal()*V.transpose()@A[p]@V

    # arrays used to track which base vectors have been sorted
	remaining_basis = list(range(n))
	sorted_basis = []

    # the sorting process: find C_ij's that are nonzero and group the base vectors v_i and v_j together
	while len(remaining_basis) > 0:
		current_block = [remaining_basis[0]]
		current_block_size = 1
		if len(remaining_basis) > 1:
			for idx in remaining_basis[1:]:
				if np.abs(C[remaining_basis[0],idx]) > threshold:
					current_block.append(idx)
					current_block_size = current_block_size + 1

		for idx in current_block:
			sorted_basis.append(idx)
			remaining_basis.remove(idx)

		# do the following in case there are zero entries inside the block
		current_block_extra = []
		if len(remaining_basis) > 0:
			for idx in remaining_basis:
				for ind in current_block:
					if np.abs(C[ind,idx]) > threshold:
						current_block_extra.append(idx)
						current_block_size = current_block_size + 1
						break

		for idx in current_block_extra:
			sorted_basis.append(idx)
			remaining_basis.remove(idx)

		BlockSizes.append(current_block_size)

    # the sorted base vectors give the final orthogonal/unitary transformation matrix that performs SBD on A
	P = V[:,sorted_basis]

	return P, BlockSizes

## uncomment below for an example use of the algorithm
#n = 10	# size of the matrices to be simultaneously block diagonalized
## A1 and A2 are random matrices that share a predefined block structure
#A1 = block_diag(np.random.rand(3,3),np.random.rand(2,2),np.random.rand(4,4),np.random.rand(1,1))
#A2 = block_diag(np.random.rand(3,3),np.random.rand(2,2),np.random.rand(4,4),np.random.rand(1,1))
## transform A1 and A2 using a random orthogonal matrix Q
#Q = ortho_group.rvs(dim=n)
#A1 = Q.T@A1@Q
#A2 = Q.T@A2@Q
#A = [A1,A2]
## use the SBD algorithm to uncover the common block structure
#sbd(A,1e-5)
