colmap 0.1.2

//! 特征点和描述符定义
//!
//! 这个模块定义了特征点、描述符和特征匹配相关的数据结构。

use crate::core::{Point2, Result, ColmapError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;

/// 特征点 ID 类型
pub type FeatureId = u32;

/// 图像 ID 类型（重新导出 image 模块中的定义）
pub use crate::core::image::ImageId;

/// 特征点
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct Feature {
    /// 图像坐标 (x, y)
    pub point: Point2,
    /// 特征描述符
    pub descriptor: Vec<u8>,
    /// 特征点的尺度（用于 SIFT 等）
    pub scale: f32,
    /// 特征点的方向角度（弧度）
    pub angle: f32,
    /// 特征点的响应强度
    pub response: f32,
    /// 金字塔层级（用于 SIFT 等）
    pub octave: i32,
    /// 对应的 3D 点 ID（如果已三角化）
    pub point3d_id: Option<u32>,
}

impl Feature {
    /// 创建新的特征点
    pub fn new(
        point: Point2,
        descriptor: Vec<u8>,
        scale: f32,
        angle: f32,
        response: f32,
        octave: i32,
    ) -> Self {
        Self {
            point,
            descriptor,
            scale,
            angle,
            response,
            octave,
            point3d_id: None,
        }
    }

    /// 创建简单的特征点（默认参数）
    pub fn simple(point: Point2, descriptor: Vec<u8>) -> Self {
        Self::new(point, descriptor, 1.0, 0.0, 0.0, 0)
    }

    /// 检查特征点是否已三角化
    pub fn is_triangulated(&self) -> bool {
        self.point3d_id.is_some()
    }

    /// 设置对应的 3D 点 ID
    pub fn set_point3d_id(&mut self, point3d_id: u32) {
        self.point3d_id = Some(point3d_id);
    }

    /// 清除 3D 点关联
    pub fn clear_point3d_id(&mut self) {
        self.point3d_id = None;
    }

    /// 计算与另一个特征点的描述符距离
    pub fn descriptor_distance(&self, other: &Feature) -> Result<f32> {
        if self.descriptor.len() != other.descriptor.len() {
            return Err(ColmapError::FeatureExtraction(
                "Descriptor lengths do not match".to_string(),
            ));
        }

        // 计算汉明距离（对于二进制描述符）或欧几里得距离
        if self.is_binary_descriptor() {
            Ok(self.hamming_distance(&other.descriptor) as f32)
        } else {
            Ok(self.euclidean_distance(&other.descriptor))
        }
    }

    /// 检查是否为二进制描述符
    fn is_binary_descriptor(&self) -> bool {
        // 简单启发式：如果描述符长度是 8 的倍数且值只有 0 或 1，认为是二进制
        self.descriptor.len() % 8 == 0 && 
        self.descriptor.iter().all(|&x| x == 0 || x == 1)
    }

    /// 计算汉明距离
    fn hamming_distance(&self, other: &[u8]) -> u32 {
        self.descriptor
            .iter()
            .zip(other.iter())
            .map(|(&a, &b)| (a ^ b).count_ones())
            .sum()
    }

    /// 计算欧几里得距离
    fn euclidean_distance(&self, other: &[u8]) -> f32 {
        self.descriptor
            .iter()
            .zip(other.iter())
            .map(|(&a, &b)| {
                let diff = a as f32 - b as f32;
                diff * diff
            })
            .sum::<f32>()
            .sqrt()
    }
}

/// 特征匹配
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct FeatureMatch {
    /// 第一个图像 ID
    pub image1_id: ImageId,
    /// 第二个图像 ID
    pub image2_id: ImageId,
    /// 第一个图像中的特征点索引
    pub feature1_idx: usize,
    /// 第二个图像中的特征点索引
    pub feature2_idx: usize,
    /// 匹配距离
    pub distance: f32,
    /// 匹配置信度
    pub confidence: f32,
}

impl FeatureMatch {
    /// 创建新的特征匹配
    pub fn new(
        image1_id: ImageId,
        image2_id: ImageId,
        feature1_idx: usize,
        feature2_idx: usize,
        distance: f32,
    ) -> Self {
        Self {
            image1_id,
            image2_id,
            feature1_idx,
            feature2_idx,
            distance,
            confidence: 1.0 / (1.0 + distance), // 简单的置信度计算
        }
    }

    /// 设置匹配置信度
    pub fn set_confidence(&mut self, confidence: f32) {
        self.confidence = confidence.clamp(0.0, 1.0);
    }

    /// 获取图像对 ID
    pub fn image_pair(&self) -> (ImageId, ImageId) {
        (self.image1_id, self.image2_id)
    }

    /// 获取特征点索引对
    pub fn feature_pair(&self) -> (usize, usize) {
        (self.feature1_idx, self.feature2_idx)
    }

    /// 检查匹配是否有效（距离和置信度在合理范围内）
    pub fn is_valid(&self, max_distance: f32, min_confidence: f32) -> bool {
        self.distance <= max_distance && self.confidence >= min_confidence
    }
}

/// 特征匹配对（两个图像之间的所有匹配）
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureMatches {
    /// 图像对 ID
    pub image1_id: ImageId,
    pub image2_id: ImageId,
    /// 匹配列表
    pub matches: Vec<FeatureMatch>,
    /// 几何验证状态
    pub verified: bool,
}

impl FeatureMatches {
    /// 创建新的特征匹配对
    pub fn new(image1_id: ImageId, image2_id: ImageId) -> Self {
        Self {
            image1_id,
            image2_id,
            matches: Vec::new(),
            verified: false,
        }
    }

    /// 添加匹配
    pub fn add_match(&mut self, feature_match: FeatureMatch) {
        self.matches.push(feature_match);
    }

    /// 获取匹配数量
    pub fn count(&self) -> usize {
        self.matches.len()
    }

    /// 过滤匹配（根据距离和置信度）
    pub fn filter_matches(&mut self, max_distance: f32, min_confidence: f32) {
        self.matches.retain(|m| m.is_valid(max_distance, min_confidence));
    }

    /// 按距离排序匹配
    pub fn sort_by_distance(&mut self) {
        self.matches.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
    }

    /// 按置信度排序匹配
    pub fn sort_by_confidence(&mut self) {
        self.matches.sort_by(|a, b| b.confidence.partial_cmp(&a.confidence).unwrap());
    }

    /// 获取平均匹配距离
    pub fn average_distance(&self) -> f32 {
        if self.matches.is_empty() {
            return 0.0;
        }
        self.matches.iter().map(|m| m.distance).sum::<f32>() / self.matches.len() as f32
    }

    /// 获取平均匹配置信度
    pub fn average_confidence(&self) -> f32 {
        if self.matches.is_empty() {
            return 0.0;
        }
        self.matches.iter().map(|m| m.confidence).sum::<f32>() / self.matches.len() as f32
    }

    /// 标记为已验证
    pub fn mark_verified(&mut self) {
        self.verified = true;
    }
}

/// 特征数据库，管理图像的特征点
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct FeatureDatabase {
    /// 图像特征映射
    image_features: HashMap<ImageId, Vec<Feature>>,
    /// 特征匹配映射
    feature_matches: HashMap<(ImageId, ImageId), FeatureMatches>,
}

impl FeatureDatabase {
    /// 创建新的特征数据库
    pub fn new() -> Self {
        Self {
            image_features: HashMap::new(),
            feature_matches: HashMap::new(),
        }
    }

    /// 添加图像特征
    pub fn add_image_features(&mut self, image_id: ImageId, features: Vec<Feature>) {
        self.image_features.insert(image_id, features);
    }

    /// 获取图像特征
    pub fn get_image_features(&self, image_id: ImageId) -> Option<&Vec<Feature>> {
        self.image_features.get(&image_id)
    }

    /// 获取可变图像特征引用
    pub fn get_image_features_mut(&mut self, image_id: ImageId) -> Option<&mut Vec<Feature>> {
        self.image_features.get_mut(&image_id)
    }

    /// 添加特征匹配
    pub fn add_feature_matches(&mut self, matches: FeatureMatches) {
        let key = (matches.image1_id, matches.image2_id);
        self.feature_matches.insert(key, matches);
    }

    /// 获取特征匹配
    pub fn get_feature_matches(&self, image1_id: ImageId, image2_id: ImageId) -> Option<&FeatureMatches> {
        // 尝试两种顺序
        self.feature_matches.get(&(image1_id, image2_id))
            .or_else(|| self.feature_matches.get(&(image2_id, image1_id)))
    }

    /// 获取可变特征匹配引用
    pub fn get_feature_matches_mut(&mut self, image1_id: ImageId, image2_id: ImageId) -> Option<&mut FeatureMatches> {
        // 尝试两种顺序
        if self.feature_matches.contains_key(&(image1_id, image2_id)) {
            self.feature_matches.get_mut(&(image1_id, image2_id))
        } else {
            self.feature_matches.get_mut(&(image2_id, image1_id))
        }
    }

    /// 获取图像的特征数量
    pub fn feature_count(&self, image_id: ImageId) -> usize {
        self.image_features.get(&image_id).map_or(0, |f| f.len())
    }

    /// 获取已三角化的特征数量
    pub fn triangulated_count(&self, image_id: ImageId) -> usize {
        self.image_features
            .get(&image_id)
            .map_or(0, |features| features.iter().filter(|f| f.is_triangulated()).count())
    }

    /// 获取所有图像 ID
    pub fn image_ids(&self) -> Vec<ImageId> {
        self.image_features.keys().copied().collect()
    }

    /// 获取所有匹配对
    pub fn match_pairs(&self) -> Vec<(ImageId, ImageId)> {
        self.feature_matches.keys().copied().collect()
    }

    /// 移除图像特征
    pub fn remove_image_features(&mut self, image_id: ImageId) -> Option<Vec<Feature>> {
        // 同时移除相关的匹配
        self.feature_matches.retain(|(id1, id2), _| *id1 != image_id && *id2 != image_id);
        self.image_features.remove(&image_id)
    }

    /// 清空所有数据
    pub fn clear(&mut self) {
        self.image_features.clear();
        self.feature_matches.clear();
    }

    /// 获取数据库统计信息
    pub fn stats(&self) -> FeatureDatabaseStats {
        let total_features: usize = self.image_features.values().map(|f| f.len()).sum();
        let total_matches: usize = self.feature_matches.values().map(|m| m.count()).sum();
        let triangulated_features: usize = self.image_features
            .values()
            .flat_map(|features| features.iter())
            .filter(|f| f.is_triangulated())
            .count();

        FeatureDatabaseStats {
            image_count: self.image_features.len(),
            total_features,
            triangulated_features,
            match_pairs: self.feature_matches.len(),
            total_matches,
        }
    }
}

/// 特征数据库统计信息
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FeatureDatabaseStats {
    /// 图像数量
    pub image_count: usize,
    /// 总特征数量
    pub total_features: usize,
    /// 已三角化特征数量
    pub triangulated_features: usize,
    /// 匹配对数量
    pub match_pairs: usize,
    /// 总匹配数量
    pub total_matches: usize,
}

impl FeatureDatabaseStats {
    /// 获取三角化率
    pub fn triangulation_rate(&self) -> f32 {
        if self.total_features == 0 {
            0.0
        } else {
            self.triangulated_features as f32 / self.total_features as f32
        }
    }

    /// 获取平均每图像特征数
    pub fn average_features_per_image(&self) -> f32 {
        if self.image_count == 0 {
            0.0
        } else {
            self.total_features as f32 / self.image_count as f32
        }
    }

    /// 获取平均每对匹配数
    pub fn average_matches_per_pair(&self) -> f32 {
        if self.match_pairs == 0 {
            0.0
        } else {
            self.total_matches as f32 / self.match_pairs as f32
        }
    }
}

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

    #[test]
    fn test_feature_creation() {
        let point = Point2::new(100.0, 200.0);
        let descriptor = vec![1, 2, 3, 4];
        let feature = Feature::simple(point, descriptor.clone());
        
        assert_eq!(feature.point, point);
        assert_eq!(feature.descriptor, descriptor);
        assert!(!feature.is_triangulated());
    }

    #[test]
    fn test_feature_match() {
        let match1 = FeatureMatch::new(1, 2, 0, 1, 0.5);
        
        assert_eq!(match1.image_pair(), (1, 2));
        assert_eq!(match1.feature_pair(), (0, 1));
        assert!(match1.is_valid(1.0, 0.1));
        assert!(!match1.is_valid(0.1, 0.1));
    }

    #[test]
    fn test_feature_database() {
        let mut db = FeatureDatabase::new();
        
        let features = vec![
            Feature::simple(Point2::new(10.0, 20.0), vec![1, 2, 3]),
            Feature::simple(Point2::new(30.0, 40.0), vec![4, 5, 6]),
        ];
        
        db.add_image_features(1, features);
        
        assert_eq!(db.feature_count(1), 2);
        assert_eq!(db.triangulated_count(1), 0);
        
        let stats = db.stats();
        assert_eq!(stats.image_count, 1);
        assert_eq!(stats.total_features, 2);
    }

    #[test]
    fn test_descriptor_distance() {
        let feature1 = Feature::simple(Point2::new(0.0, 0.0), vec![1, 2, 3, 4]);
        let feature2 = Feature::simple(Point2::new(0.0, 0.0), vec![1, 2, 3, 5]);
        
        let distance = feature1.descriptor_distance(&feature2).unwrap();
        assert!(distance > 0.0);
    }
}