import numpy as np
from scipy.spatial.transform import Rotation as R
import open3d as o3d
from scipy.spatial import KDTree
import tf.transformations as tft
import os
import rospy
from math import acos,tan,atan,degrees,sqrt

def corner_angle_discrimination(point,data,fangxiang):
    # 构建KDTree索引
    tree = KDTree(data)
    indices = np.array(tree.query_ball_point(point, r=30, p=2)) # 返回这些点的索引数组
    distances = np.array([np.linalg.norm(point - data[index]) for index in indices])
    distance_filter = distances >= 10
    filtered_indices = indices[distance_filter]
    data_circular = data[filtered_indices]

    # 可视化
    # pt_cloud_circular = o3d.geometry.PointCloud()
    # pt_cloud_circular.points = o3d.utility.Vector3dVector(data_circular)
    # o3d.visualization.draw_geometries([pt_cloud_circular])

    # 获取与方向向量相似的点云
    data_similar = []
    for i in range(len(data_circular)):
        v = np.array(data_circular[i]) - point
        angle = acos(np.dot(v, fangxiang) / (np.linalg.norm(v) * np.linalg.norm(fangxiang)))
        angle_deg = degrees(angle)
        if angle_deg <= 15:
            data_similar.append(data_circular[i])
    
    if len(data_similar) >= 1:
        corner = True
    else:
        corner = False
    return corner
def set_yaw_angle(R_mat,yaw):
    pz1 = np.array([[1, 0, 0],
                    [0, np.cos(yaw), -np.sin(yaw)],
                    [0, np.sin(yaw),  np.cos(yaw)]])
    Z_PI = np.array([[np.cos(np.pi), -np.sin(np.pi), 0],
                    [np.sin(np.pi), np.cos(np.pi), 0],
                    [0, 0, 1]])
    R1_mat = np.matmul(R_mat,pz1)
    
    sign = check_x_axis_direction(R1_mat)
    if sign == 1:
        pass
    elif sign == -1:
        R1_mat = np.matmul(R1_mat, Z_PI)

    q = R.from_matrix(R1_mat).as_quat()
    q = q.tolist()  # 将NumPy数组转换为列表
    return q
def calculate_angle_with_xy_plane(point1,point2):
    # 计算方向向量
    dir_vector = np.array([point2[0] - point1[0], point2[1] - point1[1], point2[2] - point1[2]])
    
    # 计算方向向量在 xy 平面上的投影
    proj_vector = np.array([point2[0] - point1[0], point2[1] - point1[1], 0])
    
    # 计算夹角
    angle = acos(np.dot(dir_vector, proj_vector) / (np.linalg.norm(dir_vector) * np.linalg.norm(proj_vector)))
    angle_deg = degrees(angle)

    return angle_deg
def check_x_axis_direction(rotation_matrix):
    # 提取旋转后的 x 轴方向向量
    x_prime = rotation_matrix[:, 0]
    # 检查 x 轴的 z 分量
    if x_prime[2] > 0:
        return 1
    elif x_prime[2] < 0:
        return -1
    else:
        return 0
def load_point_cloud_from_txt(file_path):
    """导入点云数据"""
    points = []
    with open(file_path, 'r') as file:
        for line in file:
            # 假设每一行是“x y z”格式，没有多余空格
            coordinates = line.strip().split()
            if len(coordinates) == 3:  # 检查是否为有效的三维点
                point = [float(coord) for coord in coordinates]
                points.append(point)

    return np.array(points)
def load_point_cloud_from_binary_txt(file_path):
    with open(file_path, 'rb') as f:
        binary_data = f.read()

        # 将二进制数据转换为 NumPy 数组
        point_cloud = np.frombuffer(binary_data, dtype=np.float64).reshape(-1, 3)
    return point_cloud

def get_hanfeng(file_path):
    """从文件中读取焊缝起点、终点和中间点，并计算两点构成的向量"""
    with open(file_path, 'r') as file:
        # 逐行读取文件内容
        lines = file.readlines()
    start_points, end_points, midpoints , hanfeng = [],[],[],[]
    for line in lines:
        # 去除行尾的换行符并按'/'分割每一行
        points_str = line.strip().split('/')
        
        # 确保每一行都正确地分为两部分
        if len(points_str) == 2:
            point1_str = points_str[0].split(',')
            point2_str = points_str[1].split(',')
            
            # 转换字符串为浮点数列表，构造三维点
            point1 = [float(coord) for coord in point1_str]
            point2 = [float(coord) for coord in point2_str]

            midpoint = [(p2 + p1)/2 for p1, p2 in zip(point1, point2)]
            vector = [p2 - p1 for p1, p2 in zip(point1, point2)]
            
            start_points.append(point1)
            end_points.append(point2)
            midpoints.append(midpoint)  
            hanfeng.append(vector)   

    return start_points, end_points, midpoints, hanfeng

def compute_Rc2w_tc2w(file_path_Tw2c):
    """相机和基底坐标系的转换关系"""
    # 读取文件
    with open(file_path_Tw2c, 'r') as file:
        lines = file.readlines()

    lines = [line.strip().replace('[', '').replace(']', '').split() for line in lines]
    data = [[float(num) for num in line] for line in lines]
    data = np.array(data)
    # 提取Rw2c和tw2c
    Rw2c = data[:3,:3]
    tw2c = data[:3,-1]

    Rc2w = Rw2c.T
    tc2w = -np.matmul(Rc2w, tw2c)
    return Rc2w,tc2w

def compute_pose_R(fangxiang, hanfeng,start_point, end_point):
    """根据计算出的焊接时z轴负方向的向量和由起(fangxiang)点指向终点的向量(hanfeng)计算旋转矩阵"""
    x, y, z = fangxiang

    # 计算旋转角度
    rotx = -np.arctan2(y, z)
    roty = np.arctan2(x, np.sqrt(y**2 + z**2))

    # 构造旋转矩阵
    A = np.array([[1, 0, 0],
                [0, np.cos(rotx), -np.sin(rotx)],
                [0, np.sin(rotx), np.cos(rotx)]])
    B = np.array([[np.cos(roty), 0, np.sin(roty)],
                [0, 1, 0],
                [-np.sin(roty), 0, np.cos(roty)]])
    bb = np.matmul(hanfeng, np.matmul(A, B))
    # 区分平缝竖缝
    angle = calculate_angle_with_xy_plane(start_point,end_point)
    ispingfeng = abs(angle) < 30
    if ispingfeng: 
        rotz = np.arctan2(bb[0], -bb[1])
    else:
        rotz = np.arctan2(bb[1], bb[0])

    C = np.array([[np.cos(rotz), -np.sin(rotz), 0],
                [np.sin(rotz), np.cos(rotz), 0],
                [0, 0, 1]])
    D = np.array([[1, 0, 0],
                    [0, np.cos(np.pi), -np.sin(np.pi)],
                    [0, np.sin(np.pi), np.cos(np.pi)]])

    R_mat = np.matmul(A, np.matmul(B, np.matmul(C, D)))

    return R_mat, ispingfeng

def get_quaternion(data, midpoint, hanfeng, start_point, end_point,result_ping,result_shu,result_wai,file_path_Tw2c): 

    # 构建KDTree索引
    tree = KDTree(data)
    # 查询半径范围内的点,找出"data"中所有距离"midpoint"点在r以内的点，这里使用的是欧几里得距离（由p=2表示）
    indices = np.array(tree.query_ball_point(midpoint, r=30, p=2)) # 返回这些点的索引数组
    distances = np.array([np.linalg.norm(midpoint - data[index]) for index in indices])
    # 对numpy数组进行布尔索引
    # distance_filter = np.logical_and(distances >= 10, distances <= 50)
    distance_filter = distances >= 10
    # 使用布尔索引过滤出符合条件的索引
    filtered_indices = indices[distance_filter]
    # 提取符合条件的点
    data_circular = data[filtered_indices]

    # 可视化
    # pt_cloud_circular = o3d.geometry.PointCloud()
    # pt_cloud_circular.points = o3d.utility.Vector3dVector(data_circular)
    # o3d.visualization.draw_geometries([pt_cloud_circular])

    # 获取竖直方向上的点云
    data_shu = []
    for i in range(len(data_circular)):
        P = np.array(data_circular[i])
        v = P - midpoint
        anglecos = np.dot(v, hanfeng) / (np.linalg.norm(v) * np.linalg.norm(hanfeng))
        
        if abs(anglecos) <= 0.015:
            # print(anglecos)
            data_shu.append(P)
    data_shu = np.array(data_shu)
    # # 可视化
    # pt_cloud_shu = o3d.geometry.PointCloud()
    # pt_cloud_shu.points = o3d.utility.Vector3dVector(np.array(data_shu))
    # o3d.visualization.draw_geometries([pt_cloud_shu])

 ##################################################
 ##################################################
    random_point = data_shu[np.random.randint(0, len(data_shu))]
    # 与midpoint一起定义直线向量
    nohanfeng = random_point - midpoint  #data_shu中任意一点midpoint构成的向量，计算所有点和该向量的距离判断是否是外缝(只拍到一面)

    distances_to_line = []

    for i in range(data_shu.shape[0]):  # 遍历每一个点 P
        P = data_shu[i]
        A = midpoint
        u = nohanfeng
        v = P - A
        distance = np.linalg.norm(np.cross(u, v)) / np.linalg.norm(u)
        distances_to_line.append(distance)

    distances_to_line = np.array(distances_to_line)

    all_distances_less_than_10 = np.all(distances_to_line < 5)#都小于10为true，也就是data_shu只有一列，侧着拍得外缝

    if all_distances_less_than_10:

        # 随机选取data_shu中的第一个点
        random_index = np.random.randint(0, data_shu.shape[0])
        vector1 = data_shu[random_index] - midpoint
        vector1 /= np.linalg.norm(vector1)  # 单位化vector1
        fa_xian = np.cross(vector1, hanfeng)
        fa_xian /= np.linalg.norm(fa_xian) 

        point = []
        point1 = midpoint + 50 * fa_xian
        point2 = midpoint - 50 * fa_xian

        Rc2w,tc2w = compute_Rc2w_tc2w(file_path_Tw2c)

        point1_c = np.matmul(Rc2w, point1) + tc2w
        point2_c = np.matmul(Rc2w, point2) + tc2w
        a = np.linalg.norm(point1_c)
        b = np.linalg.norm(point2_c)      

        if a < b:
            point = point1
        else:
            point = point2

        faxian = point - midpoint
        vector2 =  faxian / np.linalg.norm(faxian)  # 单位化vector2

        fangxiang = vector1  + vector2*1.2

        R_mat, ispingfeng = compute_pose_R(fangxiang, hanfeng, start_point, end_point)
        if ispingfeng:
            ##偏转##
            pz1 = np.array([[1, 0, 0],
                        [0, np.cos(yaw), -np.sin(yaw)],
                        [0, np.sin(yaw),  np.cos(yaw)]])
            pz2 = np.array([[1, 0, 0],
                        [0, np.cos(-yaw), -np.sin(-yaw)],
                        [0, np.sin(-yaw),  np.cos(-yaw)]])
            Z_PI = np.array([[np.cos(np.pi), -np.sin(np.pi), 0],
                    [np.sin(np.pi), np.cos(np.pi), 0],
                    [0, 0, 1]])
            R1_mat = np.matmul(R_mat,pz1)
            R2_mat = np.matmul(R_mat,pz2)

            sign = check_x_axis_direction(R1_mat)
            if sign == 1:
                pass
            elif sign == -1:
                R1_mat = np.matmul(R1_mat, Z_PI)
                R2_mat = np.matmul(R2_mat, Z_PI)

            q1 = R.from_matrix(R1_mat).as_quat()
            q2 = R.from_matrix(R2_mat).as_quat()
            q1 = q1.tolist()  # 将NumPy数组转换为列表
            q2 = q2.tolist()
            # print("一个点云面的平外缝")
            # print(q1,q2)
            result_ping.append(("pw", start_point, end_point, q1, q2))
            
            return result_ping,result_shu,result_wai
        else:
            ##偏转##
            E = np.array([[np.cos(-pitch_of_Verticalweld), 0, np.sin(-pitch_of_Verticalweld)],
                        [0, 1, 0],
                        [-np.sin(-pitch_of_Verticalweld), 0, np.cos(-pitch_of_Verticalweld)]])

            R_mat = np.matmul(R_mat,E)
            
            q1 = R.from_matrix(R_mat).as_quat()
            q2 = R.from_matrix(R_mat).as_quat()
            q1 = q1.tolist()  # 将NumPy数组转换为列表
            q2 = q2.tolist()
            # print("即一个点云面的竖外缝")
            # print(q1,q2)
            result_wai.append(("sw", start_point, end_point, q1, q2))
            return result_ping,result_shu,result_wai

    else:
        # 选择位于平面两侧的点
        # 随机选取data_shu中的第一个点
        random_index = np.random.randint(0, data_shu.shape[0])
        vector1 = data_shu[random_index] - midpoint
        vector1 /= np.linalg.norm(vector1)  # 单位化vector1

        # 随机选取下一个点，直到找到与vector1夹角大于10度的向量
        found = False
        while not found:
            random_index = np.random.randint(0, data_shu.shape[0])
            random_point = data_shu[random_index] - midpoint
            vector2 = random_point / np.linalg.norm(random_point)  # 单位化vector2
            
            # 计算两个向量之间的夹角（弧度）
            dot_product = np.dot(vector1, vector2)
            angle_rad = np.arccos(dot_product)
            
            # 判断角度是否大于10度（转换为弧度）
            if angle_rad > np.deg2rad(20) and angle_rad < np.deg2rad(180):
                found = True
                # 如果找到了，则保存vector2
            else:
                continue

        # 计算并保存最终的fangxiang向量, vector1与vector2均为单位向量
        vector1_sim_to_z = np.dot(vector1, [0,0,1]) / (np.linalg.norm(vector1))
        vector2_sim_to_z = np.dot(vector2, [0,0,1]) / (np.linalg.norm(vector2))
        if vector1_sim_to_z > vector2_sim_to_z:
            vector_shu = vector1
            vector_ping = vector2
        else:
            vector_shu = vector2
            vector_ping = vector1
        fangxiang = vector_ping + vector_shu*tan(pitch_of_Horizontalweld)
        # fangxiang = vector1 + vector2
        fangxiang /= np.linalg.norm(fangxiang)  # 单位化fangxiang

        #########################################根据vector1和vector2判断内缝外缝#########################################
        point1 = midpoint + 50 * fangxiang 
        point2 = midpoint - 50 * fangxiang 

        Rc2w,tc2w = compute_Rc2w_tc2w(file_path_Tw2c)
        
        point1_c = np.matmul(Rc2w, point1) + tc2w
        point2_c = np.matmul(Rc2w, point2) + tc2w
        a = np.linalg.norm(point1_c)
        b = np.linalg.norm(point2_c)    

        if a < b: # 内缝
            # print("两个点云面的内缝")
            fangxiang = fangxiang
        else: # 外缝
            # print("两个点云面的外缝")
            fangxiang = -fangxiang

        R_mat, ispingfeng = compute_pose_R(fangxiang, hanfeng, start_point, end_point)
        if ispingfeng:
            ##偏转##
            corner_of_start_point = corner_angle_discrimination(start_point,data,fangxiang)
            corner_of_end_point = corner_angle_discrimination(end_point,data,fangxiang)
            pi4_pz = atan(1 / (sqrt(2) * tan(np.pi/4)))#CL公式
            yaw_pz =atan(1 / (sqrt(2) * tan(yaw)))
            if corner_of_start_point:
                q1 = set_yaw_angle(R_mat,pi4_pz)
                if not corner_of_end_point:
                    len_hanfeng = np.linalg.norm(hanfeng)
                    angle = pi4_pz-(pi4_pz)*len_hanfeng/yaw_rate
                    if angle < -yaw_pz:
                        angle =-yaw_pz
                    q2 = set_yaw_angle(R_mat,angle)
            if corner_of_end_point:
                q2 = set_yaw_angle(R_mat,-pi4_pz)
                if not corner_of_start_point:
                    len_hanfeng = np.linalg.norm(hanfeng)
                    angle = -pi4_pz+(pi4_pz)*len_hanfeng/yaw_rate
                    if angle > yaw_pz:
                        angle = yaw_pz
                    q1 = set_yaw_angle(R_mat,angle)
            if not corner_of_start_point and not corner_of_end_point:     
                q1 = set_yaw_angle(R_mat,yaw_pz)
                len_hanfeng = np.linalg.norm(hanfeng)
                angle = yaw_pz-(pi4_pz)*len_hanfeng/yaw_rate
                if angle < -yaw_pz:
                    angle =-yaw_pz
                q2 = set_yaw_angle(R_mat,angle)
            
            result_ping.append(("p", start_point, end_point, q1, q2))
            
            return result_ping,result_shu,result_wai
        else:
            ##偏转##
            E = np.array([[np.cos(-pitch_of_Verticalweld), 0, np.sin(-pitch_of_Verticalweld)],
                        [0, 1, 0],
                        [-np.sin(-pitch_of_Verticalweld), 0, np.cos(-pitch_of_Verticalweld)]])

            R_mat = np.matmul(R_mat,E)
            
            q1 = R.from_matrix(R_mat).as_quat()
            q2 = R.from_matrix(R_mat).as_quat()
            q1 = q1.tolist()  # 将NumPy数组转换为列表
            q2 = q2.tolist()
            # print(q1,q2)
            result_wai.append(("s", start_point, end_point, q1, q2))
            return result_ping,result_shu,result_wai        

def run():
    global yaw,yaw_rate,pitch_of_Horizontalweld,pitch_of_Verticalweld
    # yaw = np.pi/8
    yaw = np.pi/2-rospy.get_param("yaw")
    pitch_of_Verticalweld = rospy.get_param("pitch_of_Verticalweld")
    pitch_of_Horizontalweld = rospy.get_param("pitch_of_Horizontalweld")
    yaw_rate = rospy.get_param("yaw_rate")
    file_path = rospy.get_param("folder_path")
    # file_path = ('/home/chen/catkin_ws/data/0603')
    file_path_pointcloud = os.path.join(file_path, 'pointcloud.txt')
    file_path_points = os.path.join(file_path, 'points_guihua.txt')
    file_path_result = os.path.join(file_path, 'result.txt')
    file_path_Tw2c = os.path.join(file_path, 'Tw2c.txt')

    # 获取数据
    data = load_point_cloud_from_binary_txt(file_path_pointcloud)
    start_points, end_points, midpoints, hanfeng = get_hanfeng(file_path_points)

    # 计算位姿
    result_ping,result_shu,result_wai = [],[],[]
    for i in range(len(midpoints)):
        result_ping,result_shu,result_wai = get_quaternion(data, midpoints[i], hanfeng[i], start_points[i], end_points[i],result_ping,result_shu,result_wai, file_path_Tw2c)
        # print(q1," ",q2)
    result = result_ping + result_shu + result_wai

    # 打开文件准备写入
    with open(file_path_result, "w") as file:
        for group in result:
            # 对于每个数据组，将每个部分用逗号连接，然后在各部分间插入斜杠
            line = "/".join([",".join(map(str, sublist)) for sublist in group]) + "\n"
            file.write(line)

    rospy.loginfo("Welding Angle calculation completed")
if __name__ == '__main__':     
    
    run()