colmap 0.1.2

//! 点云融合模块
//!
//! 实现多视角深度图融合算法，生成一致的3D点云

use crate::core::{
    Image, CameraPose, CameraIntrinsics, Point2, Point3, Result, ColmapError
};
use crate::mvs::stereo::DepthMap;
use crate::mvs::depth::{NormalMap, CostMap};
use nalgebra::{Vector3, Point3 as NalgebraPoint3};
use std::collections::{HashMap, HashSet};

/// 点云融合器
#[derive(Debug)]
pub struct PointCloudFusion {
    /// 融合配置
    config: FusionConfig,
}

/// 融合配置
#[derive(Debug, Clone)]
pub struct FusionConfig {
    /// 最小视图数
    pub min_num_views: usize,
    /// 最大重投影误差
    pub max_reprojection_error: f32,
    /// 最小深度
    pub min_depth: f32,
    /// 最大深度
    pub max_depth: f32,
    /// 深度一致性阈值
    pub depth_consistency_threshold: f32,
    /// 法向量一致性阈值
    pub normal_consistency_threshold: f32,
    /// 体素大小
    pub voxel_size: f32,
    /// 最大代价阈值
    pub max_cost_threshold: f32,
}

impl Default for FusionConfig {
    fn default() -> Self {
        Self {
            min_num_views: 3,
            max_reprojection_error: 2.0,
            min_depth: 0.1,
            max_depth: 100.0,
            depth_consistency_threshold: 0.01,
            normal_consistency_threshold: 0.1,
            voxel_size: 0.01,
            max_cost_threshold: 0.3,
        }
    }
}

/// 融合点
#[derive(Debug, Clone)]
pub struct FusedPoint {
    /// 3D坐标
    pub position: Point3,
    /// 法向量
    pub normal: Vector3<f32>,
    /// 颜色 (RGB)
    pub color: [u8; 3],
    /// 置信度
    pub confidence: f32,
    /// 支持视图数
    pub num_views: usize,
}

/// 融合结果
#[derive(Debug)]
pub struct FusionResult {
    /// 融合点云
    pub points: Vec<FusedPoint>,
    /// 统计信息
    pub statistics: FusionStatistics,
}

/// 融合统计信息
#[derive(Debug)]
pub struct FusionStatistics {
    /// 输入深度图数量
    pub num_input_depth_maps: usize,
    /// 原始点数
    pub num_raw_points: usize,
    /// 融合后点数
    pub num_fused_points: usize,
    /// 过滤掉的点数
    pub num_filtered_points: usize,
}

/// 深度图视图
#[derive(Debug, Clone)]
pub struct DepthMapView {
    /// 深度图
    pub depth_map: DepthMap,
    /// 法向量图
    pub normal_map: Option<NormalMap>,
    /// 代价图
    pub cost_map: Option<CostMap>,
    /// 原始图像
    pub image: Image,
}

/// 体素网格
#[derive(Debug)]
struct VoxelGrid {
    /// 体素大小
    voxel_size: f32,
    /// 体素映射
    voxels: HashMap<(i32, i32, i32), Vec<usize>>,
}

impl PointCloudFusion {
    /// 创建新的点云融合器
    pub fn new(config: FusionConfig) -> Self {
        Self { config }
    }

    /// 融合多个深度图
    pub fn fuse_depth_maps(
        &self,
        depth_map_views: &[DepthMapView],
    ) -> Result<FusionResult> {
        if depth_map_views.is_empty() {
            return Err(ColmapError::MvsReconstruction(
                "No depth maps provided for fusion".to_string(),
            ));
        }

        println!("Starting point cloud fusion with {} depth maps", depth_map_views.len());

        // 1. 从所有深度图生成候选点
        let candidate_points = self.generate_candidate_points(depth_map_views)?;
        println!("Generated {} candidate points", candidate_points.len());

        // 2. 一致性检查和过滤
        let consistent_points = self.consistency_check(
            &candidate_points,
            depth_map_views,
        )?;
        println!("Found {} consistent points", consistent_points.len());

        // 3. 体素化和去重
        let voxelized_points = self.voxelize_points(&consistent_points)?;
        println!("Voxelized to {} points", voxelized_points.len());

        // 4. 颜色插值
        let colored_points = self.interpolate_colors(
            &voxelized_points,
            depth_map_views,
        )?;

        // 5. 构建统计信息
        let statistics = FusionStatistics {
            num_input_depth_maps: depth_map_views.len(),
            num_raw_points: candidate_points.len(),
            num_fused_points: colored_points.len(),
            num_filtered_points: candidate_points.len() - colored_points.len(),
        };

        Ok(FusionResult {
            points: colored_points,
            statistics,
        })
    }

    /// 生成候选点
    fn generate_candidate_points(
        &self,
        depth_map_views: &[DepthMapView],
    ) -> Result<Vec<CandidatePoint>> {
        let mut candidate_points = Vec::new();

        for (view_id, view) in depth_map_views.iter().enumerate() {
            let depth_map = &view.depth_map;
            
            for y in 0..depth_map.height {
                for x in 0..depth_map.width {
                    if let Some(depth) = depth_map.get_depth(x, y) {
                        // 检查深度范围
                        if depth < self.config.min_depth || depth > self.config.max_depth {
                            continue;
                        }

                        // 检查代价阈值
                        if let Some(cost_map) = &view.cost_map {
                            let cost_idx = (y * depth_map.width + x) as usize;
                            if cost_idx < cost_map.data.len() && 
                               cost_map.data[cost_idx] > self.config.max_cost_threshold {
                                continue;
                            }
                        }

                        // 反投影到3D空间
                        let point_3d = self.unproject_pixel(
                            x as f32, y as f32, depth,
                            &depth_map.intrinsics,
                            &depth_map.pose,
                        )?;

                        // 获取法向量
                        let normal = if let Some(normal_map) = &view.normal_map {
                            let normal_idx = (y * depth_map.width + x) as usize;
                            if normal_idx < normal_map.data.len() {
                                normal_map.data[normal_idx]
                            } else {
                                Vector3::new(0.0, 0.0, 1.0)
                            }
                        } else {
                            Vector3::new(0.0, 0.0, 1.0)
                        };

                        candidate_points.push(CandidatePoint {
                            position: point_3d,
                            normal,
                            source_view_id: view_id,
                            pixel_x: x,
                            pixel_y: y,
                            depth,
                            confidence: 1.0,
                        });
                    }
                }
            }
        }

        Ok(candidate_points)
    }

    /// 一致性检查
    fn consistency_check(
        &self,
        candidate_points: &[CandidatePoint],
        depth_map_views: &[DepthMapView],
    ) -> Result<Vec<ConsistentPoint>> {
        let mut consistent_points = Vec::new();

        for candidate in candidate_points {
            let mut supporting_views = Vec::new();
            let mut total_confidence = 0.0;

            // 检查在其他视图中的一致性
            for (view_id, view) in depth_map_views.iter().enumerate() {
                if view_id == candidate.source_view_id {
                    supporting_views.push(view_id);
                    total_confidence += candidate.confidence;
                    continue;
                }

                // 投影到当前视图
                if let Ok(projected) = self.project_point(
                    &candidate.position,
                    &view.depth_map.intrinsics,
                    &view.depth_map.pose,
                ) {
                    let px = projected.x.round() as u32;
                    let py = projected.y.round() as u32;

                    // 检查是否在图像范围内
                    if px < view.depth_map.width && py < view.depth_map.height
                        && let Some(view_depth) = view.depth_map.get_depth(px, py) {
                            // 计算预期深度
                            let expected_depth = self.compute_expected_depth(
                                &candidate.position,
                                &view.depth_map.pose,
                            );

                            // 深度一致性检查
                            let depth_diff = (view_depth - expected_depth).abs();
                            let relative_threshold = self.config.depth_consistency_threshold * expected_depth;
                            
                            if depth_diff < relative_threshold {
                                supporting_views.push(view_id);
                                total_confidence += 1.0;

                                // 法向量一致性检查
                                if let Some(normal_map) = &view.normal_map {
                                    let normal_idx = (py * view.depth_map.width + px) as usize;
                                    if normal_idx < normal_map.data.len() {
                                        let view_normal = normal_map.data[normal_idx];
                                        let normal_similarity = candidate.normal.dot(&view_normal);
                                        if normal_similarity > self.config.normal_consistency_threshold {
                                            total_confidence += 0.5;
                                        }
                                    }
                                }
                            }
                        }
                }
            }

            // 检查是否满足最小视图数要求
            if supporting_views.len() >= self.config.min_num_views {
                let supporting_views_len = supporting_views.len();
                let confidence = total_confidence / supporting_views_len as f32;
                consistent_points.push(ConsistentPoint {
                    position: candidate.position,
                    normal: candidate.normal,
                    supporting_views,
                    confidence,
                });
            }
        }

        Ok(consistent_points)
    }

    /// 体素化点云
    fn voxelize_points(
        &self,
        consistent_points: &[ConsistentPoint],
    ) -> Result<Vec<VoxelizedPoint>> {
        let mut voxel_grid = VoxelGrid::new(self.config.voxel_size);
        
        // 将点分配到体素
        for (point_id, point) in consistent_points.iter().enumerate() {
            let voxel_key = voxel_grid.get_voxel_key(&point.position);
            voxel_grid.add_point(voxel_key, point_id);
        }

        let mut voxelized_points = Vec::new();

        // 合并每个体素中的点
        for point_ids in voxel_grid.voxels.values() {
            if !point_ids.is_empty() {
                let merged_point = self.merge_points_in_voxel(
                    point_ids,
                    consistent_points,
                )?;
                voxelized_points.push(merged_point);
            }
        }

        Ok(voxelized_points)
    }

    /// 合并体素中的点
    fn merge_points_in_voxel(
        &self,
        point_ids: &[usize],
        consistent_points: &[ConsistentPoint],
    ) -> Result<VoxelizedPoint> {
        let mut position_sum = Vector3::new(0.0, 0.0, 0.0);
        let mut normal_sum = Vector3::new(0.0, 0.0, 0.0);
        let mut confidence_sum = 0.0;
        let mut all_supporting_views = HashSet::new();

        for &point_id in point_ids {
            let point = &consistent_points[point_id];
            let weight = point.confidence;

            position_sum += Vector3::new(
                point.position.x as f32,
                point.position.y as f32,
                point.position.z as f32,
            ) * weight;
            normal_sum += point.normal * weight;
            confidence_sum += weight;

            for &view_id in &point.supporting_views {
                all_supporting_views.insert(view_id);
            }
        }

        let merged_position = if confidence_sum > 0.0 {
            position_sum / confidence_sum
        } else {
            Vector3::new(0.0, 0.0, 0.0)
        };

        let merged_normal = if normal_sum.norm() > 0.0 {
            normal_sum.normalize()
        } else {
            Vector3::new(0.0, 0.0, 1.0)
        };

        Ok(VoxelizedPoint {
            position: Point3::new(
                merged_position.x as f64,
                merged_position.y as f64,
                merged_position.z as f64,
            ),
            normal: merged_normal,
            confidence: confidence_sum / point_ids.len() as f32,
            supporting_views: all_supporting_views.into_iter().collect(),
        })
    }

    /// 颜色插值
    fn interpolate_colors(
        &self,
        voxelized_points: &[VoxelizedPoint],
        depth_map_views: &[DepthMapView],
    ) -> Result<Vec<FusedPoint>> {
        let mut colored_points = Vec::new();

        for point in voxelized_points {
            let mut color_sum = [0.0f32; 3];
            let mut weight_sum = 0.0;

            // 从支持视图中插值颜色
            for &view_id in &point.supporting_views {
                if view_id < depth_map_views.len() {
                    let view = &depth_map_views[view_id];
                    
                    if let Ok(projected) = self.project_point(
                        &point.position,
                        &view.depth_map.intrinsics,
                        &view.depth_map.pose,
                    ) {
                        let color = self.sample_image_color(&view.image, projected.x, projected.y);
                        let weight = 1.0; // 可以基于距离或角度调整权重
                        
                        color_sum[0] += color[0] as f32 * weight;
                        color_sum[1] += color[1] as f32 * weight;
                        color_sum[2] += color[2] as f32 * weight;
                        weight_sum += weight;
                    }
                }
            }

            let final_color = if weight_sum > 0.0 {
                [
                    (color_sum[0] / weight_sum).clamp(0.0, 255.0) as u8,
                    (color_sum[1] / weight_sum).clamp(0.0, 255.0) as u8,
                    (color_sum[2] / weight_sum).clamp(0.0, 255.0) as u8,
                ]
            } else {
                [128, 128, 128] // 默认灰色
            };

            colored_points.push(FusedPoint {
                position: point.position,
                normal: point.normal,
                color: final_color,
                confidence: point.confidence,
                num_views: point.supporting_views.len(),
            });
        }

        Ok(colored_points)
    }

    /// 反投影像素到3D空间
    fn unproject_pixel(
        &self,
        x: f32,
        y: f32,
        depth: f32,
        intrinsics: &CameraIntrinsics,
        pose: &CameraPose,
    ) -> Result<Point3> {
        let x_norm = (x - intrinsics.principal_point.0 as f32) / intrinsics.focal_length.0 as f32;
            let y_norm = (y - intrinsics.principal_point.1 as f32) / intrinsics.focal_length.1 as f32;
        
        let point_camera = NalgebraPoint3::new(
            x_norm * depth,
            y_norm * depth,
            depth,
        );

        let point_camera_f64 = nalgebra::Point3::<f64>::new(
            point_camera.x as f64,
            point_camera.y as f64,
            point_camera.z as f64,
        );
        let point_world_coords = pose.rotation.inverse() *
            (point_camera_f64.coords - pose.translation);
        let point_world = nalgebra::Point3::from(point_world_coords);

        Ok(Point3::new(point_world.x, point_world.y, point_world.z))
    }

    /// 投影3D点到图像
    fn project_point(
        &self,
        point_3d: &Point3,
        intrinsics: &CameraIntrinsics,
        pose: &CameraPose,
    ) -> Result<Point2> {
        let point_world = NalgebraPoint3::new(point_3d.x, point_3d.y, point_3d.z);
        let point_camera = pose.rotation * point_world + pose.translation;

        if point_camera.z <= 0.0 {
            return Err(ColmapError::MvsReconstruction(
                "Point is behind camera".to_string(),
            ));
        }

        let x = intrinsics.focal_length.0 * (point_camera.x / point_camera.z) + intrinsics.principal_point.0;
        let y = intrinsics.focal_length.1 * (point_camera.y / point_camera.z) + intrinsics.principal_point.1;

        Ok(Point2::new(x, y))
    }

    /// 计算预期深度
    fn compute_expected_depth(
        &self,
        point_3d: &Point3,
        pose: &CameraPose,
    ) -> f32 {
        let point_world = NalgebraPoint3::new(point_3d.x, point_3d.y, point_3d.z);
        let point_camera = pose.rotation * point_world + pose.translation;
        point_camera.z as f32
    }

    /// 采样图像颜色
    fn sample_image_color(&self, image: &Image, x: f64, y: f64) -> [u8; 3] {
        // 简化实现：返回固定颜色
        // 在实际实现中，这里应该进行双线性插值
        let px = x.clamp(0.0, (image.size.0 - 1) as f64) as u32;
        let py = y.clamp(0.0, (image.size.1 - 1) as f64) as u32;
        
        // 占位符：假设图像有RGB数据
        [128, 128, 128]
    }
}

/// 候选点
#[derive(Debug, Clone)]
struct CandidatePoint {
    position: Point3,
    normal: Vector3<f32>,
    source_view_id: usize,
    pixel_x: u32,
    pixel_y: u32,
    depth: f32,
    confidence: f32,
}

/// 一致性点
#[derive(Debug, Clone)]
struct ConsistentPoint {
    position: Point3,
    normal: Vector3<f32>,
    supporting_views: Vec<usize>,
    confidence: f32,
}

/// 体素化点
#[derive(Debug, Clone)]
struct VoxelizedPoint {
    position: Point3,
    normal: Vector3<f32>,
    confidence: f32,
    supporting_views: Vec<usize>,
}

impl VoxelGrid {
    fn new(voxel_size: f32) -> Self {
        Self {
            voxel_size,
            voxels: HashMap::new(),
        }
    }

    fn get_voxel_key(&self, point: &Point3) -> (i32, i32, i32) {
        (
            (point.x / self.voxel_size as f64).floor() as i32,
            (point.y / self.voxel_size as f64).floor() as i32,
            (point.z / self.voxel_size as f64).floor() as i32,
        )
    }

    fn add_point(&mut self, voxel_key: (i32, i32, i32), point_id: usize) {
        self.voxels.entry(voxel_key).or_default().push(point_id);
    }
}