"""
Chamfer Distance 계산 모듈
3D 모델 간의 기하학적 유사성을 평가합니다.
"""

import numpy as np
import trimesh
from typing import Dict, Tuple, Optional
from scipy.spatial.distance import cdist
import open3d as o3d


class ChamferDistanceCalculator:
    """Chamfer Distance 계산을 담당하는 클래스"""
    
    def __init__(self, num_points: int = 10000):
        """
        Chamfer Distance 계산기 초기화
        
        Args:
            num_points (int): 점군 생성 시 사용할 점의 개수
        """
        self.num_points = num_points
    
    def calculate_chamfer_distance(self, model_3d: Dict, reference_3d: Dict) -> float:
        """
        두 3D 모델 간의 Chamfer Distance를 계산합니다 (개선된 버전).
        
        Args:
            model_3d (Dict): 변환된 3D 모델 정보
            reference_3d (Dict): 참조 3D 모델 정보
            
        Returns:
            float: Chamfer Distance 값 (낮을수록 유사함)
        """
        return self.chamfer_distance_optimized_improved(model_3d, reference_3d)
    
    def chamfer_distance_optimized_improved(self, model_3d: Dict, reference_3d: Dict) -> float:
        """
        개선된 최적화 Chamfer Distance 계산.
        
        Args:
            model_3d (Dict): 변환된 3D 모델 정보
            reference_3d (Dict): 참조 3D 모델 정보
            
        Returns:
            float: Chamfer Distance 값
        """
        try:
            # 고품질 점군 생성
            pc1 = self._mesh_to_pointcloud_high_quality(model_3d, self.num_points)
            pc2 = self._mesh_to_pointcloud_high_quality(reference_3d, self.num_points)
            
            if len(pc1) == 0 or len(pc2) == 0:
                return float('inf')
            
            # 점군 정렬
            pc1_aligned, pc2_aligned = self._align_pointclouds(pc1, pc2)
            
            # 최적화된 Chamfer Distance 계산
            chamfer_dist = self._compute_chamfer_distance_optimized(pc1_aligned, pc2_aligned)
            
            return chamfer_dist
            
        except Exception as e:
            print(f"Chamfer Distance 계산 중 오류: {e}")
            # 폴백: 기본 방법 사용
            return self._fallback_chamfer_distance(model_3d, reference_3d)
    
    def mesh_to_pointcloud(self, mesh: Dict, num_points: int = None) -> np.ndarray:
        """
        메시를 점군으로 변환합니다.
        
        Args:
            mesh (Dict): 3D 메시 정보
            num_points (int): 생성할 점의 개수
            
        Returns:
            np.ndarray: 점군 데이터 (N, 3)
        """
        if num_points is None:
            num_points = self.num_points
        
        # vertices 키 존재 여부 확인
        if 'vertices' not in mesh:
            return np.array([])
        
        vertices = mesh['vertices']
        faces = mesh.get('faces', [])
        
        if len(vertices) == 0:
            return np.array([])
        
        # Open3D를 사용한 고품질 점군 생성
        try:
            return self._mesh_to_pointcloud_open3d(vertices, faces, num_points)
        except Exception:
            # Open3D 실패 시 간단한 방법 사용
            return self._mesh_to_pointcloud_simple(vertices, faces, num_points)
    
    def _mesh_to_pointcloud_open3d(self, vertices: np.ndarray, faces: np.ndarray, 
                                  num_points: int) -> np.ndarray:
        """
        Open3D를 사용하여 메시를 점군으로 변환합니다.
        
        Args:
            vertices (np.ndarray): 메시 정점
            faces (np.ndarray): 메시 면
            num_points (int): 생성할 점의 개수
            
        Returns:
            np.ndarray: 점군 데이터
        """
        # Open3D 메시 생성
        mesh = o3d.geometry.TriangleMesh()
        mesh.vertices = o3d.utility.Vector3dVector(vertices)
        mesh.triangles = o3d.utility.Vector3iVector(faces)
        
        # 법선 계산
        mesh.compute_vertex_normals()
        
        # 표면에서 균등하게 점 샘플링
        pointcloud = mesh.sample_points_uniformly(number_of_points=num_points)
        
        return np.asarray(pointcloud.points)
    
    def _mesh_to_pointcloud_simple(self, vertices: np.ndarray, faces: np.ndarray, 
                                  num_points: int) -> np.ndarray:
        """
        간단한 방법으로 메시를 점군으로 변환합니다.
        
        Args:
            vertices (np.ndarray): 메시 정점
            faces (np.ndarray): 메시 면
            num_points (int): 생성할 점의 개수
            
        Returns:
            np.ndarray: 점군 데이터
        """
        if len(vertices) >= num_points:
            # 정점이 충분한 경우 랜덤 샘플링
            indices = np.random.choice(len(vertices), num_points, replace=False)
            return vertices[indices]
        else:
            # 정점이 부족한 경우 면 기반 샘플링
            return self._sample_points_from_faces(vertices, faces, num_points)
    
    def _sample_points_from_faces(self, vertices: np.ndarray, faces: np.ndarray, 
                                 num_points: int) -> np.ndarray:
        """
        면에서 점을 샘플링합니다.
        
        Args:
            vertices (np.ndarray): 메시 정점
            faces (np.ndarray): 메시 면
            num_points (int): 생성할 점의 개수
            
        Returns:
            np.ndarray: 샘플링된 점군
        """
        if len(faces) == 0:
            # 면이 없는 경우 정점에서 중복 허용 샘플링
            indices = np.random.choice(len(vertices), num_points, replace=True)
            return vertices[indices]
        
        # 각 면의 면적 계산
        face_areas = []
        for face in faces:
            v0, v1, v2 = vertices[face]
            area = 0.5 * np.linalg.norm(np.cross(v1 - v0, v2 - v0))
            face_areas.append(area)
        
        face_areas = np.array(face_areas)
        total_area = np.sum(face_areas)
        
        if total_area == 0:
            # 면적이 0인 경우 정점에서 샘플링
            indices = np.random.choice(len(vertices), num_points, replace=True)
            return vertices[indices]
        
        # 면적에 비례하여 점 샘플링
        face_probs = face_areas / total_area
        sampled_faces = np.random.choice(len(faces), num_points, p=face_probs)
        
        points = []
        for face_idx in sampled_faces:
            face = faces[face_idx]
            v0, v1, v2 = vertices[face]
            
            # 삼각형 내에서 균등하게 점 샘플링
            r1, r2 = np.random.random(2)
            if r1 + r2 > 1:
                r1, r2 = 1 - r1, 1 - r2
            
            point = v0 + r1 * (v1 - v0) + r2 * (v2 - v0)
            points.append(point)
        
        return np.array(points)
    
    def chamfer_distance_naive(self, pc1: np.ndarray, pc2: np.ndarray) -> float:
        """
        두 점군 간의 Chamfer Distance를 계산합니다.
        
        Args:
            pc1 (np.ndarray): 첫 번째 점군 (N1, 3)
            pc2 (np.ndarray): 두 번째 점군 (N2, 3)
            
        Returns:
            float: Chamfer Distance
        """
        if len(pc1) == 0 or len(pc2) == 0:
            return float('inf')
        
        # pc1의 각 점에서 pc2까지의 최단 거리
        dist1 = self._compute_point_to_set_distance(pc1, pc2)
        
        # pc2의 각 점에서 pc1까지의 최단 거리
        dist2 = self._compute_point_to_set_distance(pc2, pc1)
        
        # Chamfer Distance = 평균 최단 거리
        chamfer_dist = (np.mean(dist1) + np.mean(dist2)) / 2.0
        
        return chamfer_dist
    
    def _compute_point_to_set_distance(self, points: np.ndarray, reference_set: np.ndarray) -> np.ndarray:
        """
        점들에서 참조 집합까지의 최단 거리를 계산합니다.
        
        Args:
            points (np.ndarray): 점들 (N, 3)
            reference_set (np.ndarray): 참조 점군 (M, 3)
            
        Returns:
            np.ndarray: 각 점의 최단 거리 (N,)
        """
        # 모든 점 쌍 간의 거리 계산
        distances = cdist(points, reference_set, metric='euclidean')
        
        # 각 점에서 참조 집합까지의 최단 거리
        min_distances = np.min(distances, axis=1)
        
        return min_distances
    
    def chamfer_distance_optimized(self, pc1: np.ndarray, pc2: np.ndarray) -> float:
        """
        최적화된 Chamfer Distance 계산 (KD-Tree 사용).
        
        Args:
            pc1 (np.ndarray): 첫 번째 점군
            pc2 (np.ndarray): 두 번째 점군
            
        Returns:
            float: Chamfer Distance
        """
        try:
            from scipy.spatial import cKDTree
            
            if len(pc1) == 0 or len(pc2) == 0:
                return float('inf')
            
            # KD-Tree 구축
            tree1 = cKDTree(pc1)
            tree2 = cKDTree(pc2)
            
            # pc1의 각 점에서 pc2까지의 최단 거리
            dist1, _ = tree2.query(pc1)
            
            # pc2의 각 점에서 pc1까지의 최단 거리
            dist2, _ = tree1.query(pc2)
            
            # Chamfer Distance
            chamfer_dist = (np.mean(dist1) + np.mean(dist2)) / 2.0
            
            return chamfer_dist
            
        except ImportError:
            # scipy가 없는 경우 naive 방법 사용
            return self.chamfer_distance_naive(pc1, pc2)
    
    def calculate_symmetric_chamfer_distance(self, pc1: np.ndarray, pc2: np.ndarray) -> float:
        """
        대칭 Chamfer Distance를 계산합니다.
        
        Args:
            pc1 (np.ndarray): 첫 번째 점군
            pc2 (np.ndarray): 두 번째 점군
            
        Returns:
            float: 대칭 Chamfer Distance
        """
        # 양방향 Chamfer Distance 계산
        cd1 = self._compute_point_to_set_distance(pc1, pc2)
        cd2 = self._compute_point_to_set_distance(pc2, pc1)
        
        # 대칭 Chamfer Distance
        symmetric_cd = np.mean(cd1) + np.mean(cd2)
        
        return symmetric_cd
    
    def calculate_hausdorff_distance(self, pc1: np.ndarray, pc2: np.ndarray) -> float:
        """
        Hausdorff Distance를 계산합니다.
        
        Args:
            pc1 (np.ndarray): 첫 번째 점군
            pc2 (np.ndarray): 두 번째 점군
            
        Returns:
            float: Hausdorff Distance
        """
        if len(pc1) == 0 or len(pc2) == 0:
            return float('inf')
        
        # 모든 점 쌍 간의 거리 계산
        distances = cdist(pc1, pc2, metric='euclidean')
        
        # 각 점에서 다른 점군까지의 최단 거리
        min_dist1 = np.min(distances, axis=1)  # pc1의 각 점에서 pc2까지
        min_dist2 = np.min(distances, axis=0)  # pc2의 각 점에서 pc1까지
        
        # Hausdorff Distance = 최대 최단 거리
        hausdorff_dist = max(np.max(min_dist1), np.max(min_dist2))
        
        return hausdorff_dist
    
    def calculate_point_cloud_density(self, pointcloud: np.ndarray) -> float:
        """
        점군의 밀도를 계산합니다.
        
        Args:
            pointcloud (np.ndarray): 점군 데이터
            
        Returns:
            float: 점군 밀도
        """
        if len(pointcloud) < 2:
            return 0.0
        
        # 점군의 바운딩 박스 부피 계산
        min_coords = np.min(pointcloud, axis=0)
        max_coords = np.max(pointcloud, axis=0)
        volume = np.prod(max_coords - min_coords)
        
        if volume == 0:
            return 0.0
        
        # 밀도 = 점의 개수 / 부피
        density = len(pointcloud) / volume
        
        return density
    
    def normalize_pointcloud(self, pointcloud: np.ndarray) -> np.ndarray:
        """
        점군을 정규화합니다.
        
        Args:
            pointcloud (np.ndarray): 입력 점군
            
        Returns:
            np.ndarray: 정규화된 점군
        """
        if len(pointcloud) == 0:
            return pointcloud
        
        # 중심을 원점으로 이동
        center = np.mean(pointcloud, axis=0)
        centered = pointcloud - center
        
        # 스케일 정규화 (최대 거리를 1로)
        max_dist = np.max(np.linalg.norm(centered, axis=1))
        if max_dist > 0:
            normalized = centered / max_dist
        else:
            normalized = centered
        
        return normalized
    
    def calculate_chamfer_distance_normalized(self, model_3d: Dict, reference_3d: Dict) -> float:
        """
        정규화된 Chamfer Distance를 계산합니다.
        
        Args:
            model_3d (Dict): 변환된 3D 모델 정보
            reference_3d (Dict): 참조 3D 모델 정보
            
        Returns:
            float: 정규화된 Chamfer Distance
        """
        # 메시를 점군으로 변환
        pc1 = self.mesh_to_pointcloud(model_3d, self.num_points)
        pc2 = self.mesh_to_pointcloud(reference_3d, self.num_points)
        
        if len(pc1) == 0 or len(pc2) == 0:
            return float('inf')
        
        # 점군 정규화
        pc1_norm = self.normalize_pointcloud(pc1)
        pc2_norm = self.normalize_pointcloud(pc2)
        
        # 정규화된 Chamfer Distance 계산
        chamfer_dist = self.chamfer_distance_naive(pc1_norm, pc2_norm)
        
        return chamfer_dist
    
    def _mesh_to_pointcloud_high_quality(self, mesh: Dict, num_points: int) -> np.ndarray:
        """
        고품질 점군 생성 (Open3D 메시 전처리, 중복 제거, 법선 계산).
        
        Args:
            mesh (Dict): 3D 메시 정보
            num_points (int): 생성할 점의 개수
            
        Returns:
            np.ndarray: 고품질 점군 데이터
        """
        try:
            # vertices 키 존재 여부 확인
            if 'vertices' not in mesh:
                return np.array([])
            
            vertices = mesh['vertices']
            faces = mesh.get('faces', [])
            
            if len(vertices) == 0:
                return np.array([])
            
            # Open3D 메시 생성
            o3d_mesh = o3d.geometry.TriangleMesh()
            o3d_mesh.vertices = o3d.utility.Vector3dVector(vertices)
            o3d_mesh.triangles = o3d.utility.Vector3iVector(faces)
            
            # 메시 전처리
            o3d_mesh = self._preprocess_mesh_open3d(o3d_mesh)
            
            # 고품질 점 샘플링
            pointcloud = o3d_mesh.sample_points_uniformly(number_of_points=num_points)
            
            # 점군 후처리
            processed_pc = self._postprocess_pointcloud(np.asarray(pointcloud.points))
            
            return processed_pc
            
        except Exception as e:
            print(f"고품질 점군 생성 실패: {e}")
            # 폴백: 기본 방법 사용
            return self.mesh_to_pointcloud(mesh, num_points)
    
    def _preprocess_mesh_open3d(self, mesh: o3d.geometry.TriangleMesh) -> o3d.geometry.TriangleMesh:
        """Open3D 메시 전처리 (중복 제거, 법선 계산)."""
        # 중복 정점 제거
        mesh.remove_duplicated_vertices()
        
        # 중복 면 제거
        mesh.remove_duplicated_triangles()
        
        # 비정상적인 면 제거
        mesh.remove_degenerate_triangles()
        
        # 비연결된 정점 제거
        mesh.remove_unreferenced_vertices()
        
        # 법선 계산
        mesh.compute_vertex_normals()
        mesh.compute_triangle_normals()
        
        # 메시 정리
        mesh.remove_degenerate_triangles()
        mesh.remove_duplicated_triangles()
        mesh.remove_duplicated_vertices()
        mesh.remove_unreferenced_vertices()
        
        return mesh
    
    def _postprocess_pointcloud(self, pointcloud: np.ndarray) -> np.ndarray:
        """점군 후처리 (노이즈 제거, 정규화)."""
        if len(pointcloud) == 0:
            return pointcloud
        
        # 통계적 이상치 제거
        pointcloud = self._remove_outliers(pointcloud)
        
        # 점군 정규화
        pointcloud = self._normalize_pointcloud_improved(pointcloud)
        
        return pointcloud
    
    def _remove_outliers(self, pointcloud: np.ndarray, std_ratio: float = 2.0) -> np.ndarray:
        """통계적 이상치 제거."""
        if len(pointcloud) < 10:
            return pointcloud
        
        # 각 점에서 다른 모든 점까지의 평균 거리 계산
        distances = cdist(pointcloud, pointcloud, metric='euclidean')
        mean_distances = np.mean(distances, axis=1)
        
        # 평균과 표준편차 계산
        mean_dist = np.mean(mean_distances)
        std_dist = np.std(mean_distances)
        
        # 이상치 임계값
        threshold = mean_dist + std_ratio * std_dist
        
        # 이상치가 아닌 점들만 유지
        valid_indices = mean_distances < threshold
        
        return pointcloud[valid_indices]
    
    def _normalize_pointcloud_improved(self, pointcloud: np.ndarray) -> np.ndarray:
        """개선된 점군 정규화."""
        if len(pointcloud) == 0:
            return pointcloud
        
        # 중심점 계산
        center = np.mean(pointcloud, axis=0)
        centered = pointcloud - center
        
        # 스케일 정규화 (RMS 거리 사용)
        rms_distance = np.sqrt(np.mean(np.sum(centered**2, axis=1)))
        if rms_distance > 0:
            normalized = centered / rms_distance
        else:
            normalized = centered
        
        return normalized
    
    def _align_pointclouds(self, pc1: np.ndarray, pc2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        점군 정렬 (중심점 정렬, 스케일 정규화, ICP 정렬).
        
        Args:
            pc1 (np.ndarray): 첫 번째 점군
            pc2 (np.ndarray): 두 번째 점군
            
        Returns:
            Tuple[np.ndarray, np.ndarray]: 정렬된 점군들
        """
        # 1. 중심점 정렬
        pc1_centered, pc2_centered = self._center_align_pointclouds(pc1, pc2)
        
        # 2. 스케일 정규화
        pc1_scaled, pc2_scaled = self._scale_normalize_pointclouds(pc1_centered, pc2_centered)
        
        # 3. ICP 정렬 (선택적)
        try:
            pc1_aligned, pc2_aligned = self._icp_align_pointclouds(pc1_scaled, pc2_scaled)
            return pc1_aligned, pc2_aligned
        except Exception:
            # ICP 실패 시 스케일 정규화된 점군 반환
            return pc1_scaled, pc2_scaled
    
    def _center_align_pointclouds(self, pc1: np.ndarray, pc2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """점군들을 중심점으로 정렬."""
        center1 = np.mean(pc1, axis=0)
        center2 = np.mean(pc2, axis=0)
        
        pc1_centered = pc1 - center1
        pc2_centered = pc2 - center2
        
        return pc1_centered, pc2_centered
    
    def _scale_normalize_pointclouds(self, pc1: np.ndarray, pc2: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """점군들을 스케일 정규화."""
        # RMS 거리 계산
        rms1 = np.sqrt(np.mean(np.sum(pc1**2, axis=1)))
        rms2 = np.sqrt(np.mean(np.sum(pc2**2, axis=1)))
        
        # 정규화
        if rms1 > 0:
            pc1_scaled = pc1 / rms1
        else:
            pc1_scaled = pc1
            
        if rms2 > 0:
            pc2_scaled = pc2 / rms2
        else:
            pc2_scaled = pc2
        
        return pc1_scaled, pc2_scaled
    
    def _icp_align_pointclouds(self, pc1: np.ndarray, pc2: np.ndarray, 
                              max_iterations: int = 50) -> Tuple[np.ndarray, np.ndarray]:
        """ICP를 사용한 점군 정렬."""
        try:
            # Open3D 점군 생성
            pcd1 = o3d.geometry.PointCloud()
            pcd1.points = o3d.utility.Vector3dVector(pc1)
            
            pcd2 = o3d.geometry.PointCloud()
            pcd2.points = o3d.utility.Vector3dVector(pc2)
            
            # ICP 실행
            result = o3d.pipelines.registration.registration_icp(
                pcd1, pcd2, 
                max_correspondence_distance=0.1,
                estimation_method=o3d.pipelines.registration.TransformationEstimationPointToPoint(),
                criteria=o3d.pipelines.registration.ICPConvergenceCriteria(max_iteration=max_iterations)
            )
            
            # 변환 적용
            pcd1_transformed = pcd1.transform(result.transformation)
            
            return np.asarray(pcd1_transformed.points), pc2
            
        except Exception as e:
            print(f"ICP 정렬 실패: {e}")
            return pc1, pc2
    
    def _compute_chamfer_distance_optimized(self, pc1: np.ndarray, pc2: np.ndarray) -> float:
        """
        최적화된 Chamfer Distance 계산 (KD-Tree 기반, 배치 처리).
        
        Args:
            pc1 (np.ndarray): 첫 번째 점군
            pc2 (np.ndarray): 두 번째 점군
            
        Returns:
            float: Chamfer Distance
        """
        try:
            from scipy.spatial import cKDTree
            
            if len(pc1) == 0 or len(pc2) == 0:
                return float('inf')
            
            # 배치 크기 설정 (메모리 효율성)
            batch_size = min(1000, len(pc1), len(pc2))
            
            # KD-Tree 구축
            tree1 = cKDTree(pc1)
            tree2 = cKDTree(pc2)
            
            # 배치 처리로 거리 계산
            total_dist1 = 0.0
            total_dist2 = 0.0
            
            # pc1의 각 점에서 pc2까지의 최단 거리
            for i in range(0, len(pc1), batch_size):
                batch_pc1 = pc1[i:i+batch_size]
                dist1, _ = tree2.query(batch_pc1)
                total_dist1 += np.sum(dist1)
            
            # pc2의 각 점에서 pc1까지의 최단 거리
            for i in range(0, len(pc2), batch_size):
                batch_pc2 = pc2[i:i+batch_size]
                dist2, _ = tree1.query(batch_pc2)
                total_dist2 += np.sum(dist2)
            
            # Chamfer Distance 계산
            chamfer_dist = (total_dist1 / len(pc1) + total_dist2 / len(pc2)) / 2.0
            
            return chamfer_dist
            
        except ImportError:
            # scipy가 없는 경우 기본 방법 사용
            return self.chamfer_distance_naive(pc1, pc2)
        except Exception as e:
            print(f"최적화된 Chamfer Distance 계산 실패: {e}")
            return self.chamfer_distance_naive(pc1, pc2)
    
    def _fallback_chamfer_distance(self, model_3d: Dict, reference_3d: Dict) -> float:
        """폴백 Chamfer Distance 계산."""
        try:
            # 기본 점군 생성
            pc1 = self.mesh_to_pointcloud(model_3d, min(self.num_points, 5000))
            pc2 = self.mesh_to_pointcloud(reference_3d, min(self.num_points, 5000))
            
            if len(pc1) == 0 or len(pc2) == 0:
                return float('inf')
            
            # 기본 정규화
            pc1_norm = self.normalize_pointcloud(pc1)
            pc2_norm = self.normalize_pointcloud(pc2)
            
            # 기본 Chamfer Distance 계산
            return self.chamfer_distance_naive(pc1_norm, pc2_norm)
            
        except Exception as e:
            print(f"폴백 Chamfer Distance 계산 실패: {e}")
            return float('inf')
