apex-io 0.2.0

use nalgebra::{Matrix3, Matrix6, Quaternion, UnitQuaternion, Vector3};

#[cfg(feature = "visualization")]
use rerun::external::glam::{Quat, Vec3};

use std::{
    collections, fmt,
    fmt::{Display, Formatter},
    io,
    path::Path,
};
use thiserror::Error;
use tracing::error;

// Import manifold types from apex-manifolds crate
use apex_manifolds::{se2::SE2, se3::SE3};

// Module declarations
pub mod bal;
pub mod g2o;
pub mod logger;
pub mod toro;
pub mod utils;

pub mod rosbag;

#[cfg(feature = "dds")]
pub mod dds;

pub use logger::init_logger;
pub use utils::{DatasetRegistry, ensure_ba_dataset, ensure_odometry_dataset};

/// Default base directory for odometry (pose graph) datasets relative to the workspace root.
pub const ODOMETRY_DATA_DIR: &str = "data/odometry";

/// Directory for 2D odometry datasets (`data/odometry/2d`).
pub const ODOMETRY_DATA_DIR_2D: &str = "data/odometry/2d";

/// Directory for 3D odometry datasets (`data/odometry/3d`).
pub const ODOMETRY_DATA_DIR_3D: &str = "data/odometry/3d";

/// Default directory for bundle adjustment datasets relative to the workspace root.
pub const BUNDLE_ADJUSTMENT_DATA_DIR: &str = "data/bundle_adjustment";

// Re-exports
pub use bal::{BalCamera, BalDataset, BalLoader, BalObservation, BalPoint};
pub use g2o::G2oLoader;
pub use toro::ToroLoader;

/// Errors that can occur during graph file parsing
#[derive(Error, Debug)]
pub enum IoError {
    #[error("IO error: {0}")]
    Io(#[from] io::Error),

    #[error("Parse error at line {line}: {message}")]
    Parse { line: usize, message: String },

    #[error("Unsupported vertex type: {0}")]
    UnsupportedVertexType(String),

    #[error("Unsupported edge type: {0}")]
    UnsupportedEdgeType(String),

    #[error("Invalid number format at line {line}: {value}")]
    InvalidNumber { line: usize, value: String },

    #[error("Missing required fields at line {line}")]
    MissingFields { line: usize },

    #[error("Duplicate vertex ID: {id}")]
    DuplicateVertex { id: usize },

    #[error("Invalid quaternion at line {line}: norm = {norm:.6}, expected ~1.0")]
    InvalidQuaternion { line: usize, norm: f64 },

    #[error("Unsupported file format: {0}")]
    UnsupportedFormat(String),

    #[error("Failed to create file '{path}': {reason}")]
    FileCreationFailed { path: String, reason: String },
}

impl IoError {
    /// Log the error using tracing::error and return self for chaining
    pub fn log(self) -> Self {
        error!("{}", self);
        self
    }

    /// Log the error with source error information using tracing::error and return self for chaining
    pub fn log_with_source<E: std::fmt::Debug>(self, source_error: E) -> Self {
        error!("{} | Source: {:?}", self, source_error);
        self
    }
}

#[derive(Clone, PartialEq)]
pub struct VertexSE2 {
    pub id: usize,
    pub pose: SE2,
}
impl VertexSE2 {
    pub fn new(id: usize, x: f64, y: f64, theta: f64) -> Self {
        Self {
            id,
            pose: SE2::from_xy_angle(x, y, theta),
        }
    }

    pub fn from_vector(id: usize, vector: Vector3<f64>) -> Self {
        Self {
            id,
            pose: SE2::from_xy_angle(vector[0], vector[1], vector[2]),
        }
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn x(&self) -> f64 {
        self.pose.x()
    }

    pub fn y(&self) -> f64 {
        self.pose.y()
    }

    pub fn theta(&self) -> f64 {
        self.pose.angle()
    }
}

impl Display for VertexSE2 {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "VertexSE2 [ id: {}, pose: {} ]", self.id, self.pose)
    }
}

impl VertexSE2 {
    /// Convert to Rerun 2D position with scaling
    ///
    /// **Note:** Requires the `visualization` feature to be enabled.
    ///
    /// # Arguments
    /// * `scale` - Scale factor to apply to position
    ///
    /// # Returns
    /// 2D position array [x, y] compatible with Rerun Points2D
    pub fn to_rerun_position_2d(&self, scale: f32) -> [f32; 2] {
        [(self.x() as f32) * scale, (self.y() as f32) * scale]
    }

    /// Convert to Rerun 3D position with scaling and specified height
    ///
    /// **Note:** Requires the `visualization` feature to be enabled.
    ///
    /// # Arguments
    /// * `scale` - Scale factor to apply to X and Y
    /// * `height` - Z coordinate for the 2D point in 3D space
    ///
    /// # Returns
    /// 3D position compatible with Rerun Transform3D or Points3D
    #[cfg(feature = "visualization")]
    pub fn to_rerun_position_3d(&self, scale: f32, height: f32) -> Vec3 {
        Vec3::new((self.x() as f32) * scale, (self.y() as f32) * scale, height)
    }
}

/// SE3 vertex with ID (x, y, z, qx, qy, qz, qw)
#[derive(Clone, PartialEq)]
pub struct VertexSE3 {
    pub id: usize,
    pub pose: SE3,
}

impl VertexSE3 {
    pub fn new(id: usize, translation: Vector3<f64>, rotation: UnitQuaternion<f64>) -> Self {
        Self {
            id,
            pose: SE3::new(translation, rotation),
        }
    }

    pub fn from_vector(id: usize, vector: [f64; 7]) -> Self {
        let translation = Vector3::from([vector[0], vector[1], vector[2]]);
        let rotation = UnitQuaternion::from_quaternion(Quaternion::from([
            vector[3], vector[4], vector[5], vector[6],
        ]));
        Self::new(id, translation, rotation)
    }

    pub fn from_translation_quaternion(
        id: usize,
        translation: Vector3<f64>,
        quaternion: Quaternion<f64>,
    ) -> Self {
        Self {
            id,
            pose: SE3::from_translation_quaternion(translation, quaternion),
        }
    }

    pub fn id(&self) -> usize {
        self.id
    }

    pub fn translation(&self) -> Vector3<f64> {
        self.pose.translation()
    }

    pub fn rotation(&self) -> UnitQuaternion<f64> {
        self.pose.rotation_quaternion()
    }

    pub fn x(&self) -> f64 {
        self.pose.x()
    }

    pub fn y(&self) -> f64 {
        self.pose.y()
    }

    pub fn z(&self) -> f64 {
        self.pose.z()
    }
}

impl Display for VertexSE3 {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(f, "VertexSE3 [ id: {}, pose: {} ]", self.id, self.pose)
    }
}

impl VertexSE3 {
    /// Convert to Rerun 3D transform components (position and rotation) with scaling
    ///
    /// **Note:** Requires the `visualization` feature to be enabled.
    ///
    /// # Arguments
    /// * `scale` - Scale factor to apply to position
    ///
    /// # Returns
    /// Tuple of (position, rotation) compatible with Rerun Transform3D
    #[cfg(feature = "visualization")]
    pub fn to_rerun_transform(&self, scale: f32) -> (Vec3, Quat) {
        // Extract translation and convert to glam Vec3
        let trans = self.translation();
        let position = Vec3::new(trans.x as f32, trans.y as f32, trans.z as f32) * scale;

        // Extract rotation quaternion and convert to glam Quat
        let rot = self.rotation();
        let nq = rot.as_ref();
        let rotation = Quat::from_xyzw(nq.i as f32, nq.j as f32, nq.k as f32, nq.w as f32);

        (position, rotation)
    }
}

/// 2D edge constraint between two SE2 vertices
#[derive(Clone, PartialEq)]
pub struct EdgeSE2 {
    pub from: usize,
    pub to: usize,
    pub measurement: SE2,          // Relative transformation
    pub information: Matrix3<f64>, // 3x3 information matrix
}

impl EdgeSE2 {
    pub fn new(
        from: usize,
        to: usize,
        dx: f64,
        dy: f64,
        dtheta: f64,
        information: Matrix3<f64>,
    ) -> Self {
        Self {
            from,
            to,
            measurement: SE2::from_xy_angle(dx, dy, dtheta),
            information,
        }
    }
}

impl Display for EdgeSE2 {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "EdgeSE2 [ from: {}, to: {}, measurement: {}, information: {} ]",
            self.from, self.to, self.measurement, self.information
        )
    }
}

/// 3D edge constraint between two SE3 vertices
#[derive(Clone, PartialEq)]
pub struct EdgeSE3 {
    pub from: usize,
    pub to: usize,
    pub measurement: SE3,          // Relative transformation
    pub information: Matrix6<f64>, // 6x6 information matrix
}

impl EdgeSE3 {
    pub fn new(
        from: usize,
        to: usize,
        translation: Vector3<f64>,
        rotation: UnitQuaternion<f64>,
        information: Matrix6<f64>,
    ) -> Self {
        Self {
            from,
            to,
            measurement: SE3::new(translation, rotation),
            information,
        }
    }
}

impl Display for EdgeSE3 {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "EdgeSE3 [ from: {}, to: {}, measurement: {}, information: {} ]",
            self.from, self.to, self.measurement, self.information
        )
    }
}

/// Main graph structure containing vertices and edges
#[derive(Clone)]
pub struct Graph {
    pub vertices_se2: collections::HashMap<usize, VertexSE2>,
    pub vertices_se3: collections::HashMap<usize, VertexSE3>,
    pub edges_se2: Vec<EdgeSE2>,
    pub edges_se3: Vec<EdgeSE3>,
}

impl Graph {
    pub fn new() -> Self {
        Self {
            vertices_se2: collections::HashMap::new(),
            vertices_se3: collections::HashMap::new(),
            edges_se2: Vec::new(),
            edges_se3: Vec::new(),
        }
    }

    pub fn vertex_count(&self) -> usize {
        self.vertices_se2.len() + self.vertices_se3.len()
    }

    pub fn edge_count(&self) -> usize {
        self.edges_se2.len() + self.edges_se3.len()
    }

    // Note: The from_optimized_variables() method has been moved to apex-solver
    // as it depends on VariableEnum from the core module.
}

impl Default for Graph {
    fn default() -> Self {
        Self::new()
    }
}

impl Display for Graph {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "Graph [[ vertices_se2: {} (count: {}), vertices_se3: {} (count: {}), edges_se2: {} (count: {}), edges_se3: {} (count: {}) ]]",
            self.vertices_se2
                .values()
                .map(|v| format!("{}", v))
                .collect::<Vec<_>>()
                .join(", "),
            self.vertices_se2.len(),
            self.vertices_se3
                .values()
                .map(|v| format!("{}", v))
                .collect::<Vec<_>>()
                .join(", "),
            self.vertices_se3.len(),
            self.edges_se2
                .iter()
                .map(|e| format!("{}", e))
                .collect::<Vec<_>>()
                .join(", "),
            self.edges_se2.len(),
            self.edges_se3
                .iter()
                .map(|e| format!("{}", e))
                .collect::<Vec<_>>()
                .join(", "),
            self.edges_se3.len()
        )
    }
}

/// Trait for graph file loaders and writers
pub trait GraphLoader {
    /// Load a graph from a file
    fn load<P: AsRef<Path>>(path: P) -> Result<Graph, IoError>;

    /// Write a graph to a file
    fn write<P: AsRef<Path>>(graph: &Graph, path: P) -> Result<(), IoError>;
}

/// Convenience function to load any supported format based on file extension
pub fn load_graph<P: AsRef<Path>>(path: P) -> Result<Graph, IoError> {
    let path_ref = path.as_ref();
    let extension = path_ref
        .extension()
        .and_then(|ext| ext.to_str())
        .ok_or_else(|| {
            IoError::UnsupportedFormat("No file extension".to_string())
                .log_with_source(format!("File path: {:?}", path_ref))
        })?;

    match extension.to_lowercase().as_str() {
        "g2o" => G2oLoader::load(path),
        "graph" => ToroLoader::load(path),
        _ => Err(
            IoError::UnsupportedFormat(format!("Unsupported extension: {extension}"))
                .log_with_source(format!("File path: {:?}", path_ref)),
        ),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use nalgebra::{Matrix3, Matrix6, Quaternion, UnitQuaternion, Vector3};
    use std::{error, io::Write};
    use tempfile::NamedTempFile;

    #[test]
    fn test_load_simple_graph() -> Result<(), IoError> {
        let mut temp_file = NamedTempFile::new().map_err(|e| {
            IoError::FileCreationFailed {
                path: "temp_file".to_string(),
                reason: e.to_string(),
            }
            .log()
        })?;
        writeln!(temp_file, "VERTEX_SE2 0 0.0 0.0 0.0")?;
        writeln!(temp_file, "VERTEX_SE2 1 1.0 0.0 0.0")?;
        writeln!(temp_file, "# This is a comment")?;
        writeln!(temp_file)?; // Empty line
        writeln!(temp_file, "VERTEX_SE3:QUAT 2 0.0 0.0 0.0 0.0 0.0 0.0 1.0")?;

        let graph = G2oLoader::load(temp_file.path())?;

        assert_eq!(graph.vertices_se2.len(), 2);
        assert_eq!(graph.vertices_se3.len(), 1);
        assert!(graph.vertices_se2.contains_key(&0));
        assert!(graph.vertices_se2.contains_key(&1));
        assert!(graph.vertices_se3.contains_key(&2));

        Ok(())
    }

    #[test]
    fn test_load_m3500() -> Result<(), Box<dyn error::Error>> {
        let path = utils::ensure_odometry_dataset("M3500")?;
        let graph = G2oLoader::load(&path)?;
        assert!(!graph.vertices_se2.is_empty());
        Ok(())
    }

    #[test]
    fn test_load_sphere2500() -> Result<(), Box<dyn error::Error>> {
        let path = utils::ensure_odometry_dataset("sphere2500")?;
        let graph = G2oLoader::load(&path)?;
        assert!(!graph.vertices_se3.is_empty());
        Ok(())
    }

    #[test]
    fn test_duplicate_vertex_error() -> Result<(), io::Error> {
        let mut temp_file = NamedTempFile::new()?;
        writeln!(temp_file, "VERTEX_SE2 0 0.0 0.0 0.0")?;
        writeln!(temp_file, "VERTEX_SE2 0 1.0 0.0 0.0")?; // Duplicate ID

        let result = G2oLoader::load(temp_file.path());
        assert!(matches!(result, Err(IoError::DuplicateVertex { id: 0 })));

        Ok(())
    }

    #[test]
    fn test_toro_loader() -> Result<(), IoError> {
        let mut temp_file = NamedTempFile::new().map_err(|e| {
            IoError::FileCreationFailed {
                path: "temp_file".to_string(),
                reason: e.to_string(),
            }
            .log()
        })?;
        writeln!(temp_file, "VERTEX2 0 0.0 0.0 0.0")?;
        writeln!(temp_file, "VERTEX2 1 1.0 0.0 0.0")?;

        let graph = ToroLoader::load(temp_file.path()).map_err(|e| e.log())?;
        assert_eq!(graph.vertices_se2.len(), 2);

        Ok(())
    }

    #[test]
    #[cfg(feature = "visualization")]
    fn test_se3_to_rerun() {
        let vertex = VertexSE3::new(0, Vector3::new(1.0, 2.0, 3.0), UnitQuaternion::identity());

        let (pos, rot) = vertex.to_rerun_transform(0.1);

        assert!((pos.x - 0.1).abs() < 1e-6);
        assert!((pos.y - 0.2).abs() < 1e-6);
        assert!((pos.z - 0.3).abs() < 1e-6);
        assert!((rot.w - 1.0).abs() < 1e-6);
    }

    #[test]
    fn test_se2_to_rerun_2d() {
        let vertex = VertexSE2::new(0, 10.0, 20.0, 0.5);

        let pos = vertex.to_rerun_position_2d(0.1);

        assert!((pos[0] - 1.0).abs() < 1e-6);
        assert!((pos[1] - 2.0).abs() < 1e-6);
    }

    #[test]
    #[cfg(feature = "visualization")]
    fn test_se2_to_rerun_3d() {
        let vertex = VertexSE2::new(0, 10.0, 20.0, 0.5);

        let pos = vertex.to_rerun_position_3d(0.1, 5.0);

        assert!((pos.x - 1.0).abs() < 1e-6);
        assert!((pos.y - 2.0).abs() < 1e-6);
        assert!((pos.z - 5.0).abs() < 1e-6);
    }

    // -------------------------------------------------------------------------
    // VertexSE2 / VertexSE3 constructors
    // -------------------------------------------------------------------------

    #[test]
    fn test_vertex_se2_from_vector() {
        let v = VertexSE2::from_vector(5, Vector3::new(1.0, 2.0, 0.5));
        assert_eq!(v.id(), 5);
        assert!((v.x() - 1.0).abs() < 1e-12);
        assert!((v.y() - 2.0).abs() < 1e-12);
        assert!((v.theta() - 0.5).abs() < 1e-12);
    }

    #[test]
    fn test_vertex_se3_from_vector() {
        // [tx, ty, tz, qx, qy, qz, qw]
        let arr = [1.0f64, 2.0, 3.0, 0.0, 0.0, 0.0, 1.0];
        let v = VertexSE3::from_vector(7, arr);
        assert_eq!(v.id(), 7);
        assert!((v.x() - 1.0).abs() < 1e-10);
        assert!((v.y() - 2.0).abs() < 1e-10);
        assert!((v.z() - 3.0).abs() < 1e-10);
    }

    #[test]
    fn test_vertex_se3_from_translation_quaternion() {
        let trans = Vector3::new(1.0, 2.0, 3.0);
        // nalgebra Quaternion::new(w, i, j, k)
        let quat = Quaternion::new(1.0, 0.0, 0.0, 0.0);
        let v = VertexSE3::from_translation_quaternion(3, trans, quat);
        assert_eq!(v.id(), 3);
        assert!((v.x() - 1.0).abs() < 1e-10);
        assert!((v.y() - 2.0).abs() < 1e-10);
        assert!((v.z() - 3.0).abs() < 1e-10);
    }

    // -------------------------------------------------------------------------
    // Display implementations
    // -------------------------------------------------------------------------

    #[test]
    fn test_vertex_se2_display() {
        let v = VertexSE2::new(0, 1.0, 2.0, 0.5);
        let s = format!("{v}");
        assert!(
            s.contains("VertexSE2"),
            "Display should contain 'VertexSE2': {s}"
        );
        assert!(s.contains('0'), "Display should contain id: {s}");
    }

    #[test]
    fn test_vertex_se3_display() {
        let v = VertexSE3::new(1, Vector3::new(1.0, 2.0, 3.0), UnitQuaternion::identity());
        let s = format!("{v}");
        assert!(
            s.contains("VertexSE3"),
            "Display should contain 'VertexSE3': {s}"
        );
        assert!(s.contains('1'), "Display should contain id: {s}");
    }

    #[test]
    fn test_edge_se2_display() {
        let e = EdgeSE2::new(0, 1, 1.0, 0.0, 0.0, Matrix3::identity());
        let s = format!("{e}");
        assert!(
            s.contains("EdgeSE2"),
            "Display should contain 'EdgeSE2': {s}"
        );
    }

    #[test]
    fn test_edge_se3_display() {
        let e = EdgeSE3::new(
            0,
            1,
            Vector3::zeros(),
            UnitQuaternion::identity(),
            Matrix6::identity(),
        );
        let s = format!("{e}");
        assert!(
            s.contains("EdgeSE3"),
            "Display should contain 'EdgeSE3': {s}"
        );
    }

    #[test]
    fn test_graph_display() {
        let mut g = Graph::new();
        g.vertices_se2.insert(0, VertexSE2::new(0, 0.0, 0.0, 0.0));
        let s = format!("{g}");
        assert!(s.contains("Graph"), "Display should contain 'Graph': {s}");
        assert!(s.contains("count: 1"), "Display should show count: {s}");
    }

    // -------------------------------------------------------------------------
    // Graph helpers
    // -------------------------------------------------------------------------

    #[test]
    fn test_graph_default_is_empty() {
        let g = Graph::default();
        assert_eq!(g.vertex_count(), 0);
        assert_eq!(g.edge_count(), 0);
    }

    #[test]
    fn test_graph_vertex_and_edge_counts() {
        let mut g = Graph::new();
        g.vertices_se2.insert(0, VertexSE2::new(0, 0.0, 0.0, 0.0));
        g.vertices_se3.insert(
            1,
            VertexSE3::new(1, Vector3::zeros(), UnitQuaternion::identity()),
        );
        g.edges_se2
            .push(EdgeSE2::new(0, 1, 0.0, 0.0, 0.0, Matrix3::identity()));
        assert_eq!(g.vertex_count(), 2);
        assert_eq!(g.edge_count(), 1);
    }

    // -------------------------------------------------------------------------
    // load_graph() dispatch
    // -------------------------------------------------------------------------

    #[test]
    fn test_load_graph_unsupported_extension() {
        let result = load_graph("fake_path.xyz");
        assert!(
            matches!(result, Err(IoError::UnsupportedFormat(_))),
            "unknown extension should return UnsupportedFormat"
        );
    }

    #[test]
    fn test_load_graph_no_extension() {
        let result = load_graph("/tmp/no_extension_file");
        assert!(
            matches!(result, Err(IoError::UnsupportedFormat(_))),
            "path with no extension should return UnsupportedFormat"
        );
    }

    #[test]
    fn test_load_graph_toro_extension() -> Result<(), Box<dyn error::Error>> {
        let mut f = NamedTempFile::new()?;
        writeln!(f, "VERTEX2 0 0.0 0.0 0.0")?;
        writeln!(f, "VERTEX2 1 1.0 0.0 0.0")?;
        f.flush()?;
        // Rename temp file path to have .graph extension
        let toro_path = f.path().with_extension("graph");
        std::fs::copy(f.path(), &toro_path)?;
        let graph = load_graph(&toro_path)?;
        std::fs::remove_file(&toro_path)?;
        assert_eq!(graph.vertices_se2.len(), 2);
        Ok(())
    }

    #[test]
    fn test_io_error_log_returns_self() {
        let err = IoError::UnsupportedFormat("xyz".to_string());
        let returned = err.log();
        assert!(matches!(returned, IoError::UnsupportedFormat(_)));
    }

    #[test]
    fn test_io_error_log_with_source() {
        let err = IoError::UnsupportedFormat("abc".to_string());
        let source = std::io::Error::other("source");
        let returned = err.log_with_source(source);
        assert!(matches!(returned, IoError::UnsupportedFormat(_)));
    }

    #[test]
    fn test_vertex_se2_theta() {
        use std::f64::consts::PI;
        let v = VertexSE2::new(0, 1.0, 2.0, PI / 4.0);
        assert!((v.theta() - PI / 4.0).abs() < 1e-10);
    }

    #[test]
    fn test_edge_se3_new() {
        let t = Vector3::new(1.0, 2.0, 3.0);
        let r = UnitQuaternion::identity();
        let info = Matrix6::identity();
        let e = EdgeSE3::new(0, 1, t, r, info);
        assert_eq!(e.from, 0);
        assert_eq!(e.to, 1);
    }

    #[test]
    fn test_vertex_se3_new() {
        let t = Vector3::new(1.0, 2.0, 3.0);
        let r = UnitQuaternion::identity();
        let v = VertexSE3::new(5, t, r);
        assert_eq!(v.id, 5);
        assert!((v.translation() - t).norm() < 1e-10);
    }
}