"""计算工具"""
import logging

import cv2
import numpy as np
import open3d as o3d
import transforms3d as tfs

from .object_icp import object_icp

__all__ = [
    "calculate_pose_pca", "calculate_pose_icp", "calculate_grav_width",
    "rmat2quat", "quat2rmat",
]

def pca(data: np.ndarray, sort=True):
    """主成分分析 """
    center = np.mean(data, axis=0)
    centered_points = data - center    # 去中心化

    try:
        cov_matrix = np.cov(centered_points.T)  # 转置
        eigenvalues, eigenvectors = np.linalg.eigh(cov_matrix)

    except np.linalg.LinAlgError:
        return None, None

    if sort:
        sort = eigenvalues.argsort()[::-1]  # 降序排列
        eigenvalues = eigenvalues[sort]  # 索引
        eigenvectors = eigenvectors[:, sort]

    return eigenvalues, eigenvectors


def calculate_pose_pca(
        mask,
        depth_img: np.ndarray,
        intrinsics,
        **kwargs
):
    """点云主成分分析法计算位态"""
    depth_img_mask = np.zeros_like(depth_img)
    depth_img_mask[mask > 0] = depth_img[mask > 0]

    depth_o3d = o3d.geometry.Image(depth_img_mask.astype(np.uint16))

    point_cloud = o3d.geometry.PointCloud.create_from_depth_image(
        depth=depth_o3d,
        intrinsic=intrinsics,
        depth_scale=kwargs.get("depth_scale", 1000.0),
        depth_trunc=kwargs.get("depth_trunc", 3.0),
    )

    point_cloud = point_cloud_denoising(point_cloud, kwargs.get("voxel_size", 0.010))
    if point_cloud is None:
        return None

    if len(point_cloud.points) == 0:
        logging.warning("clean_pcd is empty")
        return np.eye(4)
    center = point_cloud.get_center()
    x, y, z = center

    w, v = pca(np.asarray(point_cloud.points))

    if w is None or v is None:
        logging.warning("PCA output w or v is None")
        return np.eye(4)
    
    vx, vy, vz = v[:,0], v[:,1], v[:,2]

    if vx[0] < 0:
        vx = -vx
    if vy[1] < 0:
        vy = -vy
    if not np.allclose(np.cross(vx, vy), vz):
        vz = -vz

    R = np.column_stack((vx, vy, vz))
    rmat = tfs.affines.compose(np.squeeze(np.asarray((x, y, z))), R, [1, 1, 1])

    # draw(point_cloud, rmat)
    # draw(point_cloud_1, rmat)

    return rmat


def calculate_pose_icp(
        mask,
        depth_img: np.ndarray,
        intrinsics,
        **kwargs
):
    """点云配准法计算位姿"""
    depth_img_mask = np.zeros_like(depth_img)
    depth_img_mask[mask > 0] = depth_img[mask > 0]

    depth_o3d = o3d.geometry.Image(depth_img_mask.astype(np.uint16))

    point_cloud = o3d.geometry.PointCloud.create_from_depth_image(
        depth=depth_o3d,
        intrinsic=intrinsics,
        depth_scale=kwargs.get("depth_scale", 1000.0),
        depth_trunc=kwargs.get("depth_trunc", 3.0)
    )

    point_cloud = point_cloud_denoising(point_cloud, kwargs.get("voxel_size", 0.010))
    if point_cloud is None:
        return None

    if len(point_cloud.points) == 0:
        logging.warning("clean_pcd is empty")
        return np.eye(4)

    rmat = object_icp(
        kwargs.get("source"),
        point_cloud,
        ransac_voxel_size=kwargs.get("ransac_voxel_size", 0.005),
        icp_voxel_radius=kwargs.get("icp_voxel_radius", [0.004, 0.002, 0.001]),
        icp_max_iter=kwargs.get("icp_max_iter", [50, 30, 14])
    )

    return rmat


def point_cloud_denoising(point_cloud: o3d.geometry.PointCloud, voxel_size: float = 0.010):
    """点云去噪"""
    point_cloud = point_cloud.remove_non_finite_points()
    down_pcd = point_cloud.voxel_down_sample(voxel_size=voxel_size)

    # 半径滤波
    clean_pcd, _ = down_pcd.remove_radius_outlier(
        nb_points=10,
        radius=voxel_size * 5
        )

    # 统计滤波
    clean_pcd, _ = clean_pcd.remove_statistical_outlier(
        nb_neighbors=10,
        std_ratio=2.0
        )

    # 过滤过近的点
    points = np.asarray(clean_pcd.points)
    clean_pcd.points = o3d.utility.Vector3dVector(points[points[:, 2] >= 0.2])

    # # 使用数量最大簇判定噪声强度
    # _, counts = np.unique(labels[labels >= 0], return_counts=True)
    # largest_cluster_ratio = counts.max() / len(labels)
    # if largest_cluster_ratio < 0.5:
    #     return None

    labels = np.array(clean_pcd.cluster_dbscan(eps=voxel_size * 5, min_points=10))
    if len(labels[labels >= 0]) == 0:
        return clean_pcd

    # 使用距离最近簇作为物体
    points = np.asarray(clean_pcd.points)
    cluster_label = set(labels)
    point_cloud_clusters = []
    for label in cluster_label:
        if label == -1:
            continue
        idx = np.where(labels == label)[0]
        point_cloud_cluster = clean_pcd.select_by_index(idx)
        points_cluster_z = points[idx, 2]
        z_avg = np.mean(points_cluster_z)
        if z_avg < 0.2:
            continue
        point_cloud_clusters.append((point_cloud_cluster, z_avg))

    if len(point_cloud_clusters) == 0:
        return clean_pcd

    point_cloud_clusters.sort(key=lambda x: x[1])
    clean_pcd = point_cloud_clusters[0][0]

    # 使用最近簇判断噪音强度
    largest_cluster_ratio = len(clean_pcd.points) / len(points)
    if largest_cluster_ratio < 0.2:
        return None

    return clean_pcd


def calculate_grav_width(mask, k, depth):
    """计算重心宽度"""
    mask = mask.astype(np.uint8) * 255
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if not contours:
        return 0.0
    
    c = max(contours, key=cv2.contourArea)
    box = cv2.boxPoints(cv2.minAreaRect(c))
    if np.linalg.norm(box[1] - box[0]) < np.linalg.norm(box[1] - box[2]):
        point_diff = box[1] - box[0]
    else:
        point_diff = box[1] - box[2]

    grab_width = depth * np.sqrt(
        point_diff[0]**2 / k[0]**2 + point_diff[1]**2 / k[4]**2
        )
    return grab_width
    

def rmat2quat(rmat):
    """Convert rotation matrix to quaternion."""
    x, y, z = rmat[0:3, 3:4].flatten()
    rw, rx, ry, rz = tfs.quaternions.mat2quat(rmat[0:3, 0:3])
    quat = [x, y, z, rw, rx, ry, rz]
    return quat


def quat2rmat(quat):
    """Convert quaternion to rotation matrix."""
    x, y, z, rw, rx, ry, rz = quat
    r = tfs.quaternions.quat2mat([rw, rx, ry, rz])
    rmat = tfs.affines.compose(np.squeeze(np.asarray((x, y, z))), r, [1, 1, 1])
    return rmat

# def draw(pcd, mat):
#     R = mat[0:3, 0:3]
#     point = mat[0:3, 3:4].flatten()
#     x, y, z = R[:, 0], R[:, 1], R[:, 2]
#
#     points = [
#         [0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1],
#         point, point + x, point + y, point + z
#
#     ]  # 画点：原点、第一主成分、第二主成分
#     lines = [
#         [0, 1], [0, 2], [0, 3],
#         [4, 5], [4, 6], [4, 7]
#     ]  # 画出三点之间两两连线
#     colors = [
#         [1, 0, 0], [0, 1, 0], [0, 0, 1],
#         [1, 0, 0], [0, 1, 0], [0, 0, 1]
#     ]
#     line_set = o3d.geometry.LineSet(points=o3d.utility.Vector3dVector(points), lines=o3d.utility.Vector2iVector(lines))
#     line_set.colors = o3d.utility.Vector3dVector(colors)
#
#     o3d.visualization.draw_geometries([pcd, line_set])
