import math
import json
import logging

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

from .object_icp import object_icp


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,
        source=None,
        hand_eye_mat=np.eye(4),
        configs=None
):
    """计算位态"""
    if configs is None:
        configs = {
            "depth_scale": 1000.0,
            "depth_trunc": 3.0,
            "voxel_size": 0.020,
            "nb_points": 10,
            "radius": 0.1,
            "nb_neighbors": 20,
            "std_ratio": 3.0
        }
    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=configs["depth_scale"],
        depth_trunc=configs["depth_trunc"],
    )

    point_cloud = point_cloud.remove_non_finite_points()

    down_pcd = point_cloud.voxel_down_sample(voxel_size=configs["voxel_size"])
    clean_pcd, _ = down_pcd.remove_radius_outlier(nb_points=configs["nb_points"], radius=configs["radius"])
    clean_pcd, _ = clean_pcd.remove_statistical_outlier(nb_neighbors=configs["nb_neighbors"], std_ratio=configs["std_ratio"])

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

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

    if w is not None and v is not None:
        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])
        rmat = hand_eye_mat @ rmat
        rw, rx, ry, rz = tfs.quaternions.mat2quat(rmat[0:3, 0:3])
        x, y, z = rmat[0:3, 3].flatten()

        return x, y, z, rw, rx, ry, rz
    else:
        logging.warning("PCA output w or v is None")
        return 0.0, 0.0, 0.0, None, None, None, None


def calculate_pose_icp(
        mask,
        depth_img: np.ndarray,
        intrinsics,
        source,
        hand_eye_mat=np.eye(4),
        configs = None,
):
    if configs is None:
        configs = {
            "depth_scale": 1000.0,
            "depth_trunc": 2.0,
            "nb_points": 10,
            "radius": 0.1,
            "nb_neighbors": 20,
            "std_ratio": 3.0,
            "ransac_voxel_size": 0.005,
            "icp_voxel_radius": [0.004, 0.002, 0.001],
            "icp_max_iter": [50, 30, 14]
        }
    """计算位态"""
    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=configs["depth_scale"],
        depth_trunc=configs["depth_trunc"],
    )

    point_cloud = point_cloud.remove_non_finite_points()

    clean_pcd, _ = point_cloud.remove_radius_outlier(nb_points=configs["nb_points"], radius=configs["radius"])
    clean_pcd, _ = clean_pcd.remove_statistical_outlier(nb_neighbors=configs["nb_neighbors"], std_ratio=configs["std_ratio"])

    if len(clean_pcd.points) == 0:
        logging.warning("clean_pcd is empty")
        return 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0

    rmat = object_icp(
        source,
        clean_pcd,
        ransac_voxel_size=configs["ransac_voxel_size"],
        icp_voxel_radius=configs["icp_voxel_radius"],
        icp_max_iter=configs["icp_max_iter"]
    )

    rmat = hand_eye_mat @ rmat

    x, y, z = rmat[0:3, 3:4].flatten()
    rw, rx, ry, rz = tfs.quaternions.mat2quat(rmat[0:3, 0:3])

    return x, y, z, rw, rx, ry, rz


def calculate_grav_width(mask, K, depth):
    """计算重心宽度"""
    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if contours:
        box = cv2.boxPoints(cv2.minAreaRect(contours[0]))
        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
    else:
        return 0.0


