# -- coding: utf-8 --
import numpy as np
from scipy.linalg import svd
from scipy.spatial import KDTree
 

def best_fit_transform(A, B):
    """
    Calculates the least-squares best-fit transform that maps corresponding points A to B in m spatial dimensions.
    
    Input:
      A: Nxm numpy array of corresponding points
      B: Nxm numpy array of corresponding points
      
    Returns:
      R: mxm rotation matrix
      t: mx1 translation vector
    """
    assert A.shape == B.shape

    # get number of dimensions
    m = A.shape[1]

    # translate points to their centroids
    centroid_A = np.mean(A, axis=0)
    centroid_B = np.mean(B, axis=0)
    AA = A - centroid_A
    BB = B - centroid_B

    # rotation matrix
    H = np.dot(AA.T, BB)
    # U, S, Vt = np.linalg.svd(H)
    U, S, Vt = svd(H)
    R = np.dot(Vt.T, U.T)

    # special reflection case
    if np.linalg.det(R) < 0:
       Vt[m-1,:] *= -1
       R = np.dot(Vt.T, U.T)

    # translation
    t = centroid_B.T - np.dot(R,centroid_A.T)

    return R, t


def nearest_neighbor(src, dst):
    """
    Find the nearest (Euclidean) neighbor in dst for each point in src
    
    Input:
        src: Nxm array of points
        dst: Nxm array of points
        
    Output:
        distances: Euclidean distances of the nearest neighbor
        indices: dst indices of the nearest neighbor
    """

    all_dists = np.array([np.linalg.norm(src[i] - dst, axis=1) for i in range(src.shape[0])])
    indices = all_dists.argmin(axis=1)
    distances = all_dists[np.arange(all_dists.shape[0]), indices]
    return distances, indices


def icp(A, B, init_pose=None, max_iterations=1, tolerance=0.0002):
    """
    The Iterative Closest Point method: finds best-fit transform that maps points A on to points B
    
    Input:
        A: Nxm numpy array of source mD points
        B: Nxm numpy array of destination mD point
        init_pose: (m+1)x(m+1) homogeneous transformation
        max_iterations: exit algorithm after max_iterations
        tolerance: convergence criteria
        
    Output:
        T: final homogeneous transformation that maps A on to B
        distances: Euclidean distances (errors) of the nearest neighbor
        i: number of iterations to converge
    """

    assert A.shape == B.shape

    # get number of dimensions
    m = A.shape[1]

    # make points homogeneous, copy them so as to maintain the originals
    src = np.ones((m+1,A.shape[0]))
    dst = np.ones((m+1,B.shape[0]))
    src[:m,:] = np.copy(A.T)
    dst[:m,:] = np.copy(B.T)

    # apply the initial pose estimation
    if init_pose is not None:
        src = np.dot(init_pose, src)

    prev_error = 0

    for i in range(max_iterations):
        # find the nearest neighbors between the current source and destination points
        distances, indices = nearest_neighbor(src[:m,:].T, dst[:m,:].T)

        # compute the transformation between the current source and nearest destination points
        R, t = best_fit_transform(src[:m,:].T, dst[:m,indices].T)

        # update the current source
        src[:m,:] = np.dot(R, src[:m,:]) + t[:, np.newaxis]

        # check error
        mean_error = np.mean(distances)
        if abs(prev_error - mean_error) < tolerance:
            break
        prev_error = mean_error

    # calculate final transformation
    T = np.eye(m+1)
    T[:m, :m] = R
    T[:m, m] = t

    return R, t, distances, i


# def scipy_svd(A, B):
#     """
#     Input:
#     A: Nx3 points in matrix form in source space
#     B: Nx3 points in matrix form in target space
#     """
#     N = A.shape[0]
#     # Mean of source and target points
#     centroid_A = np.mean(A, axis=0)
#     centroid_B = np.mean(B, axis=0)
#     # Center points
#     A_centered = A - centroid_A
#     B_centered = B - centroid_B
#     # Dot product
#     H = np.dot(A_centered.T, B_centered)
#     # Singular Value Decomposition
#     # U, S, Vt = svd(H, full_matrices=False, lapack_driver='gesdd')
#     # U, S, Vt = svd(H)
#     U, S, Vt = np.linalg.svd(H)
#     # V is Vt transposed
#     V = Vt.T
#     # R, t
#     R = np.dot(V, U.T)
#     if np.linalg.det(R) < 0:
#         # Special case to correct for reflection
#         V[:, -1] *= -1
#         R = np.dot(V, U.T)
#     t = centroid_B.reshape(3, 1) - np.dot(R, centroid_A.reshape(3, 1))
    
#     return R, t.ravel(), None, None

# def scipy_svd(A, B):
#     """
#     Input:
#     A: Nx3 points in matrix form in source space
#     B: Nx3 points in matrix form in target space
#     """
#     N = A.shape[0]
#     # Mean of source and target points
#     centroid_A = np.mean(A, axis=0)
#     centroid_B = np.mean(B, axis=0)
#     # Center points
#     A_centered = A - centroid_A
#     B_centered = B - centroid_B
#     # Dot product
#     H = np.dot(A_centered.T, B_centered)
#     # Singular Value Decomposition
#     U, S, Vt = svd(H)
#     # V is Vt transposed
#     V = Vt.T
#     # R, t
#     R = np.dot(V, U.T)
#     if np.linalg.det(R) < 0:
#         # Special case to correct for reflection
#         V[:, -1] *= -1
#         R = np.dot(V, U.T)
#     t = centroid_B.reshape(3, 1) - np.dot(R, centroid_A.reshape(3, 1))
    
#     return R, t.ravel(), None, None

def scipy_svd(A, B):
    """
    Input:
    A: Nx3 points in matrix form in source space
    B: Nx3 points in matrix form in target space
    
    Output:
    R: Rotation matrix
    t: Translation vector
    """
    N = A.shape[0]
    
    # Mean of source and target points
    centroid_A = np.mean(A, axis=0)
    centroid_B = np.mean(B, axis=0)
    
    # Center points
    A_centered = A - centroid_A
    B_centered = B - centroid_B
    
    # Dot product
    H = np.dot(A_centered.T, B_centered)
    
    # Singular Value Decomposition
    U, S, Vt = svd(H)
    
    # Calculate rotation matrix
    R = np.dot(Vt.T, U.T)
    
    # Check for reflection and correct if necessary
    if np.linalg.det(R) < 0:
        Vt[-1, :] *= -1  # Correct the reflection by flipping the sign of the last row of Vt
        R = np.dot(Vt.T, U.T)
    
    # Calculate translation vector
    t = centroid_B - np.dot(R, centroid_A)
    
    return R, t, None, None

def np_svd(A, B):
    """
    Input:
    A: Nx3 points in matrix form in source space
    B: Nx3 points in matrix form in target space
    """
    N = A.shape[0]
    # Mean of source and target points
    centroid_A = np.mean(A, axis=0)
    centroid_B = np.mean(B, axis=0)
    # Center points
    A_centered = A - centroid_A
    B_centered = B - centroid_B
    # Dot product
    H = np.dot(A_centered.T, B_centered)
    # Singular Value Decomposition
    U, S, Vt = np.linalg.svd(H)
    # V is Vt transposed
    V = Vt.T
    # R, t
    R = np.dot(V, U.T)
    if np.linalg.det(R) < 0:
        # Special case to correct for reflection
        V[:, -1] *= -1
        R = np.dot(V, U.T)
    t = centroid_B.reshape(3, 1) - np.dot(R, centroid_A.reshape(3, 1))
    
    return R, t.ravel(), None, None


def icp_svd(source_points, target_points, max_iterations=1, tolerance=1e-6):
    """
    使用 SVD 方法执行 ICP 算法来对齐两个点云。
    
    参数:
    source_points (numpy.ndarray): 源点云，形状为 (N, 3)
    target_points (numpy.ndarray): 目标点云，形状为 (M, 3)
    max_iterations (int): 最大迭代次数，默认为 100
    tolerance (float): 收敛误差阈值，默认为 1e-6
    
    返回:
    aligned_source (numpy.ndarray): 变换后的源点云
    R (numpy.ndarray): 旋转矩阵
    t (numpy.ndarray): 平移向量
    """
    
    # 初始化误差
    prev_error = float('inf')
    source = source_points.copy()
    
    for iteration in range(max_iterations):
        # 寻找对应点
        tree = KDTree(target_points)
        distances, indices = tree.query(source)
        
        # 计算质心
        source_center = np.mean(source, axis=0)
        target_center = np.mean(target_points[indices], axis=0)
        
        # 中心化点云
        source_centered = source - source_center
        target_centered = target_points[indices] - target_center
        
        # 构建协方差矩阵并执行 SVD 分解
        H = source_centered.T @ target_centered
        # U, _, Vt = np.linalg.svd(H)
        U, _, Vt = svd(H)
        
        # 计算旋转矩阵 R
        R = Vt.T @ U.T
        
        # 检查是否有反射并修正
        if np.linalg.det(R) < 0:
            Vt[-1, :] *= -1
            R = Vt.T @ U.T
        
        # 计算平移向量 t
        t = target_center - R @ source_center
        
        # 应用变换到源点云
        source = (R @ source_centered.T).T + t
        
        # 计算误差
        current_error = np.mean(distances)
        print(f"Iteration {iteration}, Error: {current_error}")
        
        # 检查是否满足终止条件
        if abs(prev_error - current_error) < tolerance:
            break
        
        prev_error = current_error
    
    return R, t, source, None

if __name__ == "__main__":
    source_points = np.array([[0.0, 0.0, 0.0],
                            [1.0, 0.0, 0.0],
                            [0.0, 1.0, 0.0]], dtype=np.float64)

    target_points = np.array([[1.0, 0.0, 0.0],
                            [2.0, 0.0, 0.0],
                            [1.0, 1.0, 0.0]], dtype=np.float64)

    aligned_source, R, t = icp_svd(source_points, target_points)

    print("Aligned Source Points:\n", aligned_source)
    print("Rotation Matrix:\n", R)
    print("Translation Vector:\n", t)

if __name__ == "__main__":
    # Test with random data
    np.random.seed(0)
    num_points = 7
    dim = 3
    noise_sigma = .01
    outlier_ratio = 0.1

    A = np.random.rand(num_points, dim)
    T = np.eye(dim+1)
    angle = np.pi / 4
    T[:dim, :dim] = [[np.cos(angle), -np.sin(angle)],
                     [np.sin(angle), np.cos(angle)]]
    T[:-1, -1] = np.array([1, 2])
    B = np.dot(T, np.hstack((A.T, np.ones((1, num_points)))))
    B = B[:dim, :].T

    # add Gaussian noise
    B += np.random.randn(num_points, dim) * noise_sigma

    # add some outliers
    n_outliers = int(outlier_ratio * float(num_points))
    indices = np.arange(num_points)
    np.random.shuffle(indices)
    B[indices[-n_outliers:], :] += 50 * np.random.randn(n_outliers, dim)

    # run ICP
    R, t, distances, i = icp(A, B)

    print("Final Transformation Matrix:\n", T)
    print("Mean Distance:", np.mean(distances))
    print("Iterations:", i)