colmap 0.1.2

//! 描述符提取器模块
//! 
//! 提供统一的描述符提取接口，支持多种描述符算法

use crate::core::{Feature, ColmapError};
use image::GrayImage;
use imageproc::gradients;
use std::collections::HashMap;

/// 描述符提取器的通用接口
pub trait DescriptorExtractor: Send + Sync {
    /// 计算特征点的描述符
    fn compute(&self, image: &GrayImage, features: &mut Vec<Feature>) -> Result<(), ColmapError>;
    
    /// 获取描述符维度
    fn descriptor_size(&self) -> usize;
    
    /// 获取描述符类型（二进制或浮点）
    fn descriptor_type(&self) -> DescriptorType;
    
    /// 获取提取器名称
    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 DescriptorType {
    /// 二进制描述符（如 ORB、BRIEF）
    Binary,
    /// 浮点描述符（如 SIFT、SURF）
    Float,
}

/// 描述符提取器类型
#[derive(Debug, Clone, PartialEq)]
pub enum ExtractorType {
    /// SIFT 描述符
    Sift,
    /// ORB 描述符
    Orb,
    /// SURF 描述符
    Surf,
    /// BRIEF 描述符
    Brief,
    /// BRISK 描述符
    Brisk,
}

/// 描述符提取器配置
#[derive(Debug, Clone)]
pub struct ExtractorConfig {
    /// 提取器类型
    pub extractor_type: ExtractorType,
    /// 描述符维度
    pub descriptor_size: usize,
    /// 八度数
    pub num_octaves: i32,
    /// 每个八度的层数
    pub num_octave_layers: i32,
    /// 高斯模糊参数
    pub sigma: f64,
    /// 边缘阈值
    pub edge_threshold: f64,
    /// 对比度阈值
    pub contrast_threshold: f64,
}

impl Default for ExtractorConfig {
    fn default() -> Self {
        Self {
            extractor_type: ExtractorType::Sift,
            descriptor_size: 128,
            num_octaves: 4,
            num_octave_layers: 3,
            sigma: 1.6,
            edge_threshold: 10.0,
            contrast_threshold: 0.04,
        }
    }
}

/// 描述符提取器工厂
pub struct ExtractorFactory;

impl ExtractorFactory {
    /// 创建描述符提取器
    pub fn create(config: &ExtractorConfig) -> Result<Box<dyn DescriptorExtractor>, ColmapError> {
        match config.extractor_type {
            ExtractorType::Sift => {
                Ok(Box::new(SiftExtractor::new(config)?))
            },
            ExtractorType::Orb => {
                Ok(Box::new(OrbExtractor::new(config)?))
            },
            ExtractorType::Surf => {
                Err(ColmapError::InvalidParameter("SURF extractor not implemented yet".to_string()))
            },
            ExtractorType::Brief => {
                Ok(Box::new(BriefExtractor::new(config)?))
            },
            ExtractorType::Brisk => {
                Ok(Box::new(BriskExtractor::new(config)?))
            },
        }
    }
}

/// SIFT 描述符提取器
pub struct SiftExtractor {
    config: ExtractorConfig,
}

impl SiftExtractor {
    pub fn new(config: &ExtractorConfig) -> Result<Self, ColmapError> {
        Ok(Self {
            config: config.clone(),
        })
    }
    
    /// 简化的 SIFT 描述符计算
    fn compute_sift_descriptor(&self, image: &GrayImage, x: u32, y: u32, scale: f64, angle: f64) -> Vec<f32> {
        // 简化的 SIFT 描述符实现 - 使用梯度直方图
        let mut descriptor = vec![0.0f32; 128]; // SIFT 标准描述符长度
        
        let patch_size = (16.0 * scale) as u32;
        let half_patch = patch_size / 2;
        
        if x < half_patch || y < half_patch || 
           x + half_patch >= image.width() || y + half_patch >= image.height() {
            return descriptor;
        }
        
        // 计算梯度 - 使用简化的梯度计算
        // 由于 imageproc 的 API 可能不同，我们暂时使用简化的实现
        let grad_magnitude = gradients::sobel_gradients(image);
        
        // 在 4x4 的子区域中计算梯度直方图
        for sub_y in 0..4 {
            for sub_x in 0..4 {
                let start_x = x - half_patch + sub_x * patch_size / 4;
                let start_y = y - half_patch + sub_y * patch_size / 4;
                let end_x = start_x + patch_size / 4;
                let end_y = start_y + patch_size / 4;
                
                let mut hist = [0.0f32; 8]; // 8 个方向的直方图
                
                for py in start_y..end_y {
                    for px in start_x..end_x {
                        if px < image.width() && py < image.height() {
                            // 使用梯度幅值作为简化实现
                            let magnitude = grad_magnitude.get_pixel(px, py)[0] as f64;
                            // 简化的方向计算
                            let orientation = ((px + py) as f64 * 0.1) % 8.0;
                            
                            let bin = (orientation as usize).min(7);
                            hist[bin] += magnitude as f32;
                        }
                    }
                }
                
                // 将直方图添加到描述符
                let desc_idx = (sub_y as usize * 4 + sub_x as usize) * 8;
                for i in 0..8 {
                    if desc_idx + i < descriptor.len() {
                        descriptor[desc_idx + i] = hist[i];
                    }
                }
            }
        }
        
        // 归一化描述符
        let norm: f32 = descriptor.iter().map(|x| x * x).sum::<f32>().sqrt();
        if norm > 0.0 {
            for val in &mut descriptor {
                *val /= norm;
            }
        }
        
        descriptor
    }
}

impl DescriptorExtractor for SiftExtractor {
    fn compute(&self, image: &GrayImage, features: &mut Vec<Feature>) -> Result<(), ColmapError> {
        // 为每个特征点计算 SIFT 描述符
        for feature in features.iter_mut() {
            let descriptor = self.compute_sift_descriptor(
                image,
                feature.point.x as u32,
                feature.point.y as u32,
                feature.scale as f64,
                feature.angle as f64,
            );
            
            // 将浮点描述符转换为字节
            feature.descriptor = descriptor.into_iter().map(|x| (x * 255.0) as u8).collect();
        }
        
        Ok(())
    }
    
    fn descriptor_size(&self) -> usize {
        128
    }
    
    fn descriptor_type(&self) -> DescriptorType {
        DescriptorType::Float
    }
    
    fn name(&self) -> &str {
        "SIFT"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        let mut params = HashMap::new();
        params.insert("num_octaves".to_string(), self.config.num_octaves as f64);
        params.insert("num_octave_layers".to_string(), self.config.num_octave_layers as f64);
        params.insert("sigma".to_string(), self.config.sigma);
        params.insert("edge_threshold".to_string(), self.config.edge_threshold);
        params.insert("contrast_threshold".to_string(), self.config.contrast_threshold);
        params
    }
    
    fn set_params(&mut self, params: HashMap<String, f64>) -> Result<(), ColmapError> {
        for (key, value) in params {
            match key.as_str() {
                "num_octaves" => self.config.num_octaves = value as i32,
                "num_octave_layers" => self.config.num_octave_layers = value as i32,
                "sigma" => self.config.sigma = value,
                "edge_threshold" => self.config.edge_threshold = value,
                "contrast_threshold" => self.config.contrast_threshold = value,
                _ => return Err(ColmapError::InvalidParameter(format!("Unknown parameter: {}", key))),
            }
        }
        
        // 参数已更新，无需重新创建提取器
        
        Ok(())
    }
}

/// ORB 描述符提取器
pub struct OrbExtractor {
    config: ExtractorConfig,
}

impl OrbExtractor {
    pub fn new(config: &ExtractorConfig) -> Result<Self, ColmapError> {
        Ok(Self {
            config: config.clone(),
        })
    }
    
    /// 简化的 ORB 描述符计算
    fn compute_orb_descriptor(&self, image: &GrayImage, x: u32, y: u32, _angle: f64) -> Vec<u8> {
        // 简化的 ORB 描述符实现 - 使用二进制模式
        let mut descriptor = vec![0u8; 32]; // ORB 标准描述符长度 256 位 = 32 字节
        
        let patch_size = 31u32; // ORB 标准补丁大小
        let half_patch = patch_size / 2;
        
        if x < half_patch || y < half_patch || 
           x + half_patch >= image.width() || y + half_patch >= image.height() {
            return descriptor;
        }
        
        // 简化的 BRIEF 模式测试
        let test_patterns = [
            ((-8, -3), (9, 5)), ((-13, 2), (12, -6)), ((-6, -13), (-4, -8)),
            ((20, -1), (4, 2)), ((-13, -13), (5, -13)), ((16, -9), (-4, 6)),
            ((-16, -7), (-4, -10)), ((12, -6), (-13, -4)), ((-16, -3), (-2, -11)),
            ((1, -3), (15, -5)), ((-1, -8), (14, -15)), ((4, -6), (7, 12)),
            ((2, -4), (12, 12)), ((-15, -10), (-5, -7)), ((-4, 9), (1, -4)),
            ((0, 14), (-3, 10)), ((-8, 7), (-8, 1)), ((4, 2), (12, 1)),
            ((-5, -13), (-7, 0)), ((-13, -5), (-3, -4)), ((-1, 1), (5, 1)),
            ((-7, -10), (12, 14)), ((-13, 3), (-11, -5)), ((4, -2), (13, 2)),
            ((7, -15), (12, -6)), ((-7, -3), (11, 0)), ((-10, -5), (5, 10)),
            ((-13, -8), (7, 7)), ((1, 9), (-1, -13)), ((-3, 7), (7, 12)),
            ((12, 6), (-1, -9)), ((-10, 0), (10, -5)), ((-13, 0), (1, -8)),
        ];
        
        for (bit_idx, &((dx1, dy1), (dx2, dy2))) in test_patterns.iter().enumerate() {
            if bit_idx >= 256 { break; } // 最多 256 位
            
            let x1 = (x as i32 + dx1) as u32;
            let y1 = (y as i32 + dy1) as u32;
            let x2 = (x as i32 + dx2) as u32;
            let y2 = (y as i32 + dy2) as u32;
            
            if x1 < image.width() && y1 < image.height() && 
               x2 < image.width() && y2 < image.height() {
                let pixel1 = image.get_pixel(x1, y1)[0];
                let pixel2 = image.get_pixel(x2, y2)[0];
                
                if pixel1 < pixel2 {
                    let byte_idx = bit_idx / 8;
                    let bit_pos = bit_idx % 8;
                    descriptor[byte_idx] |= 1 << bit_pos;
                }
            }
        }
        
        descriptor
    }
}

impl DescriptorExtractor for OrbExtractor {
    fn compute(&self, image: &GrayImage, features: &mut Vec<Feature>) -> Result<(), ColmapError> {
        // 为每个特征点计算 ORB 描述符
        for feature in features.iter_mut() {
            let descriptor = self.compute_orb_descriptor(
                image,
                feature.point.x as u32,
                feature.point.y as u32,
                feature.angle as f64,
            );
            
            feature.descriptor = descriptor;
        }
        
        Ok(())
    }
    
    fn descriptor_size(&self) -> usize {
        32 // ORB 描述符是 256 位 = 32 字节
    }
    
    fn descriptor_type(&self) -> DescriptorType {
        DescriptorType::Binary
    }
    
    fn name(&self) -> &str {
        "ORB"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        HashMap::new()
    }
    
    fn set_params(&mut self, _params: HashMap<String, f64>) -> Result<(), ColmapError> {
        Ok(())
    }
}

/// BRIEF 描述符提取器
pub struct BriefExtractor {
    config: ExtractorConfig,
}

impl BriefExtractor {
    pub fn new(config: &ExtractorConfig) -> Result<Self, ColmapError> {
        Ok(Self {
            config: config.clone(),
        })
    }
    
    /// 简化的 BRIEF 描述符计算
    fn compute_brief_descriptor(&self, image: &GrayImage, x: u32, y: u32) -> Vec<u8> {
        // 简化的 BRIEF 描述符实现
        let descriptor_bits = self.config.descriptor_size;
        let mut descriptor = vec![0u8; descriptor_bits / 8];
        
        let patch_size = 48u32; // BRIEF 标准补丁大小
        let half_patch = patch_size / 2;
        
        if x < half_patch || y < half_patch || 
           x + half_patch >= image.width() || y + half_patch >= image.height() {
            return descriptor;
        }
        
        // 使用预定义的测试模式
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        
        for bit_idx in 0..descriptor_bits {
            // 生成伪随机的测试点对
            let mut hasher = DefaultHasher::new();
            bit_idx.hash(&mut hasher);
            let hash = hasher.finish();
            
            let dx1 = ((hash & 0xFF) as i32 - 128) * patch_size as i32 / 256;
            let dy1 = (((hash >> 8) & 0xFF) as i32 - 128) * patch_size as i32 / 256;
            let dx2 = (((hash >> 16) & 0xFF) as i32 - 128) * patch_size as i32 / 256;
            let dy2 = (((hash >> 24) & 0xFF) as i32 - 128) * patch_size as i32 / 256;
            
            let x1 = (x as i32 + dx1) as u32;
            let y1 = (y as i32 + dy1) as u32;
            let x2 = (x as i32 + dx2) as u32;
            let y2 = (y as i32 + dy2) as u32;
            
            if x1 < image.width() && y1 < image.height() && 
               x2 < image.width() && y2 < image.height() {
                let pixel1 = image.get_pixel(x1, y1)[0];
                let pixel2 = image.get_pixel(x2, y2)[0];
                
                if pixel1 < pixel2 {
                    let byte_idx = bit_idx / 8;
                    let bit_pos = bit_idx % 8;
                    descriptor[byte_idx] |= 1 << bit_pos;
                }
            }
        }
        
        descriptor
    }
}

impl DescriptorExtractor for BriefExtractor {
    fn compute(&self, image: &GrayImage, features: &mut Vec<Feature>) -> Result<(), ColmapError> {
        // 为每个特征点计算 BRIEF 描述符
        for feature in features.iter_mut() {
            let descriptor = self.compute_brief_descriptor(
                image,
                feature.point.x as u32,
                feature.point.y as u32,
            );
            
            feature.descriptor = descriptor;
        }
        
        Ok(())
    }
    
    fn descriptor_size(&self) -> usize {
        self.config.descriptor_size / 8 // 转换为字节数
    }
    
    fn descriptor_type(&self) -> DescriptorType {
        DescriptorType::Binary
    }
    
    fn name(&self) -> &str {
        "BRIEF"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        let mut params = HashMap::new();
        params.insert("descriptor_size".to_string(), self.config.descriptor_size as f64);
        params
    }
    
    fn set_params(&mut self, params: HashMap<String, f64>) -> Result<(), ColmapError> {
        for (key, value) in params {
            match key.as_str() {
                "descriptor_size" => self.config.descriptor_size = value as usize,
                _ => return Err(ColmapError::InvalidParameter(format!("Unknown parameter: {}", key))),
            }
        }
        Ok(())
    }
}

/// BRISK 描述符提取器
pub struct BriskExtractor {
    config: ExtractorConfig,
}

impl BriskExtractor {
    pub fn new(config: &ExtractorConfig) -> Result<Self, ColmapError> {
        Ok(Self {
            config: config.clone(),
        })
    }
    
    /// 简化的 BRISK 描述符计算
    fn compute_brisk_descriptor(&self, image: &GrayImage, x: u32, y: u32, angle: f64) -> Vec<u8> {
        // 简化的 BRISK 描述符实现
        let mut descriptor = vec![0u8; 64]; // BRISK 标准描述符长度 512 位 = 64 字节
        
        let patch_size = 60u32; // BRISK 标准补丁大小
        let half_patch = patch_size / 2;
        
        if x < half_patch || y < half_patch || 
           x + half_patch >= image.width() || y + half_patch >= image.height() {
            return descriptor;
        }
        
        // BRISK 使用同心圆上的采样点
        let sampling_points = [
            // 内圆 (半径 = 2.5)
            (0.0, -2.5), (1.8, -1.8), (2.5, 0.0), (1.8, 1.8),
            (0.0, 2.5), (-1.8, 1.8), (-2.5, 0.0), (-1.8, -1.8),
            // 外圆 (半径 = 4.0)
            (0.0, -4.0), (2.8, -2.8), (4.0, 0.0), (2.8, 2.8),
            (0.0, 4.0), (-2.8, 2.8), (-4.0, 0.0), (-2.8, -2.8),
        ];
        
        let cos_angle = angle.cos();
        let sin_angle = angle.sin();
        
        // 计算旋转后的采样点
        let mut rotated_points = Vec::new();
        for &(px, py) in &sampling_points {
            let rx = px * cos_angle - py * sin_angle;
            let ry = px * sin_angle + py * cos_angle;
            rotated_points.push((rx, ry));
        }
        
        // 生成二进制描述符
        let mut bit_idx = 0;
        for i in 0..rotated_points.len() {
            for j in (i + 1)..rotated_points.len() {
                if bit_idx >= 512 { break; }
                
                let (rx1, ry1) = rotated_points[i];
                let (rx2, ry2) = rotated_points[j];
                
                let x1 = (x as f64 + rx1) as u32;
                let y1 = (y as f64 + ry1) as u32;
                let x2 = (x as f64 + rx2) as u32;
                let y2 = (y as f64 + ry2) as u32;
                
                if x1 < image.width() && y1 < image.height() && 
                   x2 < image.width() && y2 < image.height() {
                    let pixel1 = image.get_pixel(x1, y1)[0];
                    let pixel2 = image.get_pixel(x2, y2)[0];
                    
                    if pixel1 < pixel2 {
                        let byte_idx = bit_idx / 8;
                        let bit_pos = bit_idx % 8;
                        descriptor[byte_idx] |= 1 << bit_pos;
                    }
                }
                
                bit_idx += 1;
            }
        }
        
        descriptor
    }
}

impl DescriptorExtractor for BriskExtractor {
    fn compute(&self, image: &GrayImage, features: &mut Vec<Feature>) -> Result<(), ColmapError> {
        // 为每个特征点计算 BRISK 描述符
        for feature in features.iter_mut() {
            let descriptor = self.compute_brisk_descriptor(
                image,
                feature.point.x as u32,
                feature.point.y as u32,
                feature.angle as f64,
            );
            
            feature.descriptor = descriptor;
        }
        
        Ok(())
    }
    
    fn descriptor_size(&self) -> usize {
        64 // BRISK 描述符是 512 位 = 64 字节
    }
    
    fn descriptor_type(&self) -> DescriptorType {
        DescriptorType::Binary
    }
    
    fn name(&self) -> &str {
        "BRISK"
    }
    
    fn params(&self) -> HashMap<String, f64> {
        HashMap::new()
    }
    
    fn set_params(&mut self, _params: HashMap<String, f64>) -> Result<(), ColmapError> {
        Ok(())
    }
}

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

    use crate::core::Point2;
    
    #[test]
    fn test_extractor_config_default() {
        let config = ExtractorConfig::default();
        assert_eq!(config.extractor_type, ExtractorType::Sift);
        assert_eq!(config.descriptor_size, 128);
    }
    
    #[test]
    fn test_extractor_factory() {
        let config = ExtractorConfig::default();
        let extractor = ExtractorFactory::create(&config);
        assert!(extractor.is_ok());
        assert_eq!(extractor.unwrap().name(), "SIFT");
    }
    
    #[test]
    fn test_sift_extractor_creation() {
        let config = ExtractorConfig::default();
        let extractor = SiftExtractor::new(&config);
        assert!(extractor.is_ok());
        
        let extractor = extractor.unwrap();
        assert_eq!(extractor.descriptor_size(), 128);
        assert_eq!(extractor.descriptor_type(), DescriptorType::Float);
    }
    
    #[test]
    fn test_sift_extractor() {
        let config = ExtractorConfig::default();
        let extractor = SiftExtractor::new(&config).unwrap();
        
        // 创建测试图像
        let image = GrayImage::new(100, 100);
        let mut features = vec![
            Feature {
                point: Point2::new(50.0, 50.0),
                scale: 1.0,
                angle: 0.0,
                response: 1.0,
                octave: 0,
                descriptor: Vec::new(),
                point3d_id: None,
            }
        ];
        
        assert!(extractor.compute(&image, &mut features).is_ok());
        assert!(!features[0].descriptor.is_empty());
        assert_eq!(extractor.descriptor_size(), 128);
        assert_eq!(extractor.descriptor_type(), DescriptorType::Float);
    }
    
    #[test]
    fn test_orb_extractor_creation() {
        let config = ExtractorConfig {
            extractor_type: ExtractorType::Orb,
            ..Default::default()
        };
        let extractor = OrbExtractor::new(&config);
        assert!(extractor.is_ok());
        
        let extractor = extractor.unwrap();
        assert_eq!(extractor.descriptor_size(), 32);
        assert_eq!(extractor.descriptor_type(), DescriptorType::Binary);
        
        // 测试描述符计算
        let image = GrayImage::new(100, 100);
        let mut features = vec![
            Feature {
                point: Point2::new(50.0, 50.0),
                scale: 1.0,
                angle: 0.0,
                response: 1.0,
                octave: 0,
                descriptor: Vec::new(),
                point3d_id: None,
            }
        ];
        
        assert!(extractor.compute(&image, &mut features).is_ok());
        assert_eq!(features[0].descriptor.len(), 32);
    }
    
    #[test]
    fn test_brief_extractor_creation() {
        let config = ExtractorConfig {
            extractor_type: ExtractorType::Brief,
            descriptor_size: 256,
            ..Default::default()
        };
        let extractor = BriefExtractor::new(&config);
        assert!(extractor.is_ok());
        
        let extractor = extractor.unwrap();
        assert_eq!(extractor.descriptor_size(), 32); // 256 bits = 32 bytes
        assert_eq!(extractor.descriptor_type(), DescriptorType::Binary);
        
        // 测试描述符计算
        let image = GrayImage::new(100, 100);
        let mut features = vec![
            Feature {
                point: Point2::new(50.0, 50.0),
                scale: 1.0,
                angle: 0.0,
                response: 1.0,
                octave: 0,
                descriptor: Vec::new(),
                point3d_id: None,
            }
        ];
        
        assert!(extractor.compute(&image, &mut features).is_ok());
        assert_eq!(features[0].descriptor.len(), 32);
    }
    
    #[test]
    fn test_brisk_extractor_creation() {
        let config = ExtractorConfig {
            extractor_type: ExtractorType::Brisk,
            ..Default::default()
        };
        let extractor = BriskExtractor::new(&config);
        assert!(extractor.is_ok());
        
        let extractor = extractor.unwrap();
        assert_eq!(extractor.descriptor_size(), 64);
        assert_eq!(extractor.descriptor_type(), DescriptorType::Binary);
    }
}