import os
import time
import trimesh
import numpy as np

tol_norm = 1e-4


def test_align(align):
    """
    Test an "align two vectors" function
    """
    is_rigid = trimesh.transformations.is_rigid

    # start with some edge cases and make sure the transform works
    target = np.array([0, 0, -1], dtype=np.float64)
    vectors = np.vstack((
        trimesh.unitize(np.random.random((1000, 3)) - .5),
        np.random.random((1000, 3)) - .5,
        [-target, target],
        trimesh.util.generate_basis(target),
        [[7.12106798e-07, -7.43194705e-08, 1.00000000e+00],
         [0, 0, -1],
         [1e-4, 1e-4, -1]]))
    norms = []

    unitized = trimesh.unitize(vectors)
    for unit_dest, dest in zip(unitized[-10:], vectors[-10:]):
        for unit, vector in zip(unitized, vectors):
            T, a = align(vector, dest, return_angle=True)
            if not is_rigid(T):
                print('ERR: NOT RIGID!!\n', unit, dest)
            # assert is_rigid(T)
            assert np.isclose(np.linalg.det(T), 1.0)
            # rotate vector with transform
            check = np.dot(T[:3, :3], unit)
            # compare to target vector
            norm = np.linalg.norm(check - unit_dest)
            norms.append(norm)
            assert norm < tol_norm

    norms = np.array(norms)
    print('err.ptp: {}\nerr.std: {}\nerr.mean: {}\nerr.median: {}'.format(
        norms.ptp(), norms.std(), norms.mean(), np.median(norms)))

    # these vectors should be perpendicular and zero
    angles = [align(i, target, return_angle=True)[1]
              for i in trimesh.util.generate_basis(target)]
    assert np.allclose(
        angles, [np.pi / 2, np.pi / 2, 0.0])

    # generate angles from 0 to 180 degrees
    angles = np.linspace(0.0, np.pi / 1e7, 10000)
    # generate on- plane vectors
    vectors = np.column_stack((np.cos(angles),
                               np.sin(angles),
                               np.zeros(len(angles))))

    T = align([0, 0, -1], [-1e-17, 1e-17, 1])
    assert np.isclose(np.linalg.det(T), 1.0)

    T = align([0, 0, -1], [-1e-4, 1e-4, 1])
    assert np.isclose(np.linalg.det(T), 1.0)

    vector_1 = np.array([7.12106798e-07, -7.43194705e-08, 1.00000000e+00])
    vector_2 = np.array([0, 0, -1])
    T, angle = align(vector_1, vector_2, return_angle=True)
    assert np.isclose(np.linalg.det(T), 1.0)


def generate_basis(vector):
    vector = np.array(vector)

    u = np.linalg.svd(vector[:, np.newaxis])[0]
    if np.linalg.det(u) < 0.0:
        u[:, -1] *= -1
    return u


def align_svd(a, b, return_angle=False):
    """
    Find the rotation matrix that transforms one 3D vector
    to another.

    Parameters
    ------------
    a : (3,) float
      Unit vector
    b : (3,) float
      Unit vector
    return_angle : bool
      Return the angle between vectors or not

    Returns
    -------------
    matrix : (4, 4) float
      Homogenous transform to rotate from `a` to `b`
    angle : float
      If `return_angle` angle in radians between `a` and `b`

    """
    a = np.array(a, dtype=np.float64)
    b = np.array(b, dtype=np.float64)
    if a.shape != (3,) or b.shape != (3,):
        raise ValueError('vectors must be (3,)!')

    # find the SVD of the two vectors
    au = np.linalg.svd(a.reshape((-1, 1)))[0]
    bu = np.linalg.svd(b.reshape((-1, 1)))[0]

    if np.linalg.det(au) < 0:
        au[:, -1] *= -1.0
    if np.linalg.det(bu) < 0:
        bu[:, -1] *= -1.0

    # put rotation into homogenous transformation
    matrix = np.eye(4)
    matrix[:3, :3] = bu.dot(au.T)

    if return_angle:
        # projection of a onto b
        # first row of SVD result is normalized source vector
        dot = np.dot(au[0], bu[0])

        # clip to avoid floating point error
        angle = np.arccos(np.clip(dot, -1.0, 1.0))
        if dot < -1e-5:
            angle += np.pi
        return matrix, angle

    return matrix


if __name__ == '__main__':
    trimesh.util.attach_to_log()

    print('\n\nalign_svd\n-------------------')
    tic = [time.time()]
    test_align(align_svd)
    tic.append(time.time())
    print('elapsed: {}s'.format(tic[1] - tic[0]))

    print('\n\ntrimesh.geometry.align_vectors\n---------------------------------------')

    tic.append(time.time())
    test_align(trimesh.geometry.align_vectors)
    tic.append(time.time())

    print('elapsed: {}s'.format(tic[3] - tic[2]))