colmap 0.1.2

//! 相机模型和参数定义
//!
//! 这个模块定义了相机的内参、外参和各种畸变模型。

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

/// 相机 ID 类型
pub type CameraId = u32;

/// 畸变模型枚举
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum DistortionModel {
    /// 无畸变
    None,
    /// 径向畸变 (k1, k2, k3)
    Radial([f64; 3]),
    /// 径向+切向畸变 (k1, k2, p1, p2, k3)
    RadialTangential([f64; 5]),
    /// 鱼眼畸变 (k1, k2, k3, k4)
    Fisheye([f64; 4]),
}

impl DistortionModel {
    /// 获取畸变参数数量
    pub fn param_count(&self) -> usize {
        match self {
            DistortionModel::None => 0,
            DistortionModel::Radial(_) => 3,
            DistortionModel::RadialTangential(_) => 5,
            DistortionModel::Fisheye(_) => 4,
        }
    }

    /// 应用畸变到归一化坐标
    pub fn distort(&self, point: Point2) -> Point2 {
        match self {
            DistortionModel::None => point,
            DistortionModel::Radial(params) => self.apply_radial_distortion(point, params),
            DistortionModel::RadialTangential(params) => {
                self.apply_radial_tangential_distortion(point, params)
            }
            DistortionModel::Fisheye(params) => self.apply_fisheye_distortion(point, params),
        }
    }

    /// 去除畸变（迭代方法）
    pub fn undistort(&self, point: Point2) -> Point2 {
        match self {
            DistortionModel::None => point,
            _ => {
                // 使用牛顿迭代法去畸变
                let mut undistorted = point;
                for _ in 0..10 {
                    let distorted = self.distort(undistorted);
                    let error = point - distorted;
                    if error.norm() < 1e-10 {
                        break;
                    }
                    undistorted += error;
                }
                undistorted
            }
        }
    }

    fn apply_radial_distortion(&self, point: Point2, params: &[f64; 3]) -> Point2 {
        let [k1, k2, k3] = *params;
        let r2 = point.x * point.x + point.y * point.y;
        let r4 = r2 * r2;
        let r6 = r4 * r2;
        let radial_factor = 1.0 + k1 * r2 + k2 * r4 + k3 * r6;
        Point2::new(point.x * radial_factor, point.y * radial_factor)
    }

    fn apply_radial_tangential_distortion(&self, point: Point2, params: &[f64; 5]) -> Point2 {
        let [k1, k2, p1, p2, k3] = *params;
        let x = point.x;
        let y = point.y;
        let r2 = x * x + y * y;
        let r4 = r2 * r2;
        let r6 = r4 * r2;

        // 径向畸变
        let radial_factor = 1.0 + k1 * r2 + k2 * r4 + k3 * r6;

        // 切向畸变
        let dx_tangential = 2.0 * p1 * x * y + p2 * (r2 + 2.0 * x * x);
        let dy_tangential = p1 * (r2 + 2.0 * y * y) + 2.0 * p2 * x * y;

        Point2::new(
            x * radial_factor + dx_tangential,
            y * radial_factor + dy_tangential,
        )
    }

    fn apply_fisheye_distortion(&self, point: Point2, params: &[f64; 4]) -> Point2 {
        let [k1, k2, k3, k4] = *params;
        let r = (point.x * point.x + point.y * point.y).sqrt();

        if r < 1e-8 {
            return point;
        }

        let theta = r.atan();
        let theta2 = theta * theta;
        let theta4 = theta2 * theta2;
        let theta6 = theta4 * theta2;
        let theta8 = theta4 * theta4;

        let theta_d = theta * (1.0 + k1 * theta2 + k2 * theta4 + k3 * theta6 + k4 * theta8);
        let scale = theta_d / r;

        Point2::new(point.x * scale, point.y * scale)
    }
}

/// 相机内参
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CameraIntrinsics {
    /// 焦距 (fx, fy)
    pub focal_length: (f64, f64),
    /// 主点 (cx, cy)
    pub principal_point: (f64, f64),
    /// 畸变模型
    pub distortion: DistortionModel,
}

impl CameraIntrinsics {
    /// 创建新的相机内参
    pub fn new(fx: f64, fy: f64, cx: f64, cy: f64, distortion: DistortionModel) -> Self {
        Self {
            focal_length: (fx, fy),
            principal_point: (cx, cy),
            distortion,
        }
    }

    /// 创建简单的针孔相机模型
    pub fn pinhole(fx: f64, fy: f64, cx: f64, cy: f64) -> Self {
        Self::new(fx, fy, cx, cy, DistortionModel::None)
    }

    /// 将 3D 点投影到图像平面
    pub fn project(&self, point: Point2) -> Point2 {
        // 应用畸变
        let distorted = self.distortion.distort(point);

        // 应用内参矩阵
        Point2::new(
            self.focal_length.0 * distorted.x + self.principal_point.0,
            self.focal_length.1 * distorted.y + self.principal_point.1,
        )
    }

    /// 将图像点反投影到归一化坐标
    pub fn unproject(&self, pixel: Point2) -> Point2 {
        // 转换到归一化坐标
        let normalized = Point2::new(
            (pixel.x - self.principal_point.0) / self.focal_length.0,
            (pixel.y - self.principal_point.1) / self.focal_length.1,
        );

        // 去除畸变
        self.distortion.undistort(normalized)
    }

    /// 获取内参矩阵
    pub fn matrix(&self) -> nalgebra::Matrix3<f64> {
        nalgebra::Matrix3::new(
            self.focal_length.0,
            0.0,
            self.principal_point.0,
            0.0,
            self.focal_length.1,
            self.principal_point.1,
            0.0,
            0.0,
            1.0,
        )
    }
}

/// 相机外参（姿态）
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct CameraPose {
    /// 旋转（四元数表示）
    pub rotation: Quaternion,
    /// 平移向量
    pub translation: Vector3,
}

impl CameraPose {
    /// 创建新的相机姿态
    pub fn new(rotation: Quaternion, translation: Vector3) -> Self {
        Self {
            rotation,
            translation,
        }
    }

    /// 创建单位姿态（无旋转无平移）
    pub fn identity() -> Self {
        Self {
            rotation: Quaternion::identity(),
            translation: Vector3::zeros(),
        }
    }

    /// 从旋转矩阵和平移向量创建
    pub fn from_matrix(rotation: nalgebra::Matrix3<f64>, translation: Vector3) -> Result<Self> {
        let rotation = Quaternion::from_matrix(&rotation);
        Ok(Self::new(rotation, translation))
    }

    /// 获取变换矩阵
    pub fn matrix(&self) -> nalgebra::Matrix4<f64> {
        let mut transform = nalgebra::Matrix4::identity();
        transform
            .fixed_view_mut::<3, 3>(0, 0)
            .copy_from(self.rotation.to_rotation_matrix().matrix());
        transform
            .fixed_view_mut::<3, 1>(0, 3)
            .copy_from(&self.translation);
        transform
    }

    /// 获取逆变换
    pub fn inverse(&self) -> Self {
        let inv_rotation = self.rotation.inverse();
        let inv_translation = -(inv_rotation * self.translation);
        Self::new(inv_rotation, inv_translation)
    }

    /// 组合两个变换
    pub fn compose(&self, other: &CameraPose) -> Self {
        let rotation = self.rotation * other.rotation;
        let translation = self.rotation * other.translation + self.translation;
        Self::new(rotation, translation)
    }

    /// 将 3D 点从世界坐标系变换到相机坐标系
    pub fn transform_point(&self, point: nalgebra::Point3<f64>) -> nalgebra::Point3<f64> {
        let rotated = self.rotation * point.coords;
        nalgebra::Point3::from(rotated + self.translation)
    }

    /// 将 3D 点从相机坐标系变换到世界坐标系
    pub fn inverse_transform_point(&self, point: nalgebra::Point3<f64>) -> nalgebra::Point3<f64> {
        let translated = point.coords - self.translation;
        let rotated = self.rotation.inverse() * translated;
        nalgebra::Point3::from(rotated)
    }

    /// 获取 3x4 投影矩阵 [R|t]
    pub fn to_matrix3x4(&self) -> nalgebra::Matrix3x4<f64> {
        let mut rt = nalgebra::Matrix3x4::zeros();
        rt.fixed_view_mut::<3, 3>(0, 0)
            .copy_from(self.rotation.to_rotation_matrix().matrix());
        rt.fixed_view_mut::<3, 1>(0, 3).copy_from(&self.translation);
        rt
    }
}

/// 完整的相机模型
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct Camera {
    /// 相机 ID
    pub id: CameraId,
    /// 相机内参
    pub intrinsics: CameraIntrinsics,
    /// 相机外参（可选，只有在重建中注册的相机才有）
    pub pose: Option<CameraPose>,
    /// 图像尺寸 (width, height)
    pub image_size: (u32, u32),
    /// 相机型号名称
    pub model_name: String,
}

impl Camera {
    /// 创建新的相机
    pub fn new(
        id: CameraId,
        intrinsics: CameraIntrinsics,
        image_size: (u32, u32),
        model_name: String,
    ) -> Self {
        Self {
            id,
            intrinsics,
            pose: None,
            image_size,
            model_name,
        }
    }

    /// 设置相机姿态
    pub fn set_pose(&mut self, pose: CameraPose) {
        self.pose = Some(pose);
    }

    /// 检查相机是否已注册（有姿态信息）
    pub fn is_registered(&self) -> bool {
        self.pose.is_some()
    }

    /// 将 3D 点投影到图像平面
    pub fn project_point(&self, world_point: nalgebra::Point3<f64>) -> Result<Point2> {
        let pose = self
            .pose
            .as_ref()
            .ok_or_else(|| ColmapError::Math("Camera pose not set".to_string()))?;

        // 变换到相机坐标系
        let camera_point = pose.transform_point(world_point);

        // 检查点是否在相机前方
        if camera_point.z <= 0.0 {
            return Err(ColmapError::Math("Point behind camera".to_string()));
        }

        // 投影到归一化平面
        let normalized = Point2::new(
            camera_point.x / camera_point.z,
            camera_point.y / camera_point.z,
        );

        // 应用内参
        Ok(self.intrinsics.project(normalized))
    }

    /// 将图像点反投影为射线方向
    pub fn unproject_ray(&self, pixel: Point2) -> Result<Vector3> {
        let pose = self
            .pose
            .as_ref()
            .ok_or_else(|| ColmapError::Math("Camera pose not set".to_string()))?;

        // 反投影到归一化坐标
        let normalized = self.intrinsics.unproject(pixel);

        // 构造射线方向（相机坐标系）
        let ray_camera = Vector3::new(normalized.x, normalized.y, 1.0).normalize();

        // 变换到世界坐标系
        Ok(pose.rotation.inverse() * ray_camera)
    }

    /// 获取相机中心（世界坐标系）
    pub fn center(&self) -> Option<nalgebra::Point3<f64>> {
        self.pose
            .as_ref()
            .map(|pose| pose.inverse_transform_point(nalgebra::Point3::origin()))
    }

    /// 获取相机标定矩阵
    pub fn calibration_matrix(&self) -> nalgebra::Matrix3<f64> {
        self.intrinsics.matrix()
    }

    /// 将图像坐标转换为归一化坐标
    pub fn image_to_normalized(&self, pixel: &Point2) -> Result<Point2> {
        Ok(self.intrinsics.unproject(*pixel))
    }

    /// 将世界坐标点投影到图像坐标
    pub fn world_to_image(&self, world_point: &nalgebra::Point3<f64>) -> Result<Point2> {
        let pose = self
            .pose
            .as_ref()
            .ok_or_else(|| ColmapError::Math("相机姿态未设置".to_string()))?;

        // 将世界坐标转换到相机坐标系
        let camera_point = pose.transform_point(*world_point);

        // 检查点是否在相机前方
        if camera_point.z <= 0.0 {
            return Err(ColmapError::Math("点在相机后方".to_string()));
        }

        // 投影到归一化平面
        let normalized = Point2::new(
            camera_point.x / camera_point.z,
            camera_point.y / camera_point.z,
        );

        // 应用畸变
        let distorted = self.intrinsics.distortion.distort(normalized);

        // 转换到图像坐标
        Ok(self.intrinsics.project(distorted))
    }
}

/// 相机数据库，管理多个相机
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CameraDatabase {
    cameras: HashMap<CameraId, Camera>,
    next_id: CameraId,
}

impl CameraDatabase {
    /// 创建新的相机数据库
    pub fn new() -> Self {
        Self {
            cameras: HashMap::new(),
            next_id: 1,
        }
    }

    /// 添加相机
    pub fn add_camera(&mut self, mut camera: Camera) -> CameraId {
        if camera.id == 0 {
            camera.id = self.next_id;
            self.next_id += 1;
        }
        let id = camera.id;
        self.cameras.insert(id, camera);
        id
    }

    /// 获取相机
    pub fn get_camera(&self, id: CameraId) -> Option<&Camera> {
        self.cameras.get(&id)
    }

    /// 获取可变相机引用
    pub fn get_camera_mut(&mut self, id: CameraId) -> Option<&mut Camera> {
        self.cameras.get_mut(&id)
    }

    /// 获取所有相机
    pub fn cameras(&self) -> &HashMap<CameraId, Camera> {
        &self.cameras
    }

    /// 获取已注册的相机数量
    pub fn registered_count(&self) -> usize {
        self.cameras.values().filter(|c| c.is_registered()).count()
    }

    /// 移除相机
    pub fn remove_camera(&mut self, id: CameraId) -> Option<Camera> {
        self.cameras.remove(&id)
    }

    /// 清空所有相机
    pub fn clear(&mut self) {
        self.cameras.clear();
        self.next_id = 1;
    }
}

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

    #[test]
    fn test_camera_intrinsics_project_unproject() {
        let intrinsics = CameraIntrinsics::pinhole(800.0, 800.0, 320.0, 240.0);

        let normalized = Point2::new(0.1, 0.2);
        let pixel = intrinsics.project(normalized);
        let recovered = intrinsics.unproject(pixel);

        assert!((normalized - recovered).norm() < 1e-10);
    }

    #[test]
    fn test_camera_pose_transform() {
        let pose = CameraPose::identity();
        let point = Point3::new(1.0, 2.0, 3.0);

        let transformed = pose.transform_point(point);
        let recovered = pose.inverse_transform_point(transformed);

        assert!((point.coords - recovered.coords).norm() < 1e-10);
    }

    #[test]
    fn test_distortion_model() {
        let distortion = DistortionModel::Radial([0.1, 0.01, 0.001]);
        let point = Point2::new(0.5, 0.3);

        let distorted = distortion.distort(point);
        let undistorted = distortion.undistort(distorted);

        assert!((point - undistorted).norm() < 1e-6);
    }
}