import numpy as np

import matplotlib.pyplot as plt

from mpl_toolkits.mplot3d import Axes3D

from ellipsoid_fit import ellipsoid_fit, ellipsoid_plot, data_regularize

def set_axes_equal(ax: plt.Axes):

ax.set_zlim3d([z - radius, z + radius])

if __name__=='__main__':

data = np.loadtxt("mag_out.txt")

data2 = data_regularize(data, divs=8)

center, evecs, radii, v = ellipsoid_fit(data2)

data_centered = data - center.T

data_centered_regularized = data2 - center.T

a, b, c = radii

r = (a * b * c) ** (1. / 3.)

D = np.array([[r/a, 0., 0.], [0., r/b, 0.], [0., 0., r/c]])

#http://www.cs.brandeis.edu/~cs155/Lecture_07_6.pdf

#affine transformation from ellipsoid to sphere (translation excluded)

TR = evecs.dot(D).dot(evecs.T)

data_on_sphere = TR.dot(data_centered_regularized.T).T

fig = plt.figure()

ax = fig.add_subplot(111, projection='3d')

# for direction in (-1, 1):

# for point in np.diag(direction * np.max(data) * np.array([1, 1, 1])):

# ax.plot([point[0]], [point[1]], [point[2]], 'w')

ax.scatter(data_centered[:,0], data_centered[:,1], data_centered[:,2], marker='o', color='g')

# ax.scatter(data_centered_regularized[:, 0], data_centered_regularized[:, 1],

# data_centered_regularized[:, 2], marker='o', color='b')

ax.scatter(data_on_sphere[:, 0], data_on_sphere[:, 1],

data_on_sphere[:, 2], marker='o', color='r')

ellipsoid_plot([0, 0, 0], radii, evecs, ax=ax, plot_axes=True, cage_color='g')

ellipsoid_plot([0, 0, 0], [r, r, r], evecs, ax=ax, plot_axes=True, cage_color='orange')

#ax.plot([r],[0],[0],color='r',marker='o')

#ax.plot([radii[0]],[0],[0],color='b',marker='o')

set_axes_equal(ax)

plt.show()