import numpy as np
from scipy.linalg import svd, diagsvd, inv, pinv, null_space, norm
from numpy.linalg import matrix_rank

np.set_printoptions(precision=2, floatmode="fixed", suppress=True)

rng = np.random.default_rng(1234)
mean, stdev = 0, 1

# we might convince ourselves that all works for complex data as well
# then the ^H operator (conj().T) needs to be used instead of just .T
use_complex = False

M = 3  # number of rows
N = 7  # number of cols

rank = min(M, N)  # set desired rank == full row rank == independent columns
print("desired rank of A:", rank)

if use_complex:
    dtype = "complex128"
    A = np.zeros([M, N], dtype=dtype)
    for i in range(rank):
        col = rng.normal(mean, stdev, M) + 1j * rng.normal(mean, stdev, M)
        row = rng.normal(mean, stdev, N) + 1j * rng.normal(mean, stdev, N)
        A += np.outer(col, row)  # superposition of rank-1 matrices
else:
    dtype = "float64"
    A = np.zeros([M, N], dtype=dtype)
    for i in range(rank):
        col = rng.normal(mean, stdev, M)
        row = rng.normal(mean, stdev, N)
        A += np.outer(col, row)  # superposition of rank-1 matrices
# check if rng produced desired rank
print("        rank of A:", matrix_rank(A))
print("flat/fat matrix with full row rank")
print("-> matrix U contains only the column space")
print("-> left null space is only the zero vector")
print("A =\n", A)

[U, s, Vh] = svd(A, full_matrices=True)
S = diagsvd(s, M, N)  # for full SVD the matrix S has same dim as A
V = Vh.conj().T
Uh = U.conj().T

print("U =\n", U)
# number of non-zero sing values along diag must match rank
print("non-zero singular values: ", s[:rank])
print("S =\n", S)  # contains 0 Matrix right of diagonal part
print("V =\n", V)

# all stuff that is in matrix U
print("U =\n", U)

# column space C(A)
print("\ncolumn space (orthogonal to left null space):")
print(U[:, :rank])

# left null space, if empty only 0 vector
print("left null space (orthogonal to column space):")
print(U[:, rank:])  # for full row rank this is only the zero vector

print("###")

# all stuff that is in matrix V
print("\nV =\n", V)

# row space
print("\nrow space (orthogonal to null space):")
print(V[:, :rank])

# null space N(A), if empty only 0 vector
print("null space (orthogonal to row space):")
print(V[:, rank:])

[U, s, Vh] = svd(A, full_matrices=True)
V = Vh.conj().T
Uh = U.conj().T

Si = diagsvd(1 / s, N, M)  # works if array s has only non-zero entries
print("Inverse singular value matrix with bottom zero block")
print("Si =\n", Si)
# right inverse using 'inverse' SVD:
Ari = V @ Si @ U.conj().T
# right inverse using a dedicated pinv algorithm
# proper choice is done by pinv() itself
Ari_pinv = pinv(A)
print("pinv() == right inverse via SVD?", np.allclose(Ari, Ari_pinv))
print("S @ Si = \n", S @ Si, "\nyields MxM identity matrix")
print("A @ Ari = \n", A @ Ari, "\nyields MxM identity matrix")

v_row = 1 * V[:, 0, None]
v_null = 2 * V[:, rank, None]
v_tmp = v_row + v_null

# projection onto row space
P_CAH = Ari @ A
print(
    "projection matrix P_CAH projects V-space stuff to row space:\n",
    "P_CAH @ v_tmp == v_row:",
    np.allclose(P_CAH @ v_tmp, v_row),
)

# projection onto column space
# full rank and identity since we don't have a left null space
# so each column space vector is projected onto itself
P_CA = A @ Ari  # = I_MxM
print(
    "projection matrix P_CA projects a column space vector onto itself:\n",
    "P_CA @ U[:, 0] == U[:, 0]:",
    np.allclose(P_CA @ U[:, 0], U[:, 0]),
)

# projection onto null space, for flat/fat this space might be very large
P_NA = np.eye(N, N) - P_CAH
print(
    "projection matrix P_NA projects V-space stuff to null space:\n",
    "P_NA @ v_tmp == v_null:",
    np.allclose(P_NA @ v_tmp, v_null),
)

# projection onto left null space
P_NAH = np.eye(M, M) - P_CA  # == null matrix
print("projection matrix P_NAH is the MxM zero matrix\n", P_NAH)