import sys
import numpy as np
import matplotlib.pyplot as plt

def cbrt(n):
    x = np.ones((n, n))
    x = np.triu(x, 1) # ones strictly above the diagonal
    x = x[:, ::-1] # mirrored left-right

    W, V = np.linalg.eig(x)
    W = np.cbrt(W)
    z = V.dot(np.diag(W)).dot(V.T)

    return z


def keypoints(z):
    n = len(z)
    m = n // 2
    s, q, k = z[m, 0], z[m, m//2], z[m, m-2]
    return s, q, k


do_plot = True
if do_plot:
    n = 1000
    z = cbrt(n)

    l = z[range(n), range(n)]
    plt.suptitle("diagonal of cuberoot matrix, n=%d" % n)
    plt.plot(range(len(l)), l)
    plt.savefig("cbrt-diag.pdf")
    plt.show()

    zm = z[:, ::-1]
    plt.suptitle("antidiagonals of cuberoot matrix, n=%d" % n)
    legend = []
    plt.gca().set_color_cycle(['red', 'orange', 'yellow', 'green', 'blue', 'purple', 'black'])
    for offset in range(-3, 4):
        v = np.diagonal(zm, offset)
        plt.plot(range(len(v)), v)
        legend.append("offset %+d" % offset)
    plt.xlabel("n")
    plt.legend(legend)
    plt.savefig("cbrt-antidiag.pdf")
    plt.show()

    l = z[n//2, :]
    plt.suptitle("middle row of cuberoot matrix, n=%d" % n)
    plt.plot(range(len(l)), l)
    plt.savefig("cbrt-middle.pdf")
    plt.show()


three_dee = True
if three_dee:
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib import cm

    n = 20
    z = cbrt(n)

    fig = plt.figure()
    fig.suptitle("cuberoot matrix, n=%d" % n)
    ax = fig.gca(projection='3d')

    x, y = np.meshgrid(range(n), range(n))

    surf = ax.plot_surface(x, y, z, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False)
    fig.colorbar(surf, shrink=0.5, aspect=5)
    plt.savefig("cbrt-surface.pdf")
    plt.show()


ss=[]
for n in range(20, 400, 20):
    z = cbrt(n)
    s, q, k = keypoints(z)
    ss.append((n,s,q,k))

ss = np.array(ss)

plt.plot(ss[:, 0], ss[:, 1])
plt.plot(ss[:, 0], ss[:, 2])
plt.plot(ss[:, 0], ss[:, 3])
plt.suptitle("values of key matrix elements")
plt.xlabel("n")
plt.legend(["edge", "quarter", "middle"])
plt.savefig("cbrt-chart.pdf")
plt.show()