# # ls_ellipsoid and polyToParams3D functions are taken from
# http://www.juddzone.com/ALGORITHMS/least_squares_3D_ellipsoid.html
# # FindBestFittingEllipsoid and GetConvexHull are adapted from a discussion at
# https://stackoverflow.com/questions/58501545/python-fit-3d-ellipsoid-oblate-prolate-to-3d-points


from scipy.spatial import ConvexHull
import numpy as np
from numpy.linalg import eig, inv


def ls_ellipsoid(xx, yy, zz):  # finds the ellipsoid
    # least squares fit to a 3D-ellipsoid
    #  Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz  = 1
    #
    # Note that sometimes it is expressed as a solution to
    #  Ax^2 + By^2 + Cz^2 + 2Dxy + 2Exz + 2Fyz + 2Gx + 2Hy + 2Iz  = 1
    # where the last six terms have a factor of 2 in them
    # This is in anticipation of forming a matrix with the polynomial coefficients.
    # Those terms with factors of 2 are all off diagonal elements.  These contribute
    # two terms when multiplied out (symmetric) so would need to be divided by two

    # change xx from vector of length N to Nx1 matrix so we can use hstack
    x = xx[:, np.newaxis]
    y = yy[:, np.newaxis]
    z = zz[:, np.newaxis]

    #  Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz = 1
    J = np.hstack((x * x, y * y, z * z, x * y, x * z, y * z, x, y, z))
    K = np.ones_like(x)  # column of ones

    # np.hstack performs a loop over all samples and creates
    # a row in J for each x,y,z sample:
    # J[ix,0] = x[ix]*x[ix]
    # J[ix,1] = y[ix]*y[ix]
    # etc.

    JT = J.transpose()
    JTJ = np.dot(JT, J)
    InvJTJ = np.linalg.inv(JTJ);
    ABC = np.dot(InvJTJ, np.dot(JT, K))

    # Rearrange, move the 1 to the other side
    #  Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz - 1 = 0
    #    or
    #  Ax^2 + By^2 + Cz^2 +  Dxy +  Exz +  Fyz +  Gx +  Hy +  Iz + J = 0
    #  where J = -1
    eansa = np.append(ABC, -1)

    return (eansa)


def polyToParams3D(vec, printMe):  # gets 3D parameters of an ellipsoid.
    # convert the polynomial form of the 3D-ellipsoid to parameters
    # center, axes, and transformation matrix
    # vec is the vector whose elements are the polynomial
    # coefficients A..J
    # returns (center, axes, rotation matrix)

    # Algebraic form: X.T * Amat * X --> polynomial form

    if printMe: print('\npolynomial\n', vec)

    Amat = np.array(
        [
            [vec[0], vec[3] / 2.0, vec[4] / 2.0, vec[6] / 2.0],
            [vec[3] / 2.0, vec[1], vec[5] / 2.0, vec[7] / 2.0],
            [vec[4] / 2.0, vec[5] / 2.0, vec[2], vec[8] / 2.0],
            [vec[6] / 2.0, vec[7] / 2.0, vec[8] / 2.0, vec[9]]
        ])

    if printMe: print('\nAlgebraic form of polynomial\n', Amat)

    # See B.Bartoni, Preprint SMU-HEP-10-14 Multi-dimensional Ellipsoidal Fitting
    # equation 20 for the following method for finding the center
    A3 = Amat[0:3, 0:3]
    A3inv = inv(A3)
    ofs = vec[6:9] / 2.0
    center = -np.dot(A3inv, ofs)
    if printMe: print('\nCenter at:', center)

    # Center the ellipsoid at the origin
    Tofs = np.eye(4)
    Tofs[3, 0:3] = center
    R = np.dot(Tofs, np.dot(Amat, Tofs.T))
    if printMe: print('\nAlgebraic form translated to center\n', R, '\n')

    R3 = R[0:3, 0:3]
    R3test = R3 / R3[0, 0]
    # print('normed \n',R3test)
    s1 = -R[3, 3]
    R3S = R3 / s1
    (el, ec) = eig(R3S)

    recip = 1.0 / np.abs(el)
    axes = np.sqrt(recip)
    if printMe: print('\nAxes are\n', axes, '\n')

    inve = inv(ec)  # inverse is actually the transpose here
    if printMe: print('\nRotation matrix\n', inve)
    return (center, axes, inve)


def FindBestFittingEllipsoid() -> tuple:
    # get convex hull
    convex_hull = GetConvexHull()

    # fit ellipsoid on the convex hull
    # # get ellipsoid polynomial coefficients
    ellipsoid_coef = ls_ellipsoid(convex_hull[0], convex_hull[1], convex_hull[2])
    # # get ellipsoid 3D parameters
    center, axes, orientation = polyToParams3D(ellipsoid_coef, False)

    return tuple(ellipsoid_coef, center, axes, orientation)


def GetConvexHull():
    volume = np.stack((conf.x, conf.y, conf.z), axis=-1)
    hull_volume = ConvexHull(volume)
    lH = len(hull_volume.vertices)
    hull = np.zeros((lH, 3))
    for i in range(len(hull_volume.vertices)):
        hull[i] = volume[hull_volume.vertices[i]]
    convex_hull = np.transpose(hull)

    return convex_hull