"""
点云降采样模块
支持多种降采样方法，包括均匀降采样、随机降采样、体素降采样和基于距离的分级降采样

核心设计：
- 持久化：PCD格式（Open3D tensor，支持intensity字段）
- 中间传输：NumPy数组 Nx4 (x, y, z, intensity)
- 强度信息：全程保持原始值，不归一化
- 动态配置：每次降采样前重新读取配置文件
"""

import open3d as o3d
import numpy as np
import yaml
from pathlib import Path
from typing import Optional, Union

from lidar_manager.logger import get_module_logger

logger = get_module_logger(__name__)


class PointCloudDownsampler:
    """
    点云降采样器

    内部统一使用 NumPy 数组格式：
    - Nx3: (x, y, z)
    - Nx4: (x, y, z, intensity) - 强度保持原始值，不归一化

    支持的降采样方法:
    - uniform: 均匀降采样（推荐，适合连续降采样）
    - random: 随机降采样（快速，但随机性强）
    - voxel: 体素降采样（保持几何特征）
    - distance_based: 基于距离的分级降采样（适合激光雷达数据，渐进式降采样）

    核心特性：
    - 每次降采样前自动重新读取配置文件
    - 支持热更新配置，无需重启程序
    """

    def __init__(self, config_path: str = 'cfg/downsample_config.yaml'):
        """
        初始化降采样器

        参数:
            config_path: 配置文件路径（支持相对路径和绝对路径）
        """
        # 保存配置文件路径（用于动态重载）
        self.config_path = Path(config_path)

        # 首次加载配置（用于初始化信息显示）
        self.config = self._load_config_file()

        logger.info(f"降采样器初始化完成")
        logger.info(f"配置文件路径: {self.config_path.absolute()}")
        logger.info(f"默认方法: {self.config.get('default_method', 'uniform')}")

    def _load_config_file(self) -> dict:
        """
        从文件加载配置（内部方法）

        返回:
            配置字典
        """
        try:
            if not self.config_path.exists():
                logger.warning(f"配置文件不存在 ({self.config_path})，使用默认配置")
                return self._get_default_config()

            with open(self.config_path, 'r', encoding='utf-8') as f:
                config = yaml.safe_load(f)

            # 验证配置文件不为空
            if config is None:
                logger.warning(f"配置文件为空 ({self.config_path})，使用默认配置")
                return self._get_default_config()

            return config

        except Exception as e:
            logger.warning(f"加载配置文件失败 ({e})，使用默认配置", exc_info=True)
            return self._get_default_config()

    def _reload_config(self) -> dict:
        """
        重新加载配置文件（每次降采样前调用）

        返回:
            最新的配置字典
        """
        logger.debug(f"重新加载配置文件: {self.config_path.name}")
        config = self._load_config_file()

        # 显示关键配置信息
        default_method = config.get('default_method', 'uniform')
        logger.debug(f"默认方法: {default_method}")

        # 显示当前使用的参数
        if default_method == 'uniform':
            k = config.get('uniform', {}).get('k', 10)
            logger.debug(f"uniform.k = {k}")
        elif default_method == 'voxel':
            voxel_size = config.get('voxel', {}).get('voxel_size', 0.1)
            adaptive = config.get('voxel', {}).get('adaptive', True)
            logger.debug(f"voxel.voxel_size = {voxel_size}, adaptive = {adaptive}")
        elif default_method == 'random':
            ratio = config.get('random', {}).get('ratio', 0.5)
            logger.debug(f"random.ratio = {ratio}")
        elif default_method == 'distance_based':
            rules = config.get('distance_based', {}).get('distance_rules', [])
            logger.debug(f"distance_rules: {len(rules)} 个分段")

        return config

    def _get_default_config(self) -> dict:
        """
        获取默认配置

        返回:
            默认配置字典
        """
        return {
            'default_method': 'uniform',
            'uniform': {'k': 10},
            'random': {'ratio': 0.5, 'seed': 42},
            'voxel': {
                'voxel_size': 0.1,
                'adaptive': True,
                'adaptive_rules': [
                    [1000000, 0.05],
                    [500000, 0.10],
                    [200000, 0.15],
                    [100000, 0.20],
                    [50000, 0.30],
                    [0, 0.50]
                ]
            },
            'distance_based': {
                'camera_position': None,
                'distance_rules': [
                    [5.0, 0.95],
                    [10.0, 0.90],
                    [15.0, 0.85],
                    [20.0, 0.75],
                    [30.0, 0.60],
                    [40.0, 0.45],
                    [50.0, 0.30],
                    [70.0, 0.20],
                    [100.0, 0.10],
                    [999999, 0.05]
                ],
                'visualization': {
                    'show_distance_distribution': True,
                    'show_zone_info': True
                }
            },
            'continuous_downsampling': {
                'show_info': True,
                'min_points_warning': 1000
            }
        }

    # ========== NumPy <-> Open3D 转换工具 ==========

    def numpy_to_o3d(self, points: np.ndarray) -> o3d.geometry.PointCloud:
        """
        将NumPy数组转换为Open3D点云（用于降采样算法）

        参数:
            points: Nx3 或 Nx4 的NumPy数组 (x, y, z) 或 (x, y, z, intensity)

        返回:
            Open3D点云对象（强度信息临时存储在颜色通道R中）
        """
        pcd = o3d.geometry.PointCloud()
        pcd.points = o3d.utility.Vector3dVector(points[:, :3])

        # 如果有强度信息，临时存储在颜色的R通道（保持原始值）
        if points.shape[1] > 3:
            intensities = points[:, 3]
            # 不归一化，直接存储原始强度值
            colors = np.zeros((len(points), 3))
            colors[:, 0] = intensities  # R通道存储原始强度
            pcd.colors = o3d.utility.Vector3dVector(colors)

        return pcd

    def o3d_to_numpy(self, pcd: o3d.geometry.PointCloud,
                     has_intensity: bool = True) -> np.ndarray:
        """
        将Open3D点云转换为NumPy数组

        参数:
            pcd: Open3D点云对象
            has_intensity: 是否包含强度信息（从颜色R通道读取）

        返回:
            Nx3 或 Nx4 的NumPy数组
        """
        points = np.asarray(pcd.points)

        if has_intensity and pcd.has_colors():
            colors = np.asarray(pcd.colors)
            intensities = colors[:, 0]  # 从R通道恢复原始强度
            return np.column_stack((points, intensities))
        else:
            return points

    # ========== 持久化方法（PCD格式，保留原始强度）==========

    def save_point_cloud(self, points: np.ndarray, output_path: str):
        """
        保存点云为PCD格式（使用Open3D tensor，保留原始强度值）

        参数:
            points: Nx3 或 Nx4 的NumPy数组 (x, y, z) 或 (x, y, z, intensity)
            output_path: 输出文件路径 (.pcd)
        """
        # 使用 Open3D tensor point cloud 保存（支持intensity字段）
        pcd = o3d.t.geometry.PointCloud()
        pcd.point.positions = o3d.core.Tensor(points[:, :3], dtype=o3d.core.Dtype.Float32)

        # 保存强度信息（原始值，不归一化）
        if points.shape[1] > 3:
            intensities = points[:, 3:4]  # 保持2D形状 (N, 1)
            pcd.point.intensity = o3d.core.Tensor(intensities, dtype=o3d.core.Dtype.Float32)

        # 保存
        o3d.t.io.write_point_cloud(output_path, pcd)

        logger.info(f"点云已保存: {output_path}")
        logger.info(f"   点数: {len(points):,}")
        if points.shape[1] > 3:
            logger.debug(f"   强度范围: [{points[:, 3].min():.2f}, {points[:, 3].max():.2f}] (原始值)")

    def load_point_cloud(self, input_path: str) -> np.ndarray:
        """
        从PCD文件加载点云（使用Open3D tensor，保留原始强度值）

        参数:
            input_path: 输入文件路径 (.pcd)

        返回:
            NumPy数组 (Nx3或Nx4)
        """
        # 使用 Open3D tensor point cloud 加载
        pcd = o3d.t.io.read_point_cloud(input_path)

        points = pcd.point.positions.numpy()

        # 读取强度信息（如果存在）
        if "intensity" in pcd.point:
            intensities = pcd.point.intensity.numpy().flatten()
            points_with_intensity = np.column_stack((points, intensities))

            logger.debug(f"点云已加载: {input_path}, 点数: {len(points):,}")
            logger.debug(f"   强度范围: [{intensities.min():.2f}, {intensities.max():.2f}] (原始值)")

            return points_with_intensity
        else:
            logger.debug(f"点云已加载（无强度信息）: {input_path}, 点数: {len(points):,}")
            return points

    # ========== 降采样核心方法（接收/返回NumPy数组）==========

    def downsample(self,
                   points: np.ndarray,
                   method: Optional[str] = None,
                   reload_config: bool = True,
                   **kwargs) -> np.ndarray:
        """
        对点云执行降采样（NumPy接口）

        参数:
            points: 输入点云 (NumPy数组，Nx3或Nx4)
            method: 降采样方法 ('uniform', 'random', 'voxel', 'distance_based')
                   如果为None，使用配置文件中的默认方法
            reload_config: 是否在降采样前重新加载配置文件（默认True）
            **kwargs: 方法特定的参数（会覆盖配置文件中的值）

        返回:
            降采样后的点云 (NumPy数组，形状与输入一致)
        """
        if points is None or len(points) == 0:
            logger.error("输入点云为空")
            return points

        # 🔄 动态重新加载配置文件
        if reload_config:
            config = self._reload_config()
        else:
            config = self.config

        # 使用默认方法（如果未指定）
        if method is None:
            method = config.get('default_method', 'uniform')

        # 检查方法是否有效
        valid_methods = ['uniform', 'random', 'voxel', 'distance_based']
        if method not in valid_methods:
            logger.error(f"不支持的降采样方法 '{method}'，使用默认方法 'uniform'")
            method = 'uniform'

        # 记录原始点数
        original_points = len(points)
        has_intensity = points.shape[1] > 3

        # 获取配置（从最新加载的config中）
        show_info = config.get('continuous_downsampling', {}).get('show_info', True)
        min_points_warning = config.get('continuous_downsampling', {}).get('min_points_warning', 1000)

        # 距离降采样直接处理NumPy数组
        if method == 'distance_based':
            downsampled_points = self._distance_based_downsample(points, config, **kwargs)
        else:
            # 转换为Open3D格式（用于其他降采样算法）
            pcd = self.numpy_to_o3d(points)

            # 执行降采样
            if method == 'uniform':
                downsampled_pcd = self._uniform_downsample(pcd, config, **kwargs)
            elif method == 'random':
                downsampled_pcd = self._random_downsample(pcd, config, **kwargs)
            elif method == 'voxel':
                downsampled_pcd = self._voxel_downsample(pcd, config, **kwargs)

            # 转换回NumPy格式
            downsampled_points = self.o3d_to_numpy(downsampled_pcd, has_intensity=has_intensity)

        # 记录降采样后点数
        final_points = len(downsampled_points)
        reduction_ratio = (1 - final_points / original_points) * 100 if original_points > 0 else 0

        # 显示信息
        if show_info:
            logger.info("=" * 60)
            logger.info(f"降采样方法: {method.upper()}")
            logger.info(f"原始点数: {original_points:,}")
            logger.info(f"降采样后点数: {final_points:,}")
            logger.info(f"减少: {reduction_ratio:.2f}%")
            if has_intensity and final_points > 0:
                logger.debug(f"强度范围: [{downsampled_points[:, 3].min():.2f}, {downsampled_points[:, 3].max():.2f}] (原始值)")
            logger.info("=" * 60)

        # 点数过少警告
        if final_points < min_points_warning:
            logger.warning(f"点数较少 ({final_points:,} 点)，可能影响后续处理")

        return downsampled_points

    def _uniform_downsample(self,
                            point_cloud: o3d.geometry.PointCloud,
                            config: dict,
                            k: Optional[int] = None) -> o3d.geometry.PointCloud:
        """
        均匀降采样：每k个点保留1个

        参数:
            point_cloud: 输入点云
            config: 配置字典
            k: 采样间隔（None则使用配置文件中的值）

        返回:
            降采样后的点云
        """
        # 获取k值
        if k is None:
            k = config.get('uniform', {}).get('k', 10)

        if k < 1:
            logger.warning(f"k值无效 ({k})，设置为1")
            k = 1

        # 执行均匀降采样
        downsampled_cloud = point_cloud.uniform_down_sample(every_k_points=k)

        return downsampled_cloud

    def _random_downsample(self,
                           point_cloud: o3d.geometry.PointCloud,
                           config: dict,
                           ratio: Optional[float] = None,
                           seed: Optional[int] = None) -> o3d.geometry.PointCloud:
        """
        随机降采样：随机保留指定比例的点

        参数:
            point_cloud: 输入点云
            config: 配置字典
            ratio: 保留比例 0.0-1.0（None则使用配置文件中的值）
            seed: 随机种子（None则使用配置文件中的值）

        返回:
            降采样后的点云
        """
        # 获取参数
        if ratio is None:
            ratio = config.get('random', {}).get('ratio', 0.5)
        if seed is None:
            seed = config.get('random', {}).get('seed', None)

        # 检查ratio范围
        if not 0.0 < ratio <= 1.0:
            logger.warning(f"ratio值无效 ({ratio})，设置为0.5")
            ratio = 0.5

        # 设置随机种子
        if seed is not None:
            np.random.seed(seed)

        # 执行随机降采样
        downsampled_cloud = point_cloud.random_down_sample(sampling_ratio=ratio)

        return downsampled_cloud

    def _voxel_downsample(self,
                          point_cloud: o3d.geometry.PointCloud,
                          config: dict,
                          voxel_size: Optional[float] = None,
                          adaptive: Optional[bool] = None) -> o3d.geometry.PointCloud:
        """
        体素降采样：基于3D网格的降采样

        参数:
            point_cloud: 输入点云
            config: 配置字典
            voxel_size: 体素大小（None则使用配置文件中的值或自适应）
            adaptive: 是否使用自适应体素大小（None则使用配置文件中的值）

        返回:
            降采样后的点云
        """
        # 获取参数
        if adaptive is None:
            adaptive = config.get('voxel', {}).get('adaptive', True)

        # 确定体素大小
        if voxel_size is None:
            if adaptive:
                voxel_size = self._get_adaptive_voxel_size(len(point_cloud.points), config)
            else:
                voxel_size = config.get('voxel', {}).get('voxel_size', 0.1)

        # 检查voxel_size
        if voxel_size <= 0:
            logger.warning(f"voxel_size无效 ({voxel_size})，设置为0.1")
            voxel_size = 0.1

        if adaptive:
            logger.debug(f"自适应体素大小: {voxel_size}m (基于点数: {len(point_cloud.points):,})")

        # 执行体素降采样
        downsampled_cloud = point_cloud.voxel_down_sample(voxel_size=voxel_size)

        return downsampled_cloud

    def _get_adaptive_voxel_size(self, num_points: int, config: dict) -> float:
        """
        根据点云数量自适应选择体素大小

        参数:
            num_points: 点云数量
            config: 配置字典

        返回:
            推荐的体素大小
        """
        adaptive_rules = config.get('voxel', {}).get('adaptive_rules', [])

        if not adaptive_rules:
            # 没有规则，使用默认值
            return config.get('voxel', {}).get('voxel_size', 0.1)

        # 按规则查找合适的体素大小
        for threshold, voxel_size in adaptive_rules:
            if num_points >= threshold:
                return voxel_size

        # 如果没有匹配，使用最后一个规则
        return adaptive_rules[-1][1]

    def _distance_based_downsample(self,
                                   points: np.ndarray,
                                   config: dict,
                                   camera_position: Optional[np.ndarray] = None,
                                   distance_rules: Optional[list] = None,
                                   show_zone_info: Optional[bool] = None,
                                   show_distance_distribution: Optional[bool] = None) -> np.ndarray:
        """
        基于距离的分级降采样（渐进式，使用均匀采样按比例保留）

        参数:
            points: 输入点云 (Nx3 或 Nx4)
            config: 配置字典
            camera_position: 相机在LiDAR坐标系中的位置 [x, y, z]
                            - 如果为None，先从配置文件读取
                            - 如果配置也为None，则使用点云中心（不推荐）
            distance_rules: 距离规则列表 [[最大距离, 保留比例], ...]
            show_zone_info: 是否显示各区间信息
            show_distance_distribution: 是否显示距离分布

        返回:
            降采样后的点云
        """
        # 获取配置
        distance_config = config.get('distance_based', {})

        # 优先级：
        # 1. 函数参数传入的 camera_position（最高优先级，来自外参计算）
        # 2. 配置文件中的 camera_position
        # 3. 点云中心（兜底方案，不推荐）
        if camera_position is None:
            camera_position = distance_config.get('camera_position', None)

        if distance_rules is None:
            distance_rules = distance_config.get('distance_rules', [])

        viz_config = distance_config.get('visualization', {})
        if show_zone_info is None:
            show_zone_info = viz_config.get('show_zone_info', True)
        if show_distance_distribution is None:
            show_distance_distribution = viz_config.get('show_distance_distribution', True)

        # 确定相机位置
        if camera_position is None:
            # 兜底方案：使用点云中心作为相机位置
            camera_position = np.mean(points[:, :3], axis=0)
            logger.warning("未提供相机位置，使用点云中心（不推荐）")
            logger.debug(f"相机位置（点云中心）: [{camera_position[0]:.3f}, {camera_position[1]:.3f}, {camera_position[2]:.3f}]")
        else:
            camera_position = np.array(camera_position)
            logger.debug("使用外参计算的真实相机位置")
            logger.debug(f"相机位置: [{camera_position[0]:.3f}, {camera_position[1]:.3f}, {camera_position[2]:.3f}]")

        # 计算每个点到相机的距离
        distances = np.linalg.norm(points[:, :3] - camera_position, axis=1)

        # 显示距离分布
        if show_distance_distribution:
            logger.debug("距离分布统计:")
            logger.debug(f"      最小距离: {distances.min():.2f}m")
            logger.debug(f"      最大距离: {distances.max():.2f}m")
            logger.debug(f"      平均距离: {distances.mean():.2f}m")
            logger.debug(f"      中位数距离: {np.median(distances):.2f}m")

        # 按距离分区并降采样
        downsampled_points_list = []
        zone_info = []

        has_intensity = points.shape[1] > 3

        # 添加最小距离（从0开始）
        distance_thresholds = [0.0] + [rule[0] for rule in distance_rules]

        for i, (max_dist, keep_ratio) in enumerate(distance_rules):
            min_dist = distance_thresholds[i]

            # 选择当前距离区间的点
            mask = (distances >= min_dist) & (distances < max_dist)
            zone_points = points[mask]

            if len(zone_points) == 0:
                continue

            # 对当前区间进行均匀降采样（按比例）
            downsampled_zone = self._downsample_zone_by_ratio(zone_points, keep_ratio)

            downsampled_points_list.append(downsampled_zone)

            # 记录区间信息
            actual_ratio = len(downsampled_zone) / len(zone_points) if len(zone_points) > 0 else 0
            zone_info.append({
                'range': f"{min_dist:.1f}-{max_dist:.1f}m",
                'original': len(zone_points),
                'downsampled': len(downsampled_zone),
                'target_ratio': keep_ratio,
                'actual_ratio': actual_ratio
            })

        # 合并所有区间的点
        if len(downsampled_points_list) == 0:
            logger.warning("所有距离区间都为空")
            return points[:0]  # 返回空数组

        downsampled_points = np.vstack(downsampled_points_list)

        # 显示各区间信息
        if show_zone_info and len(zone_info) > 0:
            logger.debug("距离分段降采样统计:")
            logger.debug(f"   {'距离区间':<15} {'原始点数':<12} {'降采样后':<12} {'目标比例':<12} {'实际比例':<12}")
            logger.debug(f"   {'-' * 75}")
            for info in zone_info:
                logger.debug(f"   {info['range']:<15} {info['original']:<12,} {info['downsampled']:<12,} "
                            f"{info['target_ratio'] * 100:>10.1f}%  {info['actual_ratio'] * 100:>10.1f}%")
            logger.debug(f"   {'-' * 75}")
            total_original = sum(info['original'] for info in zone_info)
            total_downsampled = sum(info['downsampled'] for info in zone_info)
            total_ratio = total_downsampled / total_original if total_original > 0 else 0
            logger.debug(f"   {'总计':<15} {total_original:<12,} {total_downsampled:<12,} "
                        f"{'':>12}  {total_ratio * 100:>10.1f}%")

        return downsampled_points


    def _downsample_zone_by_ratio(self, points: np.ndarray, keep_ratio: float) -> np.ndarray:
        """
        对区间进行均匀降采样（按比例保留点）

        参数:
            points: 区间内的点云
            keep_ratio: 保留比例 (0.0-1.0)

        返回:
            降采样后的点云
        """
        if keep_ratio >= 1.0:
            return points

        if keep_ratio <= 0.0:
            # 保留至少1个点
            return points[:1] if len(points) > 0 else points

        # 计算目标点数
        n_target = max(1, int(len(points) * keep_ratio))

        # 计算采样间隔k（每k个点保留1个）
        k = max(1, int(len(points) / n_target))

        # 均匀采样
        indices = np.arange(0, len(points), k)

        # 如果采样点数超过目标，截断
        if len(indices) > n_target:
            indices = indices[:n_target]

        return points[indices]

    # ========== 信息获取 ==========

    def get_method_info(self, method: Optional[str] = None) -> dict:
        """
        获取降采样方法的信息

        参数:
            method: 方法名称（None则返回默认方法）

        返回:
            方法信息字典
        """
        # 获取最新配置
        config = self._reload_config()
        default_method = config.get('default_method', 'uniform')

        if method is None:
            method = default_method

        info = {
            'method': method,
            'is_default': method == default_method
        }

        if method == 'uniform':
            info['config'] = config.get('uniform', {})
            info['description'] = '均匀降采样：每k个点保留1个'
            info['suitable_for'] = '连续降采样，效果稳定可预测'
        elif method == 'random':
            info['config'] = config.get('random', {})
            info['description'] = '随机降采样：随机保留指定比例的点'
            info['suitable_for'] = '快速预览，但结果有随机性'
        elif method == 'voxel':
            info['config'] = config.get('voxel', {})
            info['description'] = '体素降采样：基于3D网格，保持几何特征'
            info['suitable_for'] = '保持点云几何结构，适合精细处理'
        elif method == 'distance_based':
            info['config'] = config.get('distance_based', {})
            info['description'] = '基于距离的渐进式降采样：近处保留更多细节，远处逐渐降采样'
            info['suitable_for'] = '激光雷达数据，渐进式降采样，平滑过渡'
        else:
            info['error'] = f"不支持的方法: {method}"

        return info

    def print_all_methods_info(self):
        """打印所有降采样方法的信息"""
        logger.info("=" * 60)
        logger.info("支持的降采样方法:")
        logger.info("=" * 60)

        for method in ['uniform', 'random', 'voxel', 'distance_based']:
            info = self.get_method_info(method)
            default_marker = " ⭐ (默认)" if info['is_default'] else ""
            logger.info(f"【{method.upper()}{default_marker}】")
            logger.info(f"  描述: {info['description']}")
            logger.info(f"  适用: {info['suitable_for']}")
            if 'config' in info:
                logger.info(f"  配置: {info['config']}")

        logger.info("=" * 60)