"""点云距离与法向一致性指标(倒角 L1/L2、豪斯多夫、|NC| / NC)。

modules.pipeline 与 scripts.eval 共用同一份实现, 保证推理内打印的指标与离线评测同口径。
"""
import numpy as np
from scipy.spatial import cKDTree


def calculate_distances_numpy(pcd1, pcd2):
    """
    [已修改]
    使用 NumPy 高效计算 L2倒角距离、L1倒角距离 和 豪斯多夫距离。

    参数:
    pcd1: np.ndarray, 形状为 (N, 3)，第一个点云。
    pcd2: np.ndarray, 形状为 (M, 3)，第二个点云。

    返回:
    (float, float, float): (L2倒角距离, 豪斯多夫距离, L1倒角距离)
    """
    # --- 计算所有点对之间的差异 ---
    pcd1_expanded = pcd1[:, np.newaxis, :]
    pcd2_expanded = pcd2[np.newaxis, :, :]
    diff = pcd1_expanded - pcd2_expanded

    # --- L2 (欧氏) 距离计算 ---
    dist_matrix_sq = np.sum(diff**2, axis=2)

    # --- L2 倒角距离 (Chamfer Distance L2) ---
    min_dists_1_to_2_l2 = np.sqrt(np.min(dist_matrix_sq, axis=1))
    min_dists_2_to_1_l2 = np.sqrt(np.min(dist_matrix_sq, axis=0))
    chamfer_dist_l2 = np.mean(min_dists_1_to_2_l2) + np.mean(min_dists_2_to_1_l2)

    # --- L2 豪斯多夫距离 (Hausdorff Distance) ---
    hausdorff_dist = max(np.max(min_dists_1_to_2_l2), np.max(min_dists_2_to_1_l2))
    
    # --- L1 (曼哈顿) 距离计算 ---
    dist_matrix_l1 = np.sum(np.abs(diff), axis=2)
    
    # --- L1 倒角距离 (Chamfer Distance L1) ---
    min_dists_1_to_2_l1 = np.min(dist_matrix_l1, axis=1)
    min_dists_2_to_1_l1 = np.min(dist_matrix_l1, axis=0)
    chamfer_dist_l1 = np.mean(min_dists_1_to_2_l1) + np.mean(min_dists_2_to_1_l1)

    return chamfer_dist_l2, hausdorff_dist, chamfer_dist_l1


def calculate_normal_consistency_score(pcd_gt, normal_gt, pcd_pred, normal_pred):
    """
    【修改版】根据提供的公式计算法向量一致性得分 (Normal Consistency Score)。
    返回两个值：
    1. |NC| (绝对值 NC): np.mean( |n1·n2| )
    2.  NC  (真实值 NC): np.mean( n1·n2 )
    """

    def _calculate_one_way_score(source_pcd, source_normals, target_pcd, target_normals):
        """
        辅助函数：计算从 source 到 target 的单向 NC 得分。
        返回 (mean_abs_dot, mean_real_dot)
        """
        # 1. 计算 source 到 target 的所有点对的距离矩阵
        diff = source_pcd[:, np.newaxis, :] - target_pcd[np.newaxis, :, :]
        dist_sq = np.sum(diff**2, axis=2)

        # 2. 为每个 source 点找到 target 中的最近邻点的索引
        indices = np.argmin(dist_sq, axis=1)

        # 3. 获取这些最近邻点对应的法向量
        corresponding_normals = target_normals[indices]

        # 4. 计算 source 法向量与对应 target 法向量的点积
        dot_products = np.einsum('ij,ij->i', source_normals, corresponding_normals)
        
        # 5. [修改] 返回点积绝对值的平均值 和 真实点积的平均值
        return np.mean(np.abs(dot_products)), np.mean(dot_products)

    # --- Part 1: 计算从 GT -> Pred 的得分 ---
    score_gt_to_pred_abs, score_gt_to_pred_real = _calculate_one_way_score(pcd_gt, normal_gt, pcd_pred, normal_pred)
    
    # --- Part 2: 计算从 Pred -> GT 的得分 ---
    score_pred_to_gt_abs, score_pred_to_gt_real = _calculate_one_way_score(pcd_pred, normal_pred, pcd_gt, normal_gt)
    
    # --- Final Score: 结合两部分 ---
    final_score_abs = (score_gt_to_pred_abs + score_pred_to_gt_abs) / 2.0
    final_score_real = (score_gt_to_pred_real + score_pred_to_gt_real) / 2.0
    
    # [修改] 返回两个分数
    return final_score_abs, final_score_real