import numpy as np


def eigen(im_hess):
    """
    Calculates the eigenvectors of the hessian volumes 'H'
    generated by Hessian.py

    Parameters
    ----------
    im_hess : array_like
        M x N x 4 hessian matrix - H[:,:,0] = dxx,
        H[:,:,1] = H[:,:,2] = dxy, H[:,:,3] = dyy.

    Returns
    -------
    lambda_ : array_like
        M x N x 2 image of eigenvalues.
    v1 : array_like
        M x N x 2 eigenvector for lambda_(:,:,0)
    v2 :
        M x N x 2 eigenvector for lambda_(:,:,1)

    """
    # get size of H
    sizeX = im_hess.shape[0]
    sizeY = im_hess.shape[1]

    # initialize lambda_, v1 and v2
    lambda_ = np.zeros((sizeX, sizeY, 2))
    v1 = np.zeros((sizeX, sizeY, 2))
    v2 = np.zeros((sizeX, sizeY, 2))

    # compute eigenvalues of H
    radical = np.sqrt((im_hess[:, :, 0] - im_hess[:, :, 3]) ** 2 +
                      4 * im_hess[:, :, 1] ** 2)
    lambda_[:, :, 0] = (im_hess[:, :, 0] + im_hess[:, :, 3] + radical) / 2
    lambda_[:, :, 1] = (im_hess[:, :, 0] + im_hess[:, :, 3] - radical) / 2

    # compute eigenvectors of H
    v1[:, :, 0] = 2 * im_hess[:, :, 1]
    v1[:, :, 1] = im_hess[:, :, 3] - im_hess[:, :, 0] + radical
    norms = np.sqrt(v1[:, :, 0]**2 + v1[:, :, 1]**2)

    # normalize eigenvectors of H
    with np.errstate(divide='ignore', invalid='ignore'):
        v1[:, :, 0] = np.true_divide(v1[:, :, 0], norms)
        v1[:, :, 1] = np.true_divide(v1[:, :, 1], norms)
        # check -inf inf NaN
        v1[:, :, 0][~np.isfinite(v1[:, :, 0])] = 0
        v1[:, :, 1][~np.isfinite(v1[:, :, 1])] = 0
    v2[:, :, 0] = -v1[:, :, 1]
    v2[:, :, 1] = v1[:, :, 0]

    # order by magnitude
    swap = np.where(abs(lambda_[:, :, 0]) > abs(lambda_[:, :, 1]))

    # swap between lambda_[:, :, 0] and lambda_[:, :, 1]
    lambda_[:, :, 0][swap], lambda_[:, :, 1][swap] = \
        lambda_[:, :, 1][swap], lambda_[:, :, 0][swap]

    # swap between v1 and v2
    v1[:, :, 0][swap], v2[:, :, 0][swap] = \
        v2[:, :, 0][swap], v1[:, :, 0][swap]

    v1[:, :, 1][swap], v2[:, :, 1][swap] = \
        v2[:, :, 1][swap], v1[:, :, 1][swap]

    return lambda_, v1, v2
