colmap 0.1.2

//! 特征匹配器模块
//! 
//! 提供统一的特征匹配接口，支持多种匹配算法

use crate::core::{Feature, ColmapError};
use std::collections::HashMap;

/// 特征匹配结果
#[derive(Debug, Clone)]
pub struct FeatureMatch {
    /// 第一个图像中的特征点索引
    pub query_idx: usize,
    /// 第二个图像中的特征点索引
    pub train_idx: usize,
    /// 匹配距离
    pub distance: f64,
    /// 匹配置信度
    pub confidence: f64,
}

impl FeatureMatch {
    /// 创建新的特征匹配
    pub fn new(query_idx: usize, train_idx: usize, distance: f64) -> Self {
        Self {
            query_idx,
            train_idx,
            distance,
            confidence: 1.0 / (1.0 + distance),
        }
    }
    
    /// 检查匹配是否有效
    pub fn is_valid(&self, max_distance: f64) -> bool {
        self.distance <= max_distance
    }
}

/// 特征匹配器的通用接口
pub trait FeatureMatcher: Send + Sync {
    /// 匹配两组特征点
    fn match_features(
        &self,
        features1: &[Feature],
        features2: &[Feature],
    ) -> Result<Vec<FeatureMatch>, ColmapError>;
    
    /// 获取匹配器名称
    fn name(&self) -> &str;
    
    /// 获取匹配器参数
    fn params(&self) -> HashMap<String, f64>;
    
    /// 设置匹配器参数
    fn set_params(&mut self, params: HashMap<String, f64>) -> Result<(), ColmapError>;
}

/// 匹配器类型
#[derive(Debug, Clone, PartialEq)]
pub enum MatcherType {
    /// 暴力匹配器
    BruteForce,
    /// FLANN 匹配器
    Flann,
    /// 基于比值的匹配器
    RatioTest,
    /// 交叉检查匹配器
    CrossCheck,
}

/// 距离度量类型
#[derive(Debug, Clone, PartialEq)]
pub enum DistanceType {
    /// 欧几里得距离（L2）
    L2,
    /// 曼哈顿距离（L1）
    L1,
    /// 汉明距离（用于二进制描述符）
    Hamming,
    /// 余弦距离
    Cosine,
}

/// 匹配器配置
#[derive(Debug, Clone)]
pub struct MatcherConfig {
    /// 匹配器类型
    pub matcher_type: MatcherType,
    /// 距离度量类型
    pub distance_type: DistanceType,
    /// 最大匹配距离
    pub max_distance: f64,
    /// 比值测试阈值
    pub ratio_threshold: f64,
    /// 是否启用交叉检查
    pub cross_check: bool,
    /// 最大匹配数量
    pub max_matches: usize,
    /// 最小匹配置信度
    pub min_confidence: f64,
}

impl Default for MatcherConfig {
    fn default() -> Self {
        Self {
            matcher_type: MatcherType::BruteForce,
            distance_type: DistanceType::L2,
            max_distance: 0.7,
            ratio_threshold: 0.8,
            cross_check: true,
            max_matches: 1000,
            min_confidence: 0.1,
        }
    }
}

/// 匹配器工厂
pub struct MatcherFactory;

impl MatcherFactory {
    /// 创建特征匹配器
    pub fn create(config: &MatcherConfig) -> Result<Box<dyn FeatureMatcher>, ColmapError> {
        match config.matcher_type {
            MatcherType::BruteForce => {
                Ok(Box::new(BruteForceMatcher::new(config)?))
            },
            MatcherType::Flann => {
                Ok(Box::new(FlannMatcher::new(config)?))
            },
            MatcherType::RatioTest => {
                Ok(Box::new(RatioTestMatcher::new(config)?))
            },
            MatcherType::CrossCheck => {
                Ok(Box::new(CrossCheckMatcher::new(config)?))
            },
        }
    }
}

/// 暴力匹配器
pub struct BruteForceMatcher {
    config: MatcherConfig,
}

impl BruteForceMatcher {
    pub fn new(config: &MatcherConfig) -> Result<Self, ColmapError> {
        Ok(Self {
            config: config.clone(),
        })
    }
    
    /// 计算描述符距离
    fn compute_distance(&self, desc1: &[u8], desc2: &[u8]) -> f64 {
        match self.config.distance_type {
            DistanceType::L2 => {
                desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| (*a as f64 - *b as f64).powi(2))
                    .sum::<f64>().sqrt()
            },
            DistanceType::L1 => {
                desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| (*a as f64 - *b as f64).abs())
                    .sum::<f64>()
            },
            DistanceType::Hamming => {
                desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| (a ^ b).count_ones() as f64)
                    .sum::<f64>()
            },
            DistanceType::Cosine => {
                let dot_product: f64 = desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| *a as f64 * *b as f64)
                    .sum();
                let norm1: f64 = desc1.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
                let norm2: f64 = desc2.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
                
                if norm1 == 0.0 || norm2 == 0.0 {
                    1.0
                } else {
                    1.0 - (dot_product / (norm1 * norm2))
                }
            },
        }
    }
}

impl FeatureMatcher for BruteForceMatcher {
    fn match_features(
        &self,
        features1: &[Feature],
        features2: &[Feature],
    ) -> Result<Vec<FeatureMatch>, ColmapError> {
        let mut matches = Vec::new();
        
        // 暴力匹配：对每个特征点找到最佳匹配
        for (i, feat1) in features1.iter().enumerate() {
            if feat1.descriptor.is_empty() {
                continue;
            }
            
            let mut best_distance = f64::INFINITY;
            let mut best_idx = None;
            
            for (j, feat2) in features2.iter().enumerate() {
                if feat2.descriptor.is_empty() || feat1.descriptor.len() != feat2.descriptor.len() {
                    continue;
                }
                
                let distance = self.compute_distance(&feat1.descriptor, &feat2.descriptor);
                
                if distance < best_distance && distance <= self.config.max_distance {
                    best_distance = distance;
                    best_idx = Some(j);
                }
            }
            
            if let Some(j) = best_idx {
                let match_result = FeatureMatch::new(i, j, best_distance);
                if match_result.confidence >= self.config.min_confidence {
                    matches.push(match_result);
                }
            }
        }
        
        // 按距离排序并限制匹配数量
        matches.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
        matches.truncate(self.config.max_matches);
        
        Ok(matches)
    }
    
    fn name(&self) -> &str {
        "BruteForce"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        let mut params = HashMap::new();
        params.insert("max_distance".to_string(), self.config.max_distance);
        params.insert("max_matches".to_string(), self.config.max_matches as f64);
        params.insert("min_confidence".to_string(), self.config.min_confidence);
        params
    }
    
    fn set_params(&mut self, params: HashMap<String, f64>) -> Result<(), ColmapError> {
        for (key, value) in params {
            match key.as_str() {
                "max_distance" => self.config.max_distance = value,
                "max_matches" => self.config.max_matches = value as usize,
                "min_confidence" => self.config.min_confidence = value,
                _ => return Err(ColmapError::InvalidParameter(format!("Unknown parameter: {}", key))),
            }
        }
        Ok(())
    }
}

/// FLANN 匹配器
pub struct FlannMatcher {
    config: MatcherConfig,
}

impl FlannMatcher {
    pub fn new(config: &MatcherConfig) -> Result<Self, ColmapError> {
        Ok(Self {
            config: config.clone(),
        })
    }
}

impl FeatureMatcher for FlannMatcher {
    fn match_features(
        &self,
        features1: &[Feature],
        features2: &[Feature],
    ) -> Result<Vec<FeatureMatch>, ColmapError> {
        let mut matches = Vec::new();
        
        // 使用 KD-tree 或近似最近邻搜索的简化实现
        for (i, feat1) in features1.iter().enumerate() {
            if feat1.descriptor.is_empty() {
                continue;
            }
            
            let mut best_distance = f64::INFINITY;
            let mut best_idx = None;
            
            for (j, feat2) in features2.iter().enumerate() {
                if feat2.descriptor.is_empty() || feat1.descriptor.len() != feat2.descriptor.len() {
                    continue;
                }
                
                let distance = self.compute_distance(&feat1.descriptor, &feat2.descriptor);
                
                if distance < best_distance && distance <= self.config.max_distance {
                    best_distance = distance;
                    best_idx = Some(j);
                }
            }
            
            if let Some(j) = best_idx {
                let match_result = FeatureMatch::new(i, j, best_distance);
                if match_result.confidence >= self.config.min_confidence {
                    matches.push(match_result);
                }
            }
        }
        
        // 限制匹配数量
        matches.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
        matches.truncate(self.config.max_matches);
        
        Ok(matches)
    }
    
    fn name(&self) -> &str {
        "FLANN"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        let mut params = HashMap::new();
        params.insert("max_distance".to_string(), self.config.max_distance);
        params.insert("max_matches".to_string(), self.config.max_matches as f64);
        params.insert("min_confidence".to_string(), self.config.min_confidence);
        params
    }
    
    fn set_params(&mut self, params: HashMap<String, f64>) -> Result<(), ColmapError> {
        for (key, value) in params {
            match key.as_str() {
                "max_distance" => self.config.max_distance = value,
                "max_matches" => self.config.max_matches = value as usize,
                "min_confidence" => self.config.min_confidence = value,
                _ => return Err(ColmapError::InvalidParameter(format!("Unknown parameter: {}", key))),
            }
        }
        Ok(())
    }
}

impl FlannMatcher {
    /// 计算描述符距离
    fn compute_distance(&self, desc1: &[u8], desc2: &[u8]) -> f64 {
        match self.config.distance_type {
            DistanceType::L2 => {
                desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| (*a as f64 - *b as f64).powi(2))
                    .sum::<f64>().sqrt()
            },
            DistanceType::L1 => {
                desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| (*a as f64 - *b as f64).abs())
                    .sum::<f64>()
            },
            DistanceType::Hamming => {
                desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| (a ^ b).count_ones() as f64)
                    .sum::<f64>()
            },
            DistanceType::Cosine => {
                let dot_product: f64 = desc1.iter().zip(desc2.iter())
                    .map(|(a, b)| *a as f64 * *b as f64)
                    .sum();
                let norm1: f64 = desc1.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
                let norm2: f64 = desc2.iter().map(|x| (*x as f64).powi(2)).sum::<f64>().sqrt();
                
                if norm1 == 0.0 || norm2 == 0.0 {
                    1.0
                } else {
                    1.0 - (dot_product / (norm1 * norm2))
                }
            },
        }
    }
}

/// 比值测试匹配器
pub struct RatioTestMatcher {
    config: MatcherConfig,
    base_matcher: BruteForceMatcher,
}

impl RatioTestMatcher {
    pub fn new(config: &MatcherConfig) -> Result<Self, ColmapError> {
        let base_matcher = BruteForceMatcher::new(config)?;
        
        Ok(Self {
            config: config.clone(),
            base_matcher,
        })
    }
}

impl FeatureMatcher for RatioTestMatcher {
    fn match_features(
        &self,
        features1: &[Feature],
        features2: &[Feature],
    ) -> Result<Vec<FeatureMatch>, ColmapError> {
        let mut matches = Vec::new();
        
        // 对每个特征点找到两个最佳匹配
        for (i, feat1) in features1.iter().enumerate() {
            if feat1.descriptor.is_empty() {
                continue;
            }
            
            let mut distances: Vec<(usize, f64)> = Vec::new();
            
            for (j, feat2) in features2.iter().enumerate() {
                if feat2.descriptor.is_empty() || feat1.descriptor.len() != feat2.descriptor.len() {
                    continue;
                }
                
                let distance = self.base_matcher.compute_distance(&feat1.descriptor, &feat2.descriptor);
                distances.push((j, distance));
            }
            
            // 按距离排序
            distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
            
            // 比值测试：最佳匹配距离 / 次佳匹配距离 < 阈值
            if distances.len() >= 2 {
                let best_distance = distances[0].1;
                let second_best_distance = distances[1].1;
                
                if best_distance <= self.config.max_distance &&
                   best_distance / second_best_distance < self.config.ratio_threshold {
                    let match_result = FeatureMatch::new(i, distances[0].0, best_distance);
                    if match_result.confidence >= self.config.min_confidence {
                        matches.push(match_result);
                    }
                }
            }
        }
        
        // 限制匹配数量
        matches.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
        matches.truncate(self.config.max_matches);
        
        Ok(matches)
    }
    
    fn name(&self) -> &str {
        "RatioTest"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        let mut params = self.base_matcher.params();
        params.insert("ratio_threshold".to_string(), self.config.ratio_threshold);
        params
    }
    
    fn set_params(&mut self, params: HashMap<String, f64>) -> Result<(), ColmapError> {
        if let Some(&ratio) = params.get("ratio_threshold") {
            self.config.ratio_threshold = ratio;
        }
        self.base_matcher.set_params(params)
    }
}

/// 交叉检查匹配器
pub struct CrossCheckMatcher {
    config: MatcherConfig,
    base_matcher: BruteForceMatcher,
}

impl CrossCheckMatcher {
    pub fn new(config: &MatcherConfig) -> Result<Self, ColmapError> {
        let base_matcher = BruteForceMatcher::new(config)?;
        
        Ok(Self {
            config: config.clone(),
            base_matcher,
        })
    }
}

impl FeatureMatcher for CrossCheckMatcher {
    fn match_features(
        &self,
        features1: &[Feature],
        features2: &[Feature],
    ) -> Result<Vec<FeatureMatch>, ColmapError> {
        // 正向匹配：features1 -> features2
        let forward_matches = self.base_matcher.match_features(features1, features2)?;
        
        // 反向匹配：features2 -> features1
        let backward_matches = self.base_matcher.match_features(features2, features1)?;
        
        // 交叉检查：只保留双向一致的匹配
        let mut cross_checked_matches = Vec::new();
        
        for forward_match in forward_matches {
            // 检查是否存在对应的反向匹配
            for backward_match in &backward_matches {
                if backward_match.query_idx == forward_match.train_idx &&
                   backward_match.train_idx == forward_match.query_idx {
                    // 使用较小的距离作为最终距离
                    let final_distance = forward_match.distance.min(backward_match.distance);
                    let cross_match = FeatureMatch::new(
                        forward_match.query_idx,
                        forward_match.train_idx,
                        final_distance,
                    );
                    
                    if cross_match.confidence >= self.config.min_confidence {
                        cross_checked_matches.push(cross_match);
                    }
                    break;
                }
            }
        }
        
        // 限制匹配数量
        cross_checked_matches.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap());
        cross_checked_matches.truncate(self.config.max_matches);
        
        Ok(cross_checked_matches)
    }
    
    fn name(&self) -> &str {
        "CrossCheck"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        self.base_matcher.params()
    }
    
    fn set_params(&mut self, params: HashMap<String, f64>) -> Result<(), ColmapError> {
        self.base_matcher.set_params(params)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::Feature;
    use nalgebra::Point2;
    
    fn create_test_features() -> Vec<Feature> {
        vec![
            Feature {
                point: Point2::new(10.0, 20.0),
                descriptor: vec![1, 2, 3, 4],
                response: 0.8,
                angle: 0.0,
                octave: 0,
                scale: 1.0,
                point3d_id: None,
            },
            Feature {
                point: Point2::new(30.0, 40.0),
                descriptor: vec![5, 6, 7, 8],
                response: 0.9,
                angle: 0.0,
                octave: 0,
                scale: 1.0,
                point3d_id: None,
            },
        ]
    }
    
    #[test]
    fn test_feature_match_creation() {
        let match_result = FeatureMatch::new(0, 1, 0.5);
        assert_eq!(match_result.query_idx, 0);
        assert_eq!(match_result.train_idx, 1);
        assert_eq!(match_result.distance, 0.5);
        assert!(match_result.is_valid(1.0));
        assert!(!match_result.is_valid(0.3));
    }
    
    #[test]
    fn test_matcher_config_default() {
        let config = MatcherConfig::default();
        assert_eq!(config.matcher_type, MatcherType::BruteForce);
        assert_eq!(config.distance_type, DistanceType::L2);
        assert_eq!(config.max_distance, 0.7);
        assert_eq!(config.ratio_threshold, 0.8);
    }
    
    #[test]
    fn test_matcher_factory() {
        let config = MatcherConfig::default();
        let matcher = MatcherFactory::create(&config);
        assert!(matcher.is_ok());
        assert_eq!(matcher.unwrap().name(), "BruteForce");
    }
    
    #[test]
    fn test_brute_force_matcher_creation() {
        let config = MatcherConfig::default();
        let matcher = BruteForceMatcher::new(&config);
        assert!(matcher.is_ok());
    }
    
    #[test]
    fn test_brute_force_matching() {
        let config = MatcherConfig {
            max_distance: 10.0,
            ..Default::default()
        };
        let matcher = BruteForceMatcher::new(&config).unwrap();
        
        let features1 = create_test_features();
        let features2 = create_test_features();
        
        let matches = matcher.match_features(&features1, &features2);
        assert!(matches.is_ok());
        
        let matches = matches.unwrap();
        assert!(!matches.is_empty());
    }
    
    #[test]
    fn test_distance_computation() {
        let config = MatcherConfig {
            distance_type: DistanceType::L2,
            ..Default::default()
        };
        let matcher = BruteForceMatcher::new(&config).unwrap();
        
        let desc1 = vec![1, 2, 3, 4];
        let desc2 = vec![1, 2, 3, 4];
        let distance = matcher.compute_distance(&desc1, &desc2);
        assert_eq!(distance, 0.0);
        
        let desc3 = vec![2, 3, 4, 5];
        let distance = matcher.compute_distance(&desc1, &desc3);
        assert!(distance > 0.0);
    }
}