colmap 0.1.2

//! 3D 点和重建相关的数据结构
//!
//! 这个模块定义了 3D 点、轨迹和重建结果的数据结构。

use nalgebra::{Point3, Point2};
use serde::{Deserialize, Serialize};
use std::collections::{HashMap, HashSet};
use crate::core::image::ImageId;

/// 3D 点 ID 类型
pub type Point3dId = u64;

/// 轨迹 ID 类型
pub type TrackId = u64;

/// RGB 颜色
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub struct Color {
    pub r: u8,
    pub g: u8,
    pub b: u8,
}

impl Color {
    /// 创建新颜色
    pub fn new(r: u8, g: u8, b: u8) -> Self {
        Self { r, g, b }
    }
    
    /// 黑色
    pub fn black() -> Self {
        Self::new(0, 0, 0)
    }
    
    /// 白色
    pub fn white() -> Self {
        Self::new(255, 255, 255)
    }
    
    /// 红色
    pub fn red() -> Self {
        Self::new(255, 0, 0)
    }
    
    /// 绿色
    pub fn green() -> Self {
        Self::new(0, 255, 0)
    }
    
    /// 蓝色
    pub fn blue() -> Self {
        Self::new(0, 0, 255)
    }
    
    /// 转换为浮点数组
    pub fn to_f32_array(&self) -> [f32; 3] {
        [self.r as f32 / 255.0, self.g as f32 / 255.0, self.b as f32 / 255.0]
    }
    
    /// 从浮点数组创建
    pub fn from_f32_array(rgb: [f32; 3]) -> Self {
        Self::new(
            (rgb[0].clamp(0.0, 1.0) * 255.0) as u8,
            (rgb[1].clamp(0.0, 1.0) * 255.0) as u8,
            (rgb[2].clamp(0.0, 1.0) * 255.0) as u8,
        )
    }
}

/// 观测信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Observation {
    /// 图像 ID
    pub image_id: ImageId,
    /// 特征点索引
    pub feature_idx: usize,
    /// 2D 观测点
    pub point2d: Point2<f64>,
}

impl Observation {
    /// 创建新观测
    pub fn new(image_id: ImageId, feature_idx: usize, point2d: Point2<f64>) -> Self {
        Self {
            image_id,
            feature_idx,
            point2d,
        }
    }
}

/// 3D 点数据结构
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Point3d {
    /// 点 ID
    pub id: Point3dId,
    /// 3D 坐标
    pub position: Point3<f64>,
    /// 颜色
    pub color: Color,
    /// 重投影误差
    pub error: f64,
    /// 观测列表
    pub observations: Vec<Observation>,
    /// 轨迹 ID
    pub track_id: Option<TrackId>,
}

impl Point3d {
    /// 创建新的 3D 点
    pub fn new(id: Point3dId, position: Point3<f64>) -> Self {
        Self {
            id,
            position,
            color: Color::black(),
            error: 0.0,
            observations: Vec::new(),
            track_id: None,
        }
    }
    
    /// 添加观测
    pub fn add_observation(&mut self, observation: Observation) {
        self.observations.push(observation);
    }
    
    /// 获取观测数量
    pub fn num_observations(&self) -> usize {
        self.observations.len()
    }
    
    /// 检查是否有足够的观测
    pub fn has_sufficient_observations(&self, min_observations: usize) -> bool {
        self.observations.len() >= min_observations
    }
    
    /// 获取观测的图像 ID 列表
    pub fn image_ids(&self) -> Vec<ImageId> {
        self.observations.iter().map(|obs| obs.image_id).collect()
    }
    
    /// 检查是否被指定图像观测到
    pub fn is_observed_by(&self, image_id: ImageId) -> bool {
        self.observations.iter().any(|obs| obs.image_id == image_id)
    }
    
    /// 获取在指定图像中的观测
    pub fn get_observation_in_image(&self, image_id: ImageId) -> Option<&Observation> {
        self.observations.iter().find(|obs| obs.image_id == image_id)
    }
    
    /// 设置颜色
    pub fn set_color(&mut self, color: Color) {
        self.color = color;
    }
    
    /// 设置重投影误差
    pub fn set_error(&mut self, error: f64) {
        self.error = error;
    }
    
    /// 计算平均观测点
    pub fn mean_observation(&self) -> Option<Point2<f64>> {
        if self.observations.is_empty() {
            return None;
        }
        
        let sum = self.observations.iter()
            .fold(Point2::new(0.0, 0.0), |acc, obs| acc + obs.point2d.coords);
        
        Some(Point2::from(sum / self.observations.len() as f64))
    }
}

/// 特征轨迹
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Track {
    /// 轨迹 ID
    pub id: TrackId,
    /// 观测列表
    pub observations: Vec<Observation>,
    /// 对应的 3D 点 ID（如果已三角化）
    pub point3d_id: Option<Point3dId>,
    /// 轨迹长度（观测数量）
    pub length: usize,
}

impl Track {
    /// 创建新轨迹
    pub fn new(id: TrackId) -> Self {
        Self {
            id,
            observations: Vec::new(),
            point3d_id: None,
            length: 0,
        }
    }
    
    /// 添加观测
    pub fn add_observation(&mut self, observation: Observation) {
        self.observations.push(observation);
        self.length = self.observations.len();
    }
    
    /// 检查轨迹是否有效
    pub fn is_valid(&self, min_length: usize) -> bool {
        self.length >= min_length
    }
    
    /// 检查是否已三角化
    pub fn is_triangulated(&self) -> bool {
        self.point3d_id.is_some()
    }
    
    /// 设置对应的 3D 点
    pub fn set_point3d(&mut self, point3d_id: Point3dId) {
        self.point3d_id = Some(point3d_id);
    }
    
    /// 获取观测的图像 ID 集合
    pub fn image_ids(&self) -> HashSet<ImageId> {
        self.observations.iter().map(|obs| obs.image_id).collect()
    }
}

/// 重建结果
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct Reconstruction {
    /// 3D 点存储
    pub points3d: HashMap<Point3dId, Point3d>,
    /// 轨迹存储
    pub tracks: HashMap<TrackId, Track>,
    /// 下一个可用的点 ID
    next_point_id: Point3dId,
    /// 下一个可用的轨迹 ID
    next_track_id: TrackId,
}

impl Reconstruction {
    /// 创建新的重建结果
    pub fn new() -> Self {
        Self {
            points3d: HashMap::new(),
            tracks: HashMap::new(),
            next_point_id: 1,
            next_track_id: 1,
        }
    }
    
    /// 添加 3D 点
    pub fn add_point3d(&mut self, mut point: Point3d) -> Point3dId {
        if point.id == 0 {
            point.id = self.next_point_id;
            self.next_point_id += 1;
        } else {
            self.next_point_id = self.next_point_id.max(point.id + 1);
        }
        
        let id = point.id;
        self.points3d.insert(id, point);
        id
    }
    
    /// 添加轨迹
    pub fn add_track(&mut self, mut track: Track) -> TrackId {
        if track.id == 0 {
            track.id = self.next_track_id;
            self.next_track_id += 1;
        } else {
            self.next_track_id = self.next_track_id.max(track.id + 1);
        }
        
        let id = track.id;
        self.tracks.insert(id, track);
        id
    }
    
    /// 获取 3D 点
    pub fn get_point3d(&self, id: Point3dId) -> Option<&Point3d> {
        self.points3d.get(&id)
    }
    
    /// 获取可变 3D 点引用
    pub fn get_point3d_mut(&mut self, id: Point3dId) -> Option<&mut Point3d> {
        self.points3d.get_mut(&id)
    }
    
    /// 获取轨迹
    pub fn get_track(&self, id: TrackId) -> Option<&Track> {
        self.tracks.get(&id)
    }
    
    /// 获取可变轨迹引用
    pub fn get_track_mut(&mut self, id: TrackId) -> Option<&mut Track> {
        self.tracks.get_mut(&id)
    }
    
    /// 删除 3D 点
    pub fn remove_point3d(&mut self, id: Point3dId) -> Option<Point3d> {
        self.points3d.remove(&id)
    }
    
    /// 删除轨迹
    pub fn remove_track(&mut self, id: TrackId) -> Option<Track> {
        self.tracks.remove(&id)
    }
    
    /// 获取 3D 点数量
    pub fn num_points3d(&self) -> usize {
        self.points3d.len()
    }
    
    /// 获取轨迹数量
    pub fn num_tracks(&self) -> usize {
        self.tracks.len()
    }
    
    /// 获取已三角化的轨迹数量
    pub fn num_triangulated_tracks(&self) -> usize {
        self.tracks.values().filter(|track| track.is_triangulated()).count()
    }
    
    /// 计算重建的边界框
    pub fn bounding_box(&self) -> Option<(Point3<f64>, Point3<f64>)> {
        if self.points3d.is_empty() {
            return None;
        }
        
        let mut min_point = Point3::new(f64::INFINITY, f64::INFINITY, f64::INFINITY);
        let mut max_point = Point3::new(f64::NEG_INFINITY, f64::NEG_INFINITY, f64::NEG_INFINITY);
        
        for point in self.points3d.values() {
            let pos = &point.position;
            min_point.x = min_point.x.min(pos.x);
            min_point.y = min_point.y.min(pos.y);
            min_point.z = min_point.z.min(pos.z);
            max_point.x = max_point.x.max(pos.x);
            max_point.y = max_point.y.max(pos.y);
            max_point.z = max_point.z.max(pos.z);
        }
        
        Some((min_point, max_point))
    }
    
    /// 计算重建的中心点
    pub fn center(&self) -> Option<Point3<f64>> {
        if let Some((min_point, max_point)) = self.bounding_box() {
            Some(Point3::from((min_point.coords + max_point.coords) / 2.0))
        } else {
            None
        }
    }
    
    /// 清空重建
    pub fn clear(&mut self) {
        self.points3d.clear();
        self.tracks.clear();
        self.next_point_id = 1;
        self.next_track_id = 1;
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_color() {
        let color = Color::new(128, 64, 192);
        assert_eq!(color.r, 128);
        assert_eq!(color.g, 64);
        assert_eq!(color.b, 192);
        
        let float_array = color.to_f32_array();
        let back_color = Color::from_f32_array(float_array);
        assert_eq!(color, back_color);
    }

    #[test]
    fn test_point3d() {
        let mut point = Point3d::new(1, Point3::new(1.0, 2.0, 3.0));
        assert_eq!(point.num_observations(), 0);
        
        let obs = Observation::new(1, 0, Point2::new(100.0, 200.0));
        point.add_observation(obs);
        assert_eq!(point.num_observations(), 1);
        assert!(point.is_observed_by(1));
        assert!(!point.is_observed_by(2));
    }

    #[test]
    fn test_track() {
        let mut track = Track::new(1);
        assert!(!track.is_valid(2));
        assert!(!track.is_triangulated());
        
        track.add_observation(Observation::new(1, 0, Point2::new(100.0, 200.0)));
        track.add_observation(Observation::new(2, 1, Point2::new(150.0, 250.0)));
        
        assert!(track.is_valid(2));
        assert_eq!(track.length, 2);
        
        let image_ids = track.image_ids();
        assert!(image_ids.contains(&1));
        assert!(image_ids.contains(&2));
    }

    #[test]
    fn test_reconstruction() {
        let mut recon = Reconstruction::new();
        
        let point = Point3d::new(0, Point3::new(1.0, 2.0, 3.0));
        let point_id = recon.add_point3d(point);
        assert_eq!(point_id, 1);
        assert_eq!(recon.num_points3d(), 1);
        
        let track = Track::new(0);
        let track_id = recon.add_track(track);
        assert_eq!(track_id, 1);
        assert_eq!(recon.num_tracks(), 1);
    }
}