# # ls_ellipsoid and polyToParams3D functions are taken and adapted from
# http://www.juddzone.com/ALGORITHMS/least_squares_3D_ellipsoid.html
# # FindBestFittingEllipsoid3D and GetConvexHull3D 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

from brick.general.type import site_h, array_t
from brick.component.soma import soma_t


def ls_ellipsoid(xx: array_t, yy: array_t, zz: array_t) -> array_t:  # 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
    ellipsoid_coef = np_.append(ABC, -1)

    return (ellipsoid_coef)


def polyToParams3D(ellipsoid_coef: array_t, printMe: bool = False) -> tuple:
    """
    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', ellipsoid_coef)

    Amat = np_.array(
        [
            [ellipsoid_coef[0], ellipsoid_coef[3] / 2.0, ellipsoid_coef[4] / 2.0, ellipsoid_coef[6] / 2.0],
            [ellipsoid_coef[3] / 2.0, ellipsoid_coef[1], ellipsoid_coef[5] / 2.0, ellipsoid_coef[7] / 2.0],
            [ellipsoid_coef[4] / 2.0, ellipsoid_coef[5] / 2.0, ellipsoid_coef[2], ellipsoid_coef[8] / 2.0],
            [ellipsoid_coef[6] / 2.0, ellipsoid_coef[7] / 2.0, ellipsoid_coef[8] / 2.0, ellipsoid_coef[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 = ellipsoid_coef[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]
    s1 = -R[3, 3]
    R3S = R3 / s1

    (el, ec) = eig(R3S)
    el = np_.abs(el)

    # Sorting in descending order of the eigenvalues and eigenvectors
    sorting_vec = np_.argsort(-el)
    el = el[sorting_vec]
    ec = ec[:, sorting_vec]

    # Calculating cartesian axes
    recip = 1.0 / el
    axes = np_.sqrt(recip)  # axes are in ascending order
    if printMe:
        print('\nAxes are\n', axes, '\n')

    # Calculating the orientation
    inverse = inv(ec)  # inverse is actually the transpose here
    if printMe:
        print('\nRotation matrix\n', inverse)

    # Calculating spherical axes
    # # eigenvector 1
    r1 = np_.sqrt(sum(axe**2 for axe in ec[:, 0]))
    phi1 = np_.arctan2(ec[1, 0], ec[0, 0])
    theta1 = np_.arccos(ec[2, 0] / r1)
    spherical_axes1 = (r1, theta1, phi1)
    # # eigenvector 2
    r2 = np_.sqrt(sum(axe**2 for axe in ec[:, 1]))
    phi2 = np_.arctan2(ec[1, 1], ec[0, 1])
    theta2 = np_.arccos(ec[2, 1] / r2)
    spherical_axes2 = (r2, theta2, phi2)

    spherical_axes = (spherical_axes1, spherical_axes2)

    if printMe:
        print('\nAxes in spherical coord are\n', spherical_axes, '\n')

    return center, axes, inverse, spherical_axes


def GetConvexHull3D(soma_sites: site_h) -> tuple:
    #
    soma_3D = np_.stack((soma_sites[0], soma_sites[1], soma_sites[2]), axis=-1)
    hull_3D = ConvexHull(soma_3D)
    volume_hull = hull_3D.volume

    len_hull = len(hull_3D.vertices)
    hull = np_.zeros((len_hull, 3))

    for i in range(len(hull_3D.vertices)):
        hull[i] = soma_3D[hull_3D.vertices[i]]

    convex_hull = np_.transpose(hull)

    return convex_hull, volume_hull


def FindBestFittingEllipsoid3D(soma: soma_t) -> tuple:
    """
    Find the best fitting ellipsoid for the data points in 3D based on their convex hull.
    Return Tuple[ellipsoid coefficients, ellipsoid center, ellipsoid axes, ellipsoid orientation]
    """

    # get convex hull
    convex_hull, volume_hull = GetConvexHull3D(soma.sites)

    # 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, spherical_coor = polyToParams3D(ellipsoid_coef, False)

    return ellipsoid_coef, center, axes, orientation, spherical_coor, convex_hull, volume_hull