colmap 0.1.2

//! 图像相关的数据结构
//!
//! 这个模块定义了图像、关键点和描述符的数据结构。

use nalgebra::Point2;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// 图像 ID 类型
pub type ImageId = u32;

/// 图像数据结构
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Image {
    /// 图像 ID
    pub id: ImageId,
    /// 图像文件名
    pub name: String,
    /// 相机 ID
    pub camera_id: u32,
    /// 图像尺寸 (width, height)
    pub size: (u32, u32),
    /// 相机姿态（如果已注册）
    pub pose: Option<crate::core::CameraPose>,
    /// 特征点列表
    pub features: Vec<crate::core::Feature>,
    /// 是否已注册到重建中
    pub registered: bool,
    /// 图像质量分数
    pub quality_score: f64,
}

impl Image {
    /// 创建新图像
    pub fn new(
        id: ImageId,
        name: String,
        camera_id: u32,
        size: (u32, u32),
    ) -> Self {
        Self {
            id,
            name,
            camera_id,
            size,
            pose: None,
            features: Vec::new(),
            registered: false,
            quality_score: 0.0,
        }
    }
    
    /// 设置相机姿态
    pub fn set_pose(&mut self, pose: crate::core::CameraPose) {
        self.pose = Some(pose);
        self.registered = true;
    }
    
    /// 获取相机姿态
    pub fn pose(&self) -> Option<&crate::core::CameraPose> {
        self.pose.as_ref()
    }
    
    /// 检查图像是否已注册
    pub fn is_registered(&self) -> bool {
        self.registered && self.pose.is_some()
    }
    
    /// 添加特征点
    pub fn add_feature(&mut self, feature: crate::core::Feature) {
        self.features.push(feature);
    }
    
    /// 获取特征点数量
    pub fn num_features(&self) -> usize {
        self.features.len()
    }
    
    /// 获取已三角化的特征点数量
    pub fn num_triangulated_features(&self) -> usize {
        self.features.iter().filter(|f| f.is_triangulated()).count()
    }
    
    /// 设置质量分数
    pub fn set_quality_score(&mut self, score: f64) {
        self.quality_score = score;
    }
    
    /// 获取图像宽高比
    pub fn aspect_ratio(&self) -> f64 {
        self.size.0 as f64 / self.size.1 as f64
    }
    
    /// 检查点是否在图像边界内
    pub fn contains_point(&self, point: &Point2<f64>) -> bool {
        point.x >= 0.0 && point.x < self.size.0 as f64 &&
        point.y >= 0.0 && point.y < self.size.1 as f64
    }
    
    /// 获取图像中心点
    pub fn center(&self) -> Point2<f64> {
        Point2::new(self.size.0 as f64 / 2.0, self.size.1 as f64 / 2.0)
    }
}

/// 图像对
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct ImagePair {
    /// 第一个图像 ID
    pub image1_id: ImageId,
    /// 第二个图像 ID
    pub image2_id: ImageId,
}

impl ImagePair {
    /// 创建新的图像对
    pub fn new(image1_id: ImageId, image2_id: ImageId) -> Self {
        // 确保 ID 顺序一致
        if image1_id <= image2_id {
            Self { image1_id, image2_id }
        } else {
            Self { image1_id: image2_id, image2_id: image1_id }
        }
    }
    
    /// 检查是否包含指定图像
    pub fn contains(&self, image_id: ImageId) -> bool {
        self.image1_id == image_id || self.image2_id == image_id
    }
    
    /// 获取另一个图像 ID
    pub fn other(&self, image_id: ImageId) -> Option<ImageId> {
        if self.image1_id == image_id {
            Some(self.image2_id)
        } else if self.image2_id == image_id {
            Some(self.image1_id)
        } else {
            None
        }
    }
}

/// 图像数据库
#[derive(Debug, Default, Serialize, Deserialize)]
pub struct ImageDatabase {
    /// 图像存储
    images: HashMap<ImageId, Image>,
    /// 下一个可用的图像 ID
    next_id: ImageId,
}

impl ImageDatabase {
    /// 创建新的图像数据库
    pub fn new() -> Self {
        Self {
            images: HashMap::new(),
            next_id: 1,
        }
    }
    
    /// 添加图像
    pub fn add_image(&mut self, mut image: Image) -> ImageId {
        if image.id == 0 {
            image.id = self.next_id;
            self.next_id += 1;
        } else {
            self.next_id = self.next_id.max(image.id + 1);
        }
        
        let id = image.id;
        self.images.insert(id, image);
        id
    }
    
    /// 获取图像
    pub fn get_image(&self, id: ImageId) -> Option<&Image> {
        self.images.get(&id)
    }
    
    /// 获取可变图像引用
    pub fn get_image_mut(&mut self, id: ImageId) -> Option<&mut Image> {
        self.images.get_mut(&id)
    }
    
    /// 删除图像
    pub fn remove_image(&mut self, id: ImageId) -> Option<Image> {
        self.images.remove(&id)
    }
    
    /// 获取所有图像 ID
    pub fn image_ids(&self) -> Vec<ImageId> {
        self.images.keys().copied().collect()
    }
    
    /// 获取所有图像
    pub fn images(&self) -> &HashMap<ImageId, Image> {
        &self.images
    }
    
    /// 获取图像数量
    pub fn num_images(&self) -> usize {
        self.images.len()
    }
    
    /// 获取已注册图像数量
    pub fn num_registered_images(&self) -> usize {
        self.images.values().filter(|img| img.is_registered()).count()
    }
    
    /// 检查图像是否存在
    pub fn exists(&self, id: ImageId) -> bool {
        self.images.contains_key(&id)
    }
    
    /// 根据文件名查找图像
    pub fn find_by_name(&self, name: &str) -> Option<&Image> {
        self.images.values().find(|img| img.name == name)
    }
    
    /// 清空数据库
    pub fn clear(&mut self) {
        self.images.clear();
        self.next_id = 1;
    }
}

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

    #[test]
    fn test_image_creation() {
        let image = Image::new(1, "test.jpg".to_string(), 1, (640, 480));
        assert_eq!(image.id, 1);
        assert_eq!(image.name, "test.jpg");
        assert_eq!(image.camera_id, 1);
        assert_eq!(image.size, (640, 480));
        assert!(!image.is_registered());
        assert_eq!(image.num_features(), 0);
    }

    #[test]
    fn test_image_contains_point() {
        let image = Image::new(1, "test.jpg".to_string(), 1, (640, 480));
        
        assert!(image.contains_point(&Point2::new(100.0, 100.0)));
        assert!(image.contains_point(&Point2::new(0.0, 0.0)));
        assert!(image.contains_point(&Point2::new(639.0, 479.0)));
        
        assert!(!image.contains_point(&Point2::new(-1.0, 100.0)));
        assert!(!image.contains_point(&Point2::new(640.0, 100.0)));
        assert!(!image.contains_point(&Point2::new(100.0, 480.0)));
    }

    #[test]
    fn test_image_pair() {
        let pair1 = ImagePair::new(1, 2);
        let pair2 = ImagePair::new(2, 1);
        assert_eq!(pair1, pair2); // 应该标准化顺序
        
        assert!(pair1.contains(1));
        assert!(pair1.contains(2));
        assert!(!pair1.contains(3));
        
        assert_eq!(pair1.other(1), Some(2));
        assert_eq!(pair1.other(2), Some(1));
        assert_eq!(pair1.other(3), None);
    }

    #[test]
    fn test_image_database() {
        let mut db = ImageDatabase::new();
        
        let image1 = Image::new(0, "test1.jpg".to_string(), 1, (640, 480));
        let image2 = Image::new(0, "test2.jpg".to_string(), 1, (640, 480));
        
        let id1 = db.add_image(image1);
        let id2 = db.add_image(image2);
        
        assert_eq!(id1, 1);
        assert_eq!(id2, 2);
        assert_eq!(db.num_images(), 2);
        
        assert!(db.exists(id1));
        assert!(db.exists(id2));
        assert!(!db.exists(3));
        
        let retrieved = db.get_image(id1).unwrap();
        assert_eq!(retrieved.name, "test1.jpg");
        
        let found = db.find_by_name("test2.jpg").unwrap();
        assert_eq!(found.id, id2);
    }
}