import numpy as np
import numpy.linalg as lin
import matplotlib.pyplot as plt

matrix = np.array(
        [[0,0,0,0,0,0],
        [8,5,3,2,1,0],
        [4,4,3,4,2,0],
        [4,2,1,0,0,0],
        [4,6,5,8,4,0],
        [1,5,11,9,17,25]]
)

steps = 10
eigval, eigvec = lin.eig(matrix)
diagmat = np.zeros((6,6))
np.fill_diagonal(diagmat,eigval)
diagmat0 = diagmat
sum_mat = np.zeros((6,steps))
for i in range(steps):
        A = eigvec @ diagmat @ lin.inv(eigvec)
        for j in range(6):
                sum_mat[j][i] = np.sum(A[j])
        diagmat = diagmat0 @ diagmat

x = np.arange(1,steps+1,1)
label = ['square', 'deadend', 'corner', 'parallel', 'line', 'empty']
for i in range(6):
        plt.plot(x, sum_mat[i], label=label[i])
plt.yscale('log')
plt.plot((2,8),(5**5*10,5**17*10),'--',label='2')
log_5_11=np.log(11)/np.log(5)
plt.plot((2,8),(5**(log_5_11*2)*10,5**(log_5_11*8)*10),'--',label=r'$\log 11 / \log 5$')
plt.xlabel('steps')
plt.ylabel('n (steps)')
plt.legend()
plt.savefig('kochgraph.png', dpi=300, transparent=True)