"""
EMD (Earth Mover's Distance) 계산 모듈
3D 모델 간의 분포 유사성을 평가합니다.
"""

import numpy as np
from typing import Dict
from scipy.spatial.distance import cdist


class EMCalculator:
    """Earth Mover's Distance 계산을 담당하는 클래스"""
    
    def __init__(self, num_points: int = 1000):
        """
        EMD 계산기 초기화
        
        Args:
            num_points (int): 점군 생성 시 사용할 점의 개수
        """
        self.num_points = num_points
    
    def calculate_emd(self, model_3d: Dict, reference_3d: Dict) -> float:
        """
        두 3D 모델 간의 EMD를 계산합니다.
        
        Args:
            model_3d (Dict): 변환된 3D 모델 정보
            reference_3d (Dict): 참조 3D 모델 정보
            
        Returns:
            float: EMD 값 (낮을수록 유사함)
        """
        try:
            # 입력 데이터 검증
            if not isinstance(model_3d, dict) or not isinstance(reference_3d, dict):
                print("EMD 계산: 잘못된 입력 데이터 형식")
                return float('inf')
            
            if 'vertices' not in model_3d or 'vertices' not in reference_3d:
                print("EMD 계산: vertices 키가 없음")
                return float('inf')
            
            # 메시를 점군으로 변환
            pc1 = self._mesh_to_pointcloud(model_3d)
            pc2 = self._mesh_to_pointcloud(reference_3d)
            
            if len(pc1) == 0 or len(pc2) == 0:
                print("EMD 계산: 점군 변환 실패")
                return float('inf')
            
            # EMD 계산
            emd_value = self._compute_emd(pc1, pc2)
            
            # 결과 검증
            if np.isnan(emd_value) or np.isinf(emd_value):
                print(f"EMD 계산: 잘못된 결과 값 {emd_value}")
                return float('inf')
            
            return emd_value
            
        except Exception as e:
            print(f"EMD 계산 중 오류 발생: {e}")
            return float('inf')
    
    def _mesh_to_pointcloud(self, mesh: Dict) -> np.ndarray:
        """
        메시를 점군으로 변환합니다.
        
        Args:
            mesh (Dict): 3D 메시 정보
            
        Returns:
            np.ndarray: 점군 데이터
        """
        if 'vertices' not in mesh:
            return np.array([])
        
        vertices = mesh['vertices']
        
        if len(vertices) == 0:
            return np.array([])
        
        # vertices를 numpy 배열로 변환 (리스트인 경우 처리)
        if not isinstance(vertices, np.ndarray):
            try:
                vertices = np.array(vertices, dtype=np.float32)
            except (ValueError, TypeError) as e:
                print(f"EMD: vertices 변환 실패: {e}")
                return np.array([])
        
        # vertices가 2D 배열인지 확인
        if vertices.ndim != 2 or vertices.shape[1] != 3:
            print(f"EMD: 잘못된 vertices 형태: {vertices.shape}")
            return np.array([])
        
        # 정점이 충분한 경우 랜덤 샘플링
        if len(vertices) >= self.num_points:
            indices = np.random.choice(len(vertices), self.num_points, replace=False)
            return vertices[indices]
        else:
            # 정점이 부족한 경우 중복 허용 샘플링
            indices = np.random.choice(len(vertices), self.num_points, replace=True)
            return vertices[indices]
    
    def _compute_emd(self, pc1: np.ndarray, pc2: np.ndarray) -> float:
        """
        두 점군 간의 EMD를 계산합니다.
        
        Args:
            pc1 (np.ndarray): 첫 번째 점군
            pc2 (np.ndarray): 두 번째 점군
            
        Returns:
            float: EMD 값
        """
        # 점군을 정규화
        pc1_norm = self._normalize_pointcloud(pc1)
        pc2_norm = self._normalize_pointcloud(pc2)
        
        # 근사 EMD 계산 (최근접 이웃 기반)
        distances = cdist(pc1_norm, pc2_norm, metric='euclidean')
        min_distances_1 = np.min(distances, axis=1)
        min_distances_2 = np.min(distances, axis=0)
        
        # 근사 EMD = 평균 최단 거리
        emd_value = (np.mean(min_distances_1) + np.mean(min_distances_2)) / 2.0
        
        return emd_value
    
    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