import numpy as np
import numpy.linalg as ln
import scipy.linalg

dt = int
matr = np.array([[0, 1, 0], [-4, 4, 0], [-2, 1, 2]], dtype=dt)
assert matr.shape[0] == matr.shape[1]
n = matr.shape[0]
E = np.identity(n, dtype=dt)


def print_matrix(m):
    for i in range(m.shape[0]):
        for j in range(m.shape[1]):
            print(m[i, j], end=' ')
        print()


def make_cell(l, k):
    return l * np.eye(k, dtype=dt) + np.eye(k, k=1, dtype=dt)


def get_cells_for_value(l):
    nilp = matr - l * E
    print(f"A - {l}E:")
    print_matrix(nilp)
    k = n - ln.matrix_rank(nilp)
    print(f"dim Ker(A - {l}E) = {k}")
    cells = np.ones(k, dtype=int)
    k_prev = 0
    nilp_cur = nilp
    p = 1
    while True:
        k_prev = k
        p += 1
        nilp_cur = nilp_cur @ nilp
        k = n - ln.matrix_rank(nilp_cur)
        if k == k_prev:
            break
        print(f"(A - {l}E)^{p}:")
        print_matrix(nilp_cur)
        print(f"dim Ker(A - {l}E)^{p} = {k}")
        for i in range(k - k_prev):
            cells[i] += 1
    print(f"Cells for {l}:", end=' ')
    print_matrix(cells.reshape(1, -1))
    return scipy.linalg.block_diag(*[make_cell(l, x) for x in cells])


def jordanize():
    for l in np.unique(np.round(ln.eigvals(matr), 5)):
        print(f"For eigenvalue {l}:")
        get_cells_for_value(l)


def main():
    lambdas = np.unique(np.round(ln.eigvals(matr), 5))
    print(scipy.linalg.block_diag(*[get_cells_for_value(dt(l)) for l in lambdas]))


if __name__ == '__main__':
    main()