import numpy as np
import cv2

# Try to import CuPy for GPU acceleration
try:
    import cupy as cp
    CUPY_AVAILABLE = True
except ImportError:
    cp = None
    CUPY_AVAILABLE = False

# 缓存形态学操作的核，避免重复创建
_morphology_kernels = {}

class Contact3D:
    def __init__(self, config_contact=None, verbose=False):

        # Load Config
        self.config_contact = config_contact
        self.flow_threshold = config_contact["flow_threshold"]
        self.expansion_threshold = config_contact["expansion_threshold"]
        self.downsample_factor = 1
        self.kernel_size = config_contact["kernel_size"]
        self.denoise_sigma = config_contact["denoise_sigma"]
        self.min_contour_area = config_contact["min_contour_area"]
        # Divergence threshold: only keep divergence values above this threshold, set others to zero
        self.divergence_threshold = config_contact["div_threshold"]
        self.verbose = verbose
        if self.verbose:
            print(f"Contact3D initialized with config: {self.config_contact}")

        # Data Variables
        self.foe_results = None
        # self.expansion_mask = None
        # self.contraction_mask = None
        # self.divergence_map = None
        # self.centroid = None
        # self.centroid_found = False
        # self.contour = None
        # self.all_contours = None

    def _create_foe_result(self, expansion_mask, contraction_mask, divergence_map, 
                        centroid=None, centroid_found=False, contour=None, 
                        all_contours=None):
        """
        创建标准化的FOE结果字典，统一所有函数的返回格式。
        
        Args:
            expansion_mask: 扩张区域掩码
            contraction_mask: 收缩区域掩码
            divergence_map: 散度图
            centroid: 质心坐标，可以是单个元组 (x, y) 或列表 [(x1, y1), (x2, y2), ...]
            centroid_found: 是否找到质心
            contour: 轮廓，可以是单个轮廓或轮廓列表
            all_contours: 所有轮廓（仅在return_all_contours=True时使用）
        
        Returns:
            dict: 标准化的FOE结果字典
        """
        result = {
            'expansion_mask': expansion_mask,
            'contraction_mask': contraction_mask,
            'divergence_map': divergence_map,
            'centroid': centroid if centroid is not None else [0,0],
            'centroid_found': centroid_found,
        }
        
        # 可选字段
        if contour is not None:
            result['contour'] = contour
        
        if all_contours is not None:
            result['_all_contours'] = all_contours
        
        return result
    
    def compute_foe(self, flow_field, return_all_contours=False):
        """
        计算光流场的扩张焦点(Focus of Expansion)。
        
        Args:
            flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
            min_magnitude: 最小光流幅度阈值，低于此值的点不参与计算
            expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
            downsample_factor: 下采样因子，用于平滑计算（默认10）
            return_all_contours: 是否返回所有轮廓（默认False，仅返回最大轮廓）
            kernel_size: 形态学操作的核大小
            denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
            
        Returns:
            dict: 包含扩张焦点信息的字典
        """
        H, W = flow_field.shape[:2]
        
        # 优化：使用视图而不是复制，减少内存分配
        u = flow_field[:, :, 0]
        v = flow_field[:, :, 1]
        
        # 下采样并计算散度，然后插值回原尺寸
        if self.downsample_factor > 1:
            # 使用步进切片，避免复制
            u_ds = u[::self.downsample_factor, ::self.downsample_factor]
            v_ds = v[::self.downsample_factor, ::self.downsample_factor]
            
            # 计算幅度（优化：使用square减少中间数组）
            magnitude_ds = np.sqrt(np.square(u_ds) + np.square(v_ds))
            if self.denoise_sigma > 0:
                magnitude_ds = cv2.GaussianBlur(magnitude_ds, (0, 0), self.denoise_sigma)
            
            # 优化：合并梯度计算，减少中间数组
            # 计算散度：du/dx + dv/dy
            du_dx = np.gradient(u_ds, axis=1)
            dv_dy = np.gradient(v_ds, axis=0)
            divergence_ds = du_dx + dv_dy  # 直接相加，避免中间数组
            
            if self.denoise_sigma > 0:
                divergence_ds = cv2.GaussianBlur(divergence_ds, (0, 0), self.denoise_sigma)
            
            # 插值回原尺寸
            magnitude = cv2.resize(magnitude_ds, (W, H), interpolation=cv2.INTER_CUBIC)
            divergence = cv2.resize(divergence_ds, (W, H), interpolation=cv2.INTER_CUBIC)
        else:
            # 计算幅度（优化：使用square减少中间数组）
            magnitude = np.sqrt(np.square(u) + np.square(v))
            if self.denoise_sigma > 0:
                magnitude = cv2.GaussianBlur(magnitude, (0, 0), self.denoise_sigma)
            
            # 优化：合并梯度计算
            du_dx = np.gradient(u, axis=1)
            dv_dy = np.gradient(v, axis=0)
            divergence = du_dx + dv_dy
            
            # 降噪：在计算散度之后对散度图进行平滑
            if self.denoise_sigma > 0:
                divergence = cv2.GaussianBlur(divergence, (0, 0), self.denoise_sigma)
        
        # 优化：使用原地操作 - 只保留大于阈值的散度值，其他设为0
        divergence = np.where(divergence > self.divergence_threshold, divergence, 0)
        
        # 创建掩码（使用原地操作优化）
        valid_mask = magnitude > self.flow_threshold
        expansion_mask = (valid_mask & (divergence > self.expansion_threshold))
        contraction_mask = (valid_mask & (divergence < -self.expansion_threshold))
        
        # 初始化结果字典
        result = self._create_foe_result(expansion_mask, contraction_mask, divergence)
        
        # 寻找扩张区域轮廓
        # 优化：缓存形态学操作的核
        # Convert kernel_size to tuple for use as dictionary key (lists are not hashable)
        kernel_key = tuple(self.kernel_size) if isinstance(self.kernel_size, (list, np.ndarray)) else self.kernel_size
        if kernel_key not in _morphology_kernels:
            _morphology_kernels[kernel_key] = np.ones(self.kernel_size, np.uint8)
        
        expansion_binary = expansion_mask.astype(np.uint8) * 255
        if np.any(expansion_binary):
            expansion_binary_closed = cv2.morphologyEx(expansion_binary, cv2.MORPH_CLOSE, _morphology_kernels[kernel_key])
            contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            if contours:
                result['centroid_found'] = True
                if return_all_contours:
                    result['_all_contours'] = contours
                
                largest_contour = max(contours, key=cv2.contourArea)
                result['contour'] = largest_contour
                
                M = cv2.moments(largest_contour)
                if M["m00"] > 0:
                    result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
        
        self.foe_results = result
        return self.foe_results

    def compute_foe_gpu(self, flow_field, return_all_contours=False):
        """
        使用CUDA加速计算光流场的扩张焦点(Focus of Expansion)。
        
        Args:
            flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
            flow_threshold: 最小光流幅度阈值，低于此值的点不参与计算
            expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
            downsample_factor: 下采样因子，用于平滑计算（默认10）
            return_all_contours: 是否返回所有轮廓（默认False，仅返回最大轮廓）
            kernel_size: 形态学操作的核大小
            denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
            
        Returns:
            dict: 包含扩张焦点信息的字典
            
        Note:
            如果CuPy不可用，会自动回退到CPU版本的compute_foe
        """
        # 检查CuPy是否可用
        if not CUPY_AVAILABLE:
            return self.compute_foe(flow_field, return_all_contours)
        
        H, W = flow_field.shape[:2]
        
        # 导入CuPy的高斯滤波函数（仅在需要时导入）
        try:
            from cupyx.scipy.ndimage import gaussian_filter as cp_gaussian_filter
        except ImportError:
            # 如果导入失败，回退到CPU版本
            return self.compute_foe(flow_field, return_all_contours)
        
        # 优化：使用asarray而不是array，如果已经是CuPy数组则避免复制
        if isinstance(flow_field, cp.ndarray):
            flow_gpu = flow_field
        else:
            flow_gpu = cp.asarray(flow_field, dtype=cp.float32)
        
        # 优化：使用视图而不是复制
        u_gpu = flow_gpu[:, :, 0]
        v_gpu = flow_gpu[:, :, 1]
        
        # 下采样并计算散度，然后插值回原尺寸
        if self.downsample_factor > 1:
            # 下采样（使用步进切片，避免复制）
            u_ds_gpu = u_gpu[::self.downsample_factor, ::self.downsample_factor]
            v_ds_gpu = v_gpu[::self.downsample_factor, ::self.downsample_factor]
            
            # 计算幅度（优化：使用square减少中间数组）
            magnitude_ds_gpu = cp.sqrt(cp.square(u_ds_gpu) + cp.square(v_ds_gpu))
            
            # 高斯滤波（使用CuPy的gaussian_filter）
            if self.denoise_sigma > 0:
                magnitude_ds_gpu = cp_gaussian_filter(magnitude_ds_gpu, sigma=self.denoise_sigma)
            
            # 优化：合并梯度计算，减少中间数组
            # 计算散度：du/dx + dv/dy
            du_dx_ds = cp.gradient(u_ds_gpu, axis=1)
            dv_dy_ds = cp.gradient(v_ds_gpu, axis=0)
            divergence_ds_gpu = du_dx_ds + dv_dy_ds
            
            # 对散度进行高斯滤波
            if self.denoise_sigma > 0:
                divergence_ds_gpu = cp_gaussian_filter(divergence_ds_gpu, sigma=self.denoise_sigma)
            
            # 优化：在GPU上完成插值，减少数据传输
            # 使用CuPy的zoom进行插值（更高效）
            from cupyx.scipy.ndimage import zoom as cp_zoom
            zoom_factor = (H / divergence_ds_gpu.shape[0], W / divergence_ds_gpu.shape[1])
            magnitude_gpu = cp_zoom(magnitude_ds_gpu, zoom_factor, order=3)  # 3 = cubic
            divergence_gpu = cp_zoom(divergence_ds_gpu, zoom_factor, order=3)
        else:
            # 计算幅度（优化：使用square减少中间数组）
            magnitude_gpu = cp.sqrt(cp.square(u_gpu) + cp.square(v_gpu))
            
            # 高斯滤波
            if self.denoise_sigma > 0:
                magnitude_gpu = cp_gaussian_filter(magnitude_gpu, sigma=self.denoise_sigma)
            
            # 优化：合并梯度计算
            du_dx = cp.gradient(u_gpu, axis=1)
            dv_dy = cp.gradient(v_gpu, axis=0)
            divergence_gpu = du_dx + dv_dy
            
            # 对散度进行高斯滤波
            if self.denoise_sigma > 0:
                divergence_gpu = cp_gaussian_filter(divergence_gpu, sigma=self.denoise_sigma)
        
        # 优化：在GPU上完成掩码操作，减少数据传输
        # 只保留大于阈值的散度值，其他设为0
        divergence_gpu = cp.where(divergence_gpu > self.divergence_threshold, divergence_gpu, 0)
        
        # 在GPU上创建掩码
        valid_mask_gpu = magnitude_gpu > self.flow_threshold
        expansion_mask_gpu = (valid_mask_gpu & (divergence_gpu > self.expansion_threshold))
        contraction_mask_gpu = (valid_mask_gpu & (divergence_gpu < -self.expansion_threshold))
        
        # 优化：一次性传输所有需要的数据回CPU
        magnitude = cp.asnumpy(magnitude_gpu)
        divergence = cp.asnumpy(divergence_gpu)
        expansion_mask = cp.asnumpy(expansion_mask_gpu)
        contraction_mask = cp.asnumpy(contraction_mask_gpu)
        
        # 初始化结果字典
        result = self._create_foe_result(expansion_mask, contraction_mask, divergence)
        
        # 寻找扩张区域轮廓（在CPU上进行，因为OpenCV的轮廓检测需要CPU）
        # 优化：缓存形态学操作的核
        # Convert kernel_size to tuple for use as dictionary key (lists are not hashable)
        kernel_key = tuple(self.kernel_size) if isinstance(self.kernel_size, (list, np.ndarray)) else self.kernel_size
        if kernel_key not in _morphology_kernels:
            _morphology_kernels[kernel_key] = np.ones(self.kernel_size, np.uint8)
        
        expansion_binary = expansion_mask.astype(np.uint8) * 255
        if np.any(expansion_binary):
            expansion_binary_closed = cv2.morphologyEx(expansion_binary, cv2.MORPH_CLOSE, _morphology_kernels[kernel_key])
            contours, _ = cv2.findContours(expansion_binary_closed, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
            
            if contours:
                result['centroid_found'] = True
                if return_all_contours:
                    result['_all_contours'] = contours
                
                largest_contour = max(contours, key=cv2.contourArea)
                result['contour'] = largest_contour
                
                M = cv2.moments(largest_contour)
                if M["m00"] > 0:
                    result['centroid'] = (int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"]))
        
        self.foe_results = result
        return self.foe_results

    def compute_multiple_foe(self,flow_field):
        """
        计算光流场的多个扩张焦点(Foci of Expansion)。
        
        Args:
            flow_field: 光流场，形状为 (H, W, 2) 的 numpy 数组
            flow_threshold: 最小光流幅度阈值，低于此值的点不参与计算
            expansion_threshold: 扩张判定阈值，用于区分扩张和收缩区域
            downsample_factor: 下采样因子，用于平滑计算（默认10）
            min_contour_area: 最小轮廓面积阈值，小于此值的轮廓会被忽略
            denoise_sigma: 降噪高斯滤波的标准差（默认1.5，设为0禁用降噪）
            kernel_size: 形态学操作的核大小
            
        Returns:
            dict: 包含多个扩张焦点信息的字典，包括：
                - 'foe_points': 扩张焦点坐标列表 [(x1, y1), (x2, y2), ...]
                - 'expansion_regions': 扩张区域轮廓列表
                - 'expansion_mask': 扩张区域二值掩码
                - 'contraction_mask': 收缩区域二值掩码
                - 'divergence_map': 散度图
        """
        # 调用 compute_foe_gpu 获取基础数据和所有轮廓
        foe_result = self.compute_foe_gpu(flow_field, return_all_contours=True)
        
        # 构建多焦点结果，复用基础数据
        result = self._create_foe_result(
            expansion_mask=foe_result['expansion_mask'],
            contraction_mask=foe_result['contraction_mask'],
            divergence_map=foe_result['divergence_map'],
            centroid_found=foe_result['centroid_found']
        )
        
        # 处理所有轮廓，过滤并计算焦点
        if '_all_contours' in foe_result:
            for contour in foe_result['_all_contours']:
                # 过滤掉太小的轮廓
                if cv2.contourArea(contour) >= self.min_contour_area:
                    # 计算轮廓的质心作为焦点
                    M = cv2.moments(contour)
                    if M["m00"] > 0:
                        cx = int(M["m10"] / M["m00"])
                        cy = int(M["m01"] / M["m00"])
                        result['centroid'].append((cx, cy))
                        if 'contour' not in result:
                            result['contour'] = []
                        result['contour'].append(contour)
        
        self.foe_results = result
        return self.foe_results