import numpy as np

def quadrupole_tensor(x, y, z, m=1):
    Qxx =       np.sum(m*x**2)
    Qyy =       np.sum(m*y**2)
    Qzz =       np.sum(m*z**2)
    Qxy = Qyx = np.sum(m*x*y)
    Qxz = Qzx = np.sum(m*x*z)
    Qzy = Qyz = np.sum(m*z*y)
    return np.array([[Qxx, Qxy, Qxz],
                     [Qyx, Qyy, Qyz],
                     [Qzx, Qzy, Qzz]])

def rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues):
    """
    Calculates a rotation matrix such that the eigenvector with the largest
    eigenvalue points in the x-direction, the eigenvector with the middle
    eigenvalue points in the y-direction, and so on.
    """
    return np.array([eigenvectors[:,i] for i in np.argsort(eigenvalues)[::-1]])

def generate_ellipsoid(n, power, a, b, c, euler_1, euler_2, euler_3, seed=None):
    from scipy.spatial.transform import Rotation
    if not (a >= b >= c):
        raise ValueError('Requires a >= b >= c')
    b_over_a = b/a
    c_over_a = c/a
    np.random.seed(seed)

    # Create a sphere of particles with radius a.
    r     = (np.random.rand(n)**power)*a
    theta = np.arccos(np.random.rand(n)*2-1)
    phi   = np.random.rand(n)*2*np.pi
    x = r*np.cos(phi)*np.sin(theta)
    y = r*np.sin(phi)*np.sin(theta)
    z = r*            np.cos(theta)

    # Scale the intermediate and minor axes
    y *= b_over_a
    z *= c_over_a

    # Rotate
    rotation = Rotation.from_euler('zxz', [euler_1, euler_2, euler_3])
    return rotation.apply(np.stack((x,y,z), axis=1))

if __name__ == '__main__': # Test if the above routines
    n = 65536
    power = 1.5
    a, b, c = 2.3, 0.92, 0.345 # or sorted(np.random.rand(3))[::-1]
    angles = 0.43, 1.81, 0.32  # or np.random.rand(3)*2*np.pi

    X = generate_ellipsoid(n, power, a, b, c, *angles, seed=None)
    Q = quadrupole_tensor(*X.T)
    eigenvalues, eigenvectors = np.linalg.eig(Q)
    i = np.argsort(eigenvalues)
    print('measured b/a = %.4f (expected %.4f)' % (np.sqrt(eigenvalues[i[1]]/eigenvalues[i[2]]), b/a))
    print('measured c/a = %.4f (expected %.4f)' % (np.sqrt(eigenvalues[i[0]]/eigenvalues[i[2]]), c/a))
    R = rotation_matrix_from_eigenvectors(eigenvectors, eigenvalues)
    print('Rotating vectors and recalculating quadrupole tensor...')
    X_new = (R @ X.T).T
    Q_new = quadrupole_tensor(*X_new.T)
    print('%11.4e %11.4e %11.4e' % tuple(Q_new[0,:]))
    print('%11.4e %11.4e %11.4e' % tuple(Q_new[1,:]))
    print('%11.4e %11.4e %11.4e' % tuple(Q_new[2,:]))
    print('... it should be almost diagonal.')