import numpy as np
from scipy.spatial.transform import Rotation as R
from scipy.linalg import svd
import os
import utils
import config

def rotation_matrix_to_rpy(R):
    # 确保输入是一个有效的旋转矩阵
    if not np.allclose(np.dot(R, R.T), np.eye(3)) or not np.isclose(np.linalg.det(R), 1.0):
        raise ValueError("输入不是一个有效的旋转矩阵")

    _pitch = -np.arcsin(R[2, 0])
    
    if np.abs(_pitch - np.pi / 2) < 1e-6:  # 处理接近90度的情况
        _yaw = 0
        _roll = np.arctan2(R[0, 1], R[0, 2])
    elif np.abs(_pitch + np.pi / 2) < 1e-6:  # 处理接近-90度的情况
        _yaw = 0
        _roll = -np.arctan2(R[0, 1], R[0, 2])
    else:
        _yaw = np.arctan2(R[1, 0], R[0, 0])
        _roll = np.arctan2(R[2, 1], R[2, 2])

    return _roll, _pitch, _yaw

def average_rotation_matrices(rotation_matrices, weights=None):
    # 将旋转矩阵转换为四元数
    quaternions = [R.from_matrix(matrix).as_quat() for matrix in rotation_matrices]
    
    if weights is None:
        # 如果没有提供权重，则使用等权重
        weights = np.ones(len(quaternions)) / len(quaternions)
    else:
        # 确保权重之和为1
        weights = np.array(weights) / np.sum(weights)
    
    # 对四元数进行加权平均
    avg_quaternion = np.zeros(4)
    for quat, weight in zip(quaternions, weights):
        avg_quaternion += weight * quat
    
    # 归一化四元数以确保它是单位四元数
    avg_quaternion /= np.linalg.norm(avg_quaternion)
    
    # 将平均后的四元数转换回旋转矩阵
    avg_rotation_matrix = R.from_quat(avg_quaternion).as_matrix()
    
    return avg_rotation_matrix

def addQ(q1, q2):
    return np.quaternion(q1.x + q2.x, q1.y + q2.y, q1.z + q2.z, q1.w + q2.w)

def calc_RT(lidar, camera, MAX_ITERS, lidarToCamera):
    global iteration_counter, pkg_loc, translation_sum, rotation_sum, rotation_avg_by_mult, rmse_avg
    
    if iteration_counter == 0:
        with open(pkg_loc + "/log/avg_values.txt", "w") as clean_file:
            pass
        
        translation_sum = np.zeros((3, 1), dtype=float)
        rotation_sum = np.ones(4)
        rotation_avg_by_mult = np.eye(3, dtype=float)
        rmse_avg = np.ones(3)
    
    num_points = lidar.shape[1]
    print("Number of points:", num_points)
    
    mu_lidar = np.mean(lidar, axis=0)
    mu_camera = np.mean(camera, axis=0)
    
    if iteration_counter == 0:
        print("mu_lidar:\n", mu_lidar)
        print("mu_camera:\n", mu_camera)
    
    # lidar_centered = lidar - mu_lidar[:, np.newaxis]
    # camera_centered = camera - mu_camera[:, np.newaxis]
    lidar_centered = lidar - mu_lidar
    camera_centered = camera - mu_camera
    
    if iteration_counter == 0:
        print("lidar_centered:\n", lidar_centered)
        print("camera_centered:\n", camera_centered)
    
    # cov = camera_centered @ lidar_centered.T ERROR
    H = np.dot(lidar_centered.T, camera_centered) 
    print(H)
    
    U, _, Vt = np.linalg.svd(H)

    # rotation = U @ Vt ERROR
    rotation = np.dot(Vt, U.T)
    
    if np.linalg.det(rotation) < 0:
        # rotation = U @ np.diag([1, 1, -1]) @ Vt
        rotation = np.dot(Vt, U.T)
    
    # translation = mu_camera - rotation @ mu_lidar 
    translation = mu_camera.reshape(3, 1) - np.dot(rotation, mu_lidar.reshape(3, 1))

    # averaging translation and rotation
    translation_sum += translation
    temp_q = R.from_matrix(rotation).as_quat()
    if iteration_counter == 0:
        rotation_sum = temp_q
    else:
        rotation_sum = rotation_sum + temp_q

    # averaging rotations by multiplication
    if iteration_counter == 0:
        rotation_avg_by_mult = rotation
    else:
        rotation_avg_by_mult = average_rotation_matrices([rotation_avg_by_mult, rotation])
    
    # ea = R.from_matrix(rotation).as_euler('zyx')
    ea = R.from_matrix(rotation).as_euler('xyz', degrees=False) 
    combine_ea = utils.combine_rpy(config.roll, config.pitch, config.yaw, ea[0], ea[1], ea[2])

    print("Rotation matrix:\n", rotation)
    print("Rotation in Euler :\n", ea)
    print("Combin Rotation in Euler :\n", combine_ea)
    print("Translation:\n", translation)
    
    # eltwise_error = np.sum((camera - (rotation @ lidar + translation[:, np.newaxis])) ** 2, axis=0) ERROR
    eltwise_error = np.sum((camera - (np.dot(lidar, rotation) + translation[:, np.newaxis])) ** 2, axis=0)
    error = np.sqrt(np.mean(eltwise_error))
    print("RMSE:", error)
    
    rmse_avg += error
    
    T = np.eye(4)
    T[:3, :3] = rotation
    T[:3, 3] = np.squeeze(translation)
    
    print("Rigid-body transformation:\n", T)
    
    iteration_counter += 1
    if iteration_counter % 1 == 0:
        with open(pkg_loc + "/log/avg_values.txt", "a") as log_avg_values:
            print("--------------------------------------------------------------------")
            print(f"After {iteration_counter} iterations")
            print("--------------------------------------------------------------------")
            
            print("Average translation is:\n", translation_sum / iteration_counter)
            log_avg_values.write(f"{iteration_counter}\n")
            log_avg_values.write(str(translation_sum / iteration_counter) + "\n")
            
            # rotation_sum_normalized = np.array([getattr(rotation_sum, dim) for dim in ['x', 'y', 'z', 'w']]) / iteration_counter
            rotation_sum_normalized = rotation_sum / iteration_counter
            mag = np.linalg.norm(rotation_sum_normalized)
            rotation_sum_normalized /= mag
            
            rotation_avg = R.from_quat(rotation_sum_normalized).as_matrix()
            print("Average rotation is:\n", rotation_avg)
            # final_rotation = rotation_avg @ lidarToCamera ERROR
            final_rotation = np.dot(lidarToCamera, rotation_avg)

            final_angles = R.from_matrix(final_rotation).as_euler('xyz')
            combine_final_angles = utils.combine_rpy(config.roll, config.pitch, config.yaw, \
                                                     final_angles[0], final_angles[1], final_angles[2])

            log_avg_values.write("{:.8f} {:.8f} {:.8f}\n".format(*rotation_avg[0]))
            log_avg_values.write("{:.8f} {:.8f} {:.8f}\n".format(*rotation_avg[1]))
            log_avg_values.write("{:.8f} {:.8f} {:.8f}\n".format(*rotation_avg[2]))
            
            T_avg = np.eye(4)
            T_avg[:3, :3] = rotation_avg
            T_avg[:3, 3] = np.squeeze(translation_sum / iteration_counter)
            print("Average transformation is:\n", T_avg)
            print("Final rotation is:\n", final_rotation)

            print("Final ypr is:\n", final_angles)
            print("Final Combine ypr is:\n", combine_final_angles)
            print(f"Average RMSE is: {rmse_avg / iteration_counter}")
            
            # eltwise_error_temp = np.sum((camera - (rotation_avg @ lidar + (translation_sum / iteration_counter)[:, np.newaxis])) ** 2, axis=0)
            eltwise_error_temp = np.sum((camera - (np.dot(lidar, rotation_avg) + (translation_sum / iteration_counter)[:, np.newaxis])) ** 2, axis=0)
            
            error_temp = np.sqrt(np.mean(eltwise_error_temp))
            print(f"RMSE on average transformation is: {error_temp}")
            log_avg_values.write("{:.8f}\n".format(error_temp))

    return T



# Global variables initialization
iteration_counter = 0
pkg_loc = "./record"
translation_sum = np.zeros((3, 1), dtype=float)
rotation_sum = np.ones(4)
rotation_avg_by_mult = np.eye(3, dtype=float)
rmse_avg = np.ones(3)

if __name__ == "__main__":
    all_points = utils.read_from_json(os.path.dirname(__file__) + '\\record\\points2.json')
    # records = utils.read_from_json(os.path.dirname(__file__) + '\\record\\record.json')
    carmera = np.asarray(all_points[0])
    lidar = np.asarray(all_points[1])
    lidarToCamera = R.from_euler('xyz', [config.roll, config.pitch, config.yaw], degrees=False).as_matrix()

    calc_RT(lidar, carmera, 50, lidarToCamera)