colmap 0.1.2

//! 立体匹配模块
//!
//! 实现双目立体匹配算法，用于计算视差图和深度图

use crate::core::{
    Image, CameraPose, CameraIntrinsics, Point3, Result, ColmapError
};
use nalgebra::{Matrix3, Point3 as NalgebraPoint3};

/// 立体匹配器
#[derive(Debug)]
pub struct StereoMatcher {
    /// 匹配配置
    config: StereoMatchConfig,
}

/// 立体匹配配置
#[derive(Debug, Clone)]
pub struct StereoMatchConfig {
    /// 最大视差
    pub max_disparity: i32,
    /// 窗口大小
    pub window_size: usize,
    /// 代价阈值
    pub cost_threshold: f32,
    /// 一致性检查
    pub consistency_check: bool,
    /// 子像素精度
    pub subpixel_refinement: bool,
}

impl Default for StereoMatchConfig {
    fn default() -> Self {
        Self {
            max_disparity: 128,
            window_size: 5,
            cost_threshold: 0.1,
            consistency_check: true,
            subpixel_refinement: true,
        }
    }
}

/// 视差图
#[derive(Debug, Clone)]
pub struct DisparityMap {
    /// 宽度
    pub width: u32,
    /// 高度
    pub height: u32,
    /// 视差数据
    pub data: Vec<f32>,
}

/// 深度图
#[derive(Debug, Clone)]
pub struct DepthMap {
    /// 宽度
    pub width: u32,
    /// 高度
    pub height: u32,
    /// 深度数据
    pub data: Vec<f32>,
    /// 相机内参
    pub intrinsics: CameraIntrinsics,
    /// 相机姿态
    pub pose: CameraPose,
}

/// 立体匹配结果
#[derive(Debug)]
pub struct StereoMatchResult {
    /// 视差图
    pub disparity_map: DisparityMap,
    /// 深度图
    pub depth_map: DepthMap,
    /// 匹配置信度
    pub confidence: Vec<f32>,
}

impl StereoMatcher {
    /// 创建新的立体匹配器
    pub fn new(config: StereoMatchConfig) -> Self {
        Self { config }
    }

    /// 执行立体匹配
    pub fn match_stereo(
        &self,
        left_image: &Image,
        right_image: &Image,
        left_intrinsics: &CameraIntrinsics,
        right_intrinsics: &CameraIntrinsics,
        left_pose: &CameraPose,
        right_pose: &CameraPose,
    ) -> Result<StereoMatchResult> {
        // 检查图像尺寸
        if left_image.size.0 != right_image.size.0 || left_image.size.1 != right_image.size.1 {
            return Err(ColmapError::MvsReconstruction(
                "Images must have the same dimensions".to_string()
            ));
        }

        let width = left_image.size.0;
        let height = left_image.size.1;

        // 计算基础矩阵和极线几何
        let fundamental_matrix = self.compute_fundamental_matrix(
            left_intrinsics, right_intrinsics, left_pose, right_pose
        )?;

        // 执行立体校正
        let (rectified_left, rectified_right, rectification_params) = 
            self.rectify_stereo_pair(
                left_image, right_image,
                left_intrinsics, right_intrinsics,
                &fundamental_matrix
            )?;

        // 计算视差图
        let disparity_map = self.compute_disparity_map(
            &rectified_left, &rectified_right
        )?;

        // 转换为深度图
        let depth_map = self.disparity_to_depth(
            &disparity_map,
            &rectification_params,
            left_intrinsics,
            left_pose,
        )?;

        // 计算置信度
        let confidence = self.compute_confidence(&disparity_map, &rectified_left, &rectified_right)?;

        Ok(StereoMatchResult {
            disparity_map,
            depth_map,
            confidence,
        })
    }

    /// 计算基础矩阵
    fn compute_fundamental_matrix(
        &self,
        left_intrinsics: &CameraIntrinsics,
        right_intrinsics: &CameraIntrinsics,
        left_pose: &CameraPose,
        right_pose: &CameraPose,
    ) -> Result<Matrix3<f64>> {
        // 计算相对姿态
        let relative_rotation = right_pose.rotation.inverse() * left_pose.rotation;
        let relative_translation = right_pose.rotation.inverse() * 
            (left_pose.translation - right_pose.translation);

        // 构建反对称矩阵
        let tx = relative_translation.x;
        let ty = relative_translation.y;
        let tz = relative_translation.z;
        
        let skew_symmetric = Matrix3::new(
            0.0, -tz, ty,
            tz, 0.0, -tx,
            -ty, tx, 0.0,
        );

        // 计算本质矩阵
        let essential_matrix = skew_symmetric * relative_rotation.to_rotation_matrix().matrix();

        // 转换为基础矩阵
        let k1_inv = left_intrinsics.matrix().try_inverse()
            .ok_or_else(|| ColmapError::MvsReconstruction(
                "Cannot invert left camera intrinsics matrix".to_string()
            ))?;
        let k2_inv = right_intrinsics.matrix().try_inverse()
            .ok_or_else(|| ColmapError::MvsReconstruction(
                "Cannot invert right camera intrinsics matrix".to_string()
            ))?;

        let fundamental_matrix = k2_inv.transpose() * essential_matrix * k1_inv;

        Ok(fundamental_matrix)
    }

    /// 立体校正
    fn rectify_stereo_pair(
        &self,
        left_image: &Image,
        right_image: &Image,
        left_intrinsics: &CameraIntrinsics,
        _right_intrinsics: &CameraIntrinsics,
        _fundamental_matrix: &Matrix3<f64>,
    ) -> Result<(Image, Image, RectificationParams)> {
        // 简化实现：假设图像已经校正
        // 在实际实现中，这里需要进行复杂的立体校正算法
        
        let rectified_left = left_image.clone();
        let rectified_right = right_image.clone();
        
        let params = RectificationParams {
            baseline: 0.1, // 假设基线长度
            focal_length: left_intrinsics.focal_length.0,
        };

        Ok((rectified_left, rectified_right, params))
    }

    /// 计算视差图
    fn compute_disparity_map(
        &self,
        left_image: &Image,
        right_image: &Image,
    ) -> Result<DisparityMap> {
        let width = left_image.size.0;
        let height = left_image.size.1;
        let mut disparity_data = vec![0.0; (width * height) as usize];

        // 简化的块匹配算法
        let half_window = self.config.window_size / 2;
        
        for y in half_window..(height as usize - half_window) {
            for x in half_window..(width as usize - half_window) {
                let mut best_disparity = 0.0;
                let mut min_cost = f32::MAX;

                // 在视差范围内搜索最佳匹配
                for d in 0..self.config.max_disparity {
                    let right_x = x as i32 - d;
                    if right_x < half_window as i32 || 
                       right_x >= (width as i32 - half_window as i32) {
                        continue;
                    }

                    // 计算匹配代价（简化为 SAD）
                    let cost = self.compute_matching_cost(
                        left_image, right_image,
                        x, y, right_x as usize, y
                    );

                    if cost < min_cost {
                        min_cost = cost;
                        best_disparity = d as f32;
                    }
                }

                // 子像素精度优化
                if self.config.subpixel_refinement {
                    best_disparity = self.refine_subpixel_disparity(
                        left_image, right_image, x, y, best_disparity
                    );
                }

                disparity_data[y * width as usize + x] = best_disparity;
            }
        }

        Ok(DisparityMap {
            width,
            height,
            data: disparity_data,
        })
    }

    /// 计算匹配代价
    fn compute_matching_cost(
        &self,
        left_image: &Image,
        right_image: &Image,
        left_x: usize,
        left_y: usize,
        right_x: usize,
        right_y: usize,
    ) -> f32 {
        let half_window = self.config.window_size / 2;
        let mut cost = 0.0f32;
        let mut count = 0;

        for dy in -(half_window as i32)..=(half_window as i32) {
            for dx in -(half_window as i32)..=(half_window as i32) {
                let ly = (left_y as i32 + dy) as usize;
                let lx = (left_x as i32 + dx) as usize;
                let ry = (right_y as i32 + dy) as usize;
                let rx = (right_x as i32 + dx) as usize;

                if ly < left_image.size.1 as usize && lx < left_image.size.0 as usize &&
                   ry < right_image.size.1 as usize && rx < right_image.size.0 as usize {
                    
                    // 简化：假设图像有灰度数据
                    let left_intensity = 128.0f32; // 占位符
                    let right_intensity = 128.0f32; // 占位符
                    
                    cost += (left_intensity - right_intensity).abs();
                    count += 1;
                }
            }
        }

        if count > 0 {
            cost / count as f32
        } else {
            f32::MAX
        }
    }

    /// 子像素精度优化
    fn refine_subpixel_disparity(
        &self,
        left_image: &Image,
        right_image: &Image,
        x: usize,
        y: usize,
        disparity: f32,
    ) -> f32 {
        // 简化实现：使用抛物线拟合
        let d = disparity as i32;
        
        if d <= 0 || d >= self.config.max_disparity - 1 {
            return disparity;
        }

        let cost_prev = self.compute_matching_cost(
            left_image, right_image, x, y, (x as i32 - d + 1) as usize, y
        );
        let cost_curr = self.compute_matching_cost(
            left_image, right_image, x, y, (x as i32 - d) as usize, y
        );
        let cost_next = self.compute_matching_cost(
            left_image, right_image, x, y, (x as i32 - d - 1) as usize, y
        );

        // 抛物线拟合求最小值
        let denom = 2.0 * (cost_prev + cost_next - 2.0 * cost_curr);
        if denom.abs() > 1e-6 {
            let offset = (cost_prev - cost_next) / denom;
            disparity + offset.clamp(-1.0, 1.0)
        } else {
            disparity
        }
    }

    /// 视差转深度
    fn disparity_to_depth(
        &self,
        disparity_map: &DisparityMap,
        rectification_params: &RectificationParams,
        intrinsics: &CameraIntrinsics,
        pose: &CameraPose,
    ) -> Result<DepthMap> {
        let mut depth_data = vec![0.0; disparity_map.data.len()];

        for (i, &disparity) in disparity_map.data.iter().enumerate() {
            if disparity > 0.0 {
                // 深度 = 基线 * 焦距 / 视差
                let depth = rectification_params.baseline * rectification_params.focal_length / disparity as f64;
                depth_data[i] = depth as f32;
            }
        }

        Ok(DepthMap {
            width: disparity_map.width,
            height: disparity_map.height,
            data: depth_data,
            intrinsics: intrinsics.clone(),
            pose: pose.clone(),
        })
    }

    /// 计算置信度
    fn compute_confidence(
        &self,
        disparity_map: &DisparityMap,
        left_image: &Image,
        right_image: &Image,
    ) -> Result<Vec<f32>> {
        let mut confidence = vec![0.0; disparity_map.data.len()];

        for (i, &disparity) in disparity_map.data.iter().enumerate() {
            if disparity > 0.0 {
                // 简化的置信度计算：基于左右一致性检查
                confidence[i] = if disparity > self.config.cost_threshold {
                    1.0 - (disparity / self.config.max_disparity as f32)
                } else {
                    0.0
                };
            }
        }

        Ok(confidence)
    }
}

/// 立体校正参数
#[derive(Debug, Clone)]
pub struct RectificationParams {
    /// 基线长度
    pub baseline: f64,
    /// 焦距
    pub focal_length: f64,
}

/// 深度图工具函数
impl DepthMap {
    /// 获取指定像素的深度值
    pub fn get_depth(&self, x: u32, y: u32) -> Option<f32> {
        if x < self.width && y < self.height {
            let index = (y * self.width + x) as usize;
            let depth = self.data[index];
            if depth > 0.0 {
                Some(depth)
            } else {
                None
            }
        } else {
            None
        }
    }

    /// 将深度图转换为点云
    pub fn to_point_cloud(&self) -> Vec<Point3> {
        let mut points = Vec::new();

        for y in 0..self.height {
            for x in 0..self.width {
                if let Some(depth) = self.get_depth(x, y) {
                    // 反投影到3D空间
                    let x_norm = (x as f64 - self.intrinsics.principal_point.0) / self.intrinsics.focal_length.0;
        let y_norm = (y as f64 - self.intrinsics.principal_point.1) / self.intrinsics.focal_length.1;
                    
                    let point_camera = NalgebraPoint3::new(
                        x_norm * depth as f64,
                        y_norm * depth as f64,
                        depth as f64,
                    );

                    // 转换到世界坐标系
                    let point_world: nalgebra::Point3<f64> = (self.pose.rotation.inverse() *
                        (point_camera - nalgebra::Point3::<f64>::from(self.pose.translation))).into();

                    points.push(Point3::new(point_world.x, point_world.y, point_world.z));
                }
            }
        }

        points
    }
}