pub mod probe;

use clap::Parser;
use nalgebra::{Matrix3, Quaternion, SVD, Translation3, UnitQuaternion};
use plotters::prelude::*;
use plotters::style::full_palette::GREY;
use serde::Serialize;
use sophus_lie::IsAffineGroup;
use std::error::Error;
use std::fs::File;
use std::io::{BufRead, BufReader, Write};
use std::path::{Path, PathBuf};
use std::time::{SystemTime, UNIX_EPOCH};

#[cfg(feature = "rerun")]
use rerun as rr;
#[cfg(feature = "rerun")]
use rerun::external::glam;

pub type SE3 = sophus_lie::Isometry3F64;
type SO3 = sophus_lie::Rotation3F64;

const PLOT_WIDTH: u32 = 800;
const PLOT_HEIGHT: u32 = 600;

#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    #[command(subcommand)]
    pub command: Command,
}

#[derive(clap::Subcommand, Debug)]
enum Command {
    /// Evaluate trajectories
    Trajectory {
        /// Path to estimated trajectory in TUM format
        estimated_file: PathBuf,

        /// Path to reference trajectory in TUM format
        reference_file: PathBuf,

        /// Print metrics as JSON to stdout
        #[arg(long)]
        json: bool,

        /// Export summary metrics JSON to a file. Without value, auto-generates a JSON filename.
        #[arg(
            long = "export-summary-json",
            alias = "export-results-json",
            alias = "export-results",
            value_name = "JSON_FILE",
            num_args = 0..=1,
            default_missing_value = ""
        )]
        export_results: Option<String>,

        /// Export plots (SVG) to a directory. Without value, auto-generates a directory.
        #[arg(
            long = "export-plots-dir",
            alias = "export-plots",
            value_name = "DIR",
            num_args = 0..=1,
            default_missing_value = ""
        )]
        export_plots: Option<String>,

        /// Frame delta for Relative Pose Error (RPE)
        #[arg(short, long, default_value_t = 10)]
        delta: usize,

        /// Visualize with rerun
        #[cfg(feature = "rerun")]
        #[arg(long)]
        rerun: bool,
    },
    /// Generate shell completions
    Completions {
        /// The shell to generate completions for
        #[arg(value_enum)]
        shell: clap_complete::Shell,
    },
    /// Probe system and process resources
    Probe {
        /// Target process ID or name
        target: String,

        /// Output directory for probe CSV, summary JSON, and plots
        #[arg(long, default_value = "probe_eval")]
        output_dir: PathBuf,

        /// Sampling interval in seconds
        #[arg(short, long, default_value_t = 1.0)]
        interval: f64,

        /// Print probe metrics as JSON to stdout
        #[arg(long)]
        json: bool,

        /// Visualize with rerun
        #[cfg(feature = "rerun")]
        #[arg(long)]
        rerun: bool,

        /// Strict matching: wait until a process with the exact name appears (no fuzzy match)
        #[arg(short, long)]
        strict: bool,
    },
}

fn epoch_seconds() -> u64 {
    SystemTime::now()
        .duration_since(UNIX_EPOCH)
        .map_or(0, |d| d.as_secs())
}

fn resolve_optional_path(value: Option<String>, default_path: PathBuf) -> Option<PathBuf> {
    match value {
        None => None,
        Some(s) if s.trim().is_empty() => Some(default_path),
        Some(s) => Some(PathBuf::from(s)),
    }
}

#[derive(Serialize)]
struct EvalResults {
    umeyama_scale: f64,
    umeyama_translation: [f64; 3],
    ape_trans_rmse: f64,
    ape_rot_rmse_deg: f64,
    rpe_trans_rmse: f64,
    rpe_rot_rmse_deg: f64,
}

#[derive(Debug, Clone, Copy)]
struct Pose {
    timestamp: f64,
    translation: Translation3<f64>,
    rotation: UnitQuaternion<f64>,
}

#[cfg(feature = "rerun")]
fn pose_to_rerun_transform(pose: &Pose) -> rr::Transform3D {
    let correction = nalgebra::UnitQuaternion::from_axis_angle(
        &nalgebra::Vector3::x_axis(),
        -std::f32::consts::FRAC_PI_2,
    );

    let corrected_rot = pose.rotation.cast::<f32>() * correction;

    rr::Transform3D::from_translation_rotation(
        glam::Vec3::new(
            pose.translation.x as f32,
            pose.translation.y as f32,
            pose.translation.z as f32,
        ),
        rr::datatypes::Quaternion::from_xyzw([
            corrected_rot.i,
            corrected_rot.j,
            corrected_rot.k,
            corrected_rot.w,
        ]),
    )
}

#[cfg(feature = "rerun")]
fn get_rerun_color_for_error(error: f64, max_error: f64) -> rr::Color {
    let normalized_error = if max_error > 0.0 {
        (error / max_error).clamp(0.0, 1.0) as f32
    } else {
        0.0
    };
    rr::Color::from_rgb(
        (normalized_error * 255.0) as u8,
        0,
        ((1.0 - normalized_error) * 255.0) as u8,
    )
}

#[cfg(feature = "rerun")]
fn visualize_with_rerun(
    corresponding_gt_trajectory: &[Pose],
    aligned_est_trajectory: &[Pose],
    ape_data_for_plot: &[(f64, f64)],
    rpe_data_for_plot: &[(f64, f64)],
) -> anyhow::Result<()> {
    let rec = rr::RecordingStreamBuilder::new("Trajectory Analysis").spawn()?;

    rec.log("world", &rr::ViewCoordinates::RDF())?;
    rec.log_static(
        "world/axes",
        &rr::Arrows3D::from_vectors([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]])
            .with_colors([[255, 0, 0], [0, 255, 0], [0, 0, 255]])
            .with_labels(["X", "Y", "Z"]),
    )?;

    let gt_points: Vec<glam::Vec3> = corresponding_gt_trajectory
        .iter()
        .map(|p| {
            glam::vec3(
                p.translation.x as f32,
                p.translation.y as f32,
                p.translation.z as f32,
            )
        })
        .collect();

    rec.log_static(
        "world/reference_path",
        &rr::LineStrips3D::new([gt_points])
            .with_colors([rr::Color::from_rgb(128, 128, 128)])
            .with_labels(["Reference"])
            .with_radii([rr::Radius::new_ui_points(0.85)]),
    )?;

    let max_ape_for_color = ape_data_for_plot
        .iter()
        .map(|(_, e)| *e)
        .fold(0.0, f64::max);

    let est_points: Vec<glam::Vec3> = aligned_est_trajectory
        .iter()
        .map(|p| {
            glam::vec3(
                p.translation.x as f32,
                p.translation.y as f32,
                p.translation.z as f32,
            )
        })
        .collect();
    let ape_colors: Vec<rr::Color> = ape_data_for_plot
        .iter()
        .map(|(_, err)| get_rerun_color_for_error(*err, max_ape_for_color))
        .collect();

    log_colored_trajectory(
        &rec,
        "world/estimated_path_ape",
        &est_points,
        &ape_colors,
        "Estimated (APE)",
        0.5,
    )?;

    let max_rpe_for_color = rpe_data_for_plot
        .iter()
        .map(|(_, e)| *e)
        .fold(0.0, f64::max);
    let rpe_points: Vec<glam::Vec3> = est_points[0..rpe_data_for_plot.len()].to_vec();
    let rpe_colors: Vec<rr::Color> = rpe_data_for_plot
        .iter()
        .map(|(_, err)| get_rerun_color_for_error(*err, max_rpe_for_color))
        .collect();

    log_colored_trajectory(
        &rec,
        "world/estimated_path_rpe",
        &rpe_points,
        &rpe_colors,
        "Estimated (RPE)",
        0.5,
    )?;

    for i in 0..corresponding_gt_trajectory.len() {
        let gt_pose = &corresponding_gt_trajectory[i];
        let est_pose = &aligned_est_trajectory[i];
        let (timestamp, ape) = ape_data_for_plot[i];

        rec.set_duration_secs("timestamp", timestamp);

        rec.log("world/reference_pose", &pose_to_rerun_transform(gt_pose))?;
        rec.log("world/estimated_pose", &pose_to_rerun_transform(est_pose))?;

        rec.log(
            "world/estimated_pose/camera",
            &rr::Pinhole::from_focal_length_and_resolution([300.0, 300.0], [640.0, 480.0]),
        )?;

        let gt_point = glam::vec3(
            gt_pose.translation.x as f32,
            gt_pose.translation.y as f32,
            gt_pose.translation.z as f32,
        );
        let est_point = glam::vec3(
            est_pose.translation.x as f32,
            est_pose.translation.y as f32,
            est_pose.translation.z as f32,
        );
        rec.log(
            "world/error_line",
            &rr::LineStrips3D::new([[gt_point, est_point]])
                .with_colors([rr::Color::from_rgb(255, 128, 128)]),
        )?;

        rec.log("errors/ape_translation", &rr::Scalars::new([ape]))?;
    }

    for (timestamp, rpe) in rpe_data_for_plot {
        rec.set_duration_secs("timestamp", *timestamp);
        rec.log("errors/rpe_translation", &rr::Scalars::new([*rpe]))?;
    }

    Ok(())
}

#[cfg(feature = "rerun")]
fn log_colored_trajectory(
    rec: &rr::RecordingStream,
    entity_path: &str,
    points: &[glam::Vec3],
    colors: &[rr::Color],
    label: &str,
    radius: f32,
) -> anyhow::Result<()> {
    if points.len() < 2 {
        return Ok(());
    }

    let strips = points.windows(2).map(|segment| segment.to_vec());

    let segment_colors: Vec<rr::Color> = colors[0..points.len() - 1].to_vec();

    rec.log_static(
        entity_path,
        &rr::LineStrips3D::new(strips)
            .with_colors(segment_colors)
            .with_labels([label])
            .with_radii([rr::Radius::new_ui_points(radius)]),
    )?;

    Ok(())
}

fn umeyama_alignment(
    gt_traj: &[Pose],
    est_traj: &[Pose],
) -> anyhow::Result<(f64, Matrix3<f64>, nalgebra::Vector3<f64>)> {
    if gt_traj.is_empty() || est_traj.is_empty() || gt_traj.len() != est_traj.len() {
        anyhow::bail!("Umeyama alignment requires equal-length non-empty trajectories");
    }
    let gt_points: Vec<_> = gt_traj.iter().map(|p| p.translation.vector).collect();
    let est_points: Vec<_> = est_traj.iter().map(|p| p.translation.vector).collect();

    let gt_centroid: nalgebra::Vector3<f64> =
        gt_points.iter().sum::<nalgebra::Vector3<f64>>() / gt_points.len() as f64;
    let est_centroid: nalgebra::Vector3<f64> =
        est_points.iter().sum::<nalgebra::Vector3<f64>>() / est_points.len() as f64;

    let gt_centered: Vec<_> = gt_points.iter().map(|p| p - gt_centroid).collect();
    let est_centered: Vec<_> = est_points.iter().map(|p| p - est_centroid).collect();

    let mut h = Matrix3::<f64>::zeros();
    for i in 0..gt_points.len() {
        h += gt_centered[i] * est_centered[i].transpose();
    }
    h /= gt_points.len() as f64;

    let svd = SVD::new(h, true, true);
    let u = svd
        .u
        .ok_or_else(|| anyhow::anyhow!("SVD did not return U matrix"))?;
    let vt = svd
        .v_t
        .ok_or_else(|| anyhow::anyhow!("SVD did not return V^T matrix"))?;
    let s_diag = svd.singular_values;

    let mut s_identity = Matrix3::<f64>::identity();
    if (u * vt).determinant() < 0.0 {
        s_identity[(2, 2)] = -1.0;
    }
    let r = u * s_identity * vt;

    let var_est =
        est_centered.iter().map(|p| p.norm_squared()).sum::<f64>() / est_points.len() as f64;
    let s = (s_diag.component_mul(&s_identity.diagonal())).sum() / var_est;

    let t = gt_centroid - s * r * est_centroid;

    Ok((s, r, t))
}

fn create_error_plot(
    filepath: &Path,
    title: &str,
    data: &[(f64, f64)],
) -> Result<(), Box<dyn Error>> {
    if data.is_empty() {
        return Ok(());
    }
    let root = SVGBackend::new(filepath, (PLOT_WIDTH, PLOT_HEIGHT)).into_drawing_area();
    root.fill(&WHITE)?;

    let start_t = data[0].0;

    let max_t = data.iter().last().map_or(0.0, |(t, _)| *t);
    let max_err = data.iter().fold(0.0, |max: f64, (_, e)| max.max(*e));

    let mut chart = ChartBuilder::on(&root)
        .caption(title, ("sans-serif", 40).into_font())
        .margin(10)
        .x_label_area_size(40)
        .y_label_area_size(60)
        .build_cartesian_2d(0.0..(max_t - start_t), 0.0..max_err * 1.1)?;

    chart
        .configure_mesh()
        .x_desc("Time (s)")
        .y_desc("Error")
        .draw()?;

    chart.draw_series(LineSeries::new(
        data.iter().map(|(t, e)| (*t - start_t, *e)),
        &RED,
    ))?;

    root.present()?;
    Ok(())
}

fn normalize_tum_timestamp(timestamp: f64) -> f64 {
    if timestamp.abs() >= 1.0e12 {
        timestamp / 1.0e9
    } else {
        timestamp
    }
}

fn parse_pose(line: &str) -> Result<Pose, Box<dyn Error>> {
    let parts: Vec<f64> = line
        .split_whitespace()
        .map(|s| s.parse::<f64>())
        .collect::<Result<Vec<f64>, _>>()?;

    if parts.len() != 8 {
        return Err("Each line must have 8 values".into());
    }

    let translation = Translation3::new(parts[1], parts[2], parts[3]);

    Ok(Pose {
        timestamp: normalize_tum_timestamp(parts[0]),
        translation,
        rotation: UnitQuaternion::from_quaternion(Quaternion::new(
            parts[7], parts[4], parts[5], parts[6],
        )),
    })
}

fn read_trajectory(path: &Path) -> Result<Vec<Pose>, Box<dyn Error>> {
    let file = File::open(path)?;
    let reader = BufReader::new(file);
    reader.lines().map(|line| parse_pose(&line?)).collect()
}

fn pose_to_se3(pose: &Pose) -> anyhow::Result<SE3> {
    let so3 = SO3::try_from_mat(pose.rotation.to_rotation_matrix().matrix())
        .ok_or_else(|| anyhow::anyhow!("Invalid rotation matrix in pose"))?;
    Ok(SE3::from_rotation_and_translation(
        so3,
        pose.translation.vector,
    ))
}

fn normalize_trajectory_to_first(trajectory: &[Pose]) -> anyhow::Result<Vec<Pose>> {
    let first_pose = trajectory
        .first()
        .ok_or_else(|| anyhow::anyhow!("Cannot normalize an empty trajectory"))?;
    let first_pose_inv = pose_to_se3(first_pose)?.inverse();

    trajectory
        .iter()
        .map(|pose| {
            let se3_pose = pose_to_se3(pose)?;
            let normalized_se3 = first_pose_inv.group_mul(se3_pose);
            Ok(Pose {
                timestamp: pose.timestamp,
                translation: Translation3::from(normalized_se3.translation()),
                rotation: UnitQuaternion::from_matrix(&normalized_se3.rotation().matrix()),
            })
        })
        .collect()
}

fn get_plotters_color_for_error(error: f64, max_error: f64) -> RGBColor {
    let normalized_error = if max_error > 0.0 {
        (error / max_error).clamp(0.0, 1.0)
    } else {
        0.0
    };
    let r = (normalized_error * 255.0) as u8;
    let b = ((1.0 - normalized_error) * 255.0) as u8;
    RGBColor(r, 0, b)
}

fn interpolate_pose(p1: &Pose, p2: &Pose, timestamp: f64) -> anyhow::Result<Pose> {
    if (p2.timestamp - p1.timestamp).abs() < f64::EPSILON {
        anyhow::bail!("Cannot interpolate two poses with identical timestamps");
    }
    let t = (timestamp - p1.timestamp) / (p2.timestamp - p1.timestamp);

    let p1 = SE3::from_rotation_and_translation(
        SO3::try_from_mat(p1.rotation.to_rotation_matrix().matrix())
            .ok_or_else(|| anyhow::anyhow!("Invalid rotation matrix in first pose"))?,
        p1.translation.vector,
    );

    let p2 = SE3::from_rotation_and_translation(
        SO3::try_from_mat(p2.rotation.to_rotation_matrix().matrix())
            .ok_or_else(|| anyhow::anyhow!("Invalid rotation matrix in second pose"))?,
        p2.translation.vector,
    );

    let diff = p1.interpolate(&p2, t);

    Ok(Pose {
        timestamp,
        translation: Translation3::from(diff.translation()),
        rotation: UnitQuaternion::from_matrix(&diff.rotation().matrix()),
    })
}

fn create_2d_plots(
    out_dir: &Path,
    gt_traj: &[Pose],
    est_traj: &[Pose],
    ape_data: &[(f64, f64)],
    rpe_data: &[(f64, f64)],
) -> Result<(), Box<dyn std::error::Error>> {
    let max_ape = ape_data.iter().map(|(_, e)| *e).fold(0.0, f64::max);

    let path_ape_side = out_dir.join("side_view_APE.svg");
    let root_ape_side =
        SVGBackend::new(&path_ape_side, (PLOT_WIDTH, PLOT_HEIGHT)).into_drawing_area();
    root_ape_side.fill(&WHITE)?;
    let (x_min, x_max, z_min, z_max) = gt_traj.iter().chain(est_traj.iter()).fold(
        (f64::MAX, f64::MIN, f64::MAX, f64::MIN),
        |(x_min, x_max, z_min, z_max), pose| {
            (
                x_min.min(pose.translation.x),
                x_max.max(pose.translation.x),
                z_min.min(pose.translation.z),
                z_max.max(pose.translation.z),
            )
        },
    );
    let mut chart_ape_side = ChartBuilder::on(&root_ape_side)
        .caption(
            "Side View Global Error (APE)",
            ("sans-serif", 50).into_font(),
        )
        .margin(10)
        .x_label_area_size(40)
        .y_label_area_size(50)
        .build_cartesian_2d(x_min..x_max, z_min..z_max)?;
    chart_ape_side.configure_mesh().draw()?;
    for i in 0..est_traj.len() - 1 {
        let p1 = &est_traj[i];
        let p2 = &est_traj[i + 1];
        let error = ape_data[i].1;
        let color = get_plotters_color_for_error(error, max_ape);
        chart_ape_side.draw_series(LineSeries::new(
            vec![
                (p1.translation.x, p1.translation.z),
                (p2.translation.x, p2.translation.z),
            ],
            color.stroke_width(2),
        ))?;
    }
    chart_ape_side
        .draw_series(LineSeries::new(vec![], RED.stroke_width(3)))?
        .label("High Error")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED.filled()));
    chart_ape_side
        .draw_series(LineSeries::new(vec![], BLUE.stroke_width(3)))?
        .label("Low Error")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE.filled()));
    chart_ape_side
        .configure_series_labels()
        .border_style(BLACK)
        .draw()?;

    let path_ape_top = out_dir.join("top_view_APE.svg");
    let root_ape_top =
        SVGBackend::new(&path_ape_top, (PLOT_WIDTH, PLOT_HEIGHT)).into_drawing_area();
    root_ape_top.fill(&WHITE)?;
    let (x_min_top, x_max_top, y_min_top, y_max_top) = gt_traj.iter().chain(est_traj.iter()).fold(
        (f64::MAX, f64::MIN, f64::MAX, f64::MIN),
        |(x_min, x_max, y_min, y_max), pose| {
            (
                x_min.min(pose.translation.x),
                x_max.max(pose.translation.x),
                y_min.min(pose.translation.y),
                y_max.max(pose.translation.y),
            )
        },
    );
    let mut chart_ape_top = ChartBuilder::on(&root_ape_top)
        .caption(
            "Top View Global Error (APE)",
            ("sans-serif", 50).into_font(),
        )
        .margin(10)
        .x_label_area_size(40)
        .y_label_area_size(50)
        .build_cartesian_2d(x_min_top..x_max_top, y_min_top..y_max_top)?;
    chart_ape_top.configure_mesh().draw()?;
    for i in 0..est_traj.len() - 1 {
        let p1 = &est_traj[i];
        let p2 = &est_traj[i + 1];
        let error = ape_data[i].1;
        let color = get_plotters_color_for_error(error, max_ape);
        chart_ape_top.draw_series(LineSeries::new(
            vec![
                (p1.translation.x, p1.translation.y),
                (p2.translation.x, p2.translation.y),
            ],
            color.stroke_width(2),
        ))?;
    }
    chart_ape_top
        .draw_series(LineSeries::new(vec![], RED.stroke_width(3)))?
        .label("High Error")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED.filled()));
    chart_ape_top
        .draw_series(LineSeries::new(vec![], BLUE.stroke_width(3)))?
        .label("Low Error")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE.filled()));
    chart_ape_top
        .configure_series_labels()
        .border_style(BLACK)
        .draw()?;

    let max_rpe = rpe_data.iter().map(|(_, e)| *e).fold(0.0, f64::max);

    let path_rpe_side = out_dir.join("side_view_drift_RPE.svg");
    let root_rpe_side =
        SVGBackend::new(&path_rpe_side, (PLOT_WIDTH, PLOT_HEIGHT)).into_drawing_area();
    root_rpe_side.fill(&WHITE)?;
    let mut chart_rpe_side = ChartBuilder::on(&root_rpe_side)
        .caption("Side View Drift (RPE)", ("sans-serif", 50).into_font())
        .margin(10)
        .x_label_area_size(40)
        .y_label_area_size(50)
        .build_cartesian_2d(x_min..x_max, z_min..z_max)?;
    chart_rpe_side.configure_mesh().draw()?;
    for i in 0..rpe_data.len() {
        let p1 = &est_traj[i];
        let p2 = &est_traj[i + 1];
        let error = rpe_data[i].1;
        let color = get_plotters_color_for_error(error, max_rpe);
        chart_rpe_side.draw_series(LineSeries::new(
            vec![
                (p1.translation.x, p1.translation.z),
                (p2.translation.x, p2.translation.z),
            ],
            color.stroke_width(2),
        ))?;
    }
    for i in rpe_data.len()..est_traj.len() - 1 {
        let p1 = &est_traj[i];
        let p2 = &est_traj[i + 1];
        chart_rpe_side.draw_series(LineSeries::new(
            vec![
                (p1.translation.x, p1.translation.z),
                (p2.translation.x, p2.translation.z),
            ],
            GREY.stroke_width(2),
        ))?;
    }
    chart_rpe_side
        .draw_series(LineSeries::new(vec![], RED.stroke_width(3)))?
        .label("High Drift")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED.filled()));
    chart_rpe_side
        .draw_series(LineSeries::new(vec![], BLUE.stroke_width(3)))?
        .label("Low Drift")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE.filled()));
    chart_rpe_side
        .draw_series(LineSeries::new(vec![], GREY.stroke_width(3)))?
        .label("No RPE Data")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], GREY.filled()));
    chart_rpe_side
        .configure_series_labels()
        .border_style(BLACK)
        .draw()?;

    let path_rpe_top = out_dir.join("top_view_drift_RPE.svg");
    let root_rpe_top =
        SVGBackend::new(&path_rpe_top, (PLOT_WIDTH, PLOT_HEIGHT)).into_drawing_area();
    root_rpe_top.fill(&WHITE)?;
    let mut chart_rpe_top = ChartBuilder::on(&root_rpe_top)
        .caption("Top View Drift (RPE)", ("sans-serif", 50).into_font())
        .margin(10)
        .x_label_area_size(40)
        .y_label_area_size(50)
        .build_cartesian_2d(x_min_top..x_max_top, y_min_top..y_max_top)?;
    chart_rpe_top.configure_mesh().draw()?;
    for i in 0..rpe_data.len() {
        let p1 = &est_traj[i];
        let p2 = &est_traj[i + 1];
        let error = rpe_data[i].1;
        let color = get_plotters_color_for_error(error, max_rpe);
        chart_rpe_top.draw_series(LineSeries::new(
            vec![
                (p1.translation.x, p1.translation.y),
                (p2.translation.x, p2.translation.y),
            ],
            color.stroke_width(2),
        ))?;
    }
    for i in rpe_data.len()..est_traj.len() - 1 {
        let p1 = &est_traj[i];
        let p2 = &est_traj[i + 1];
        chart_rpe_top.draw_series(LineSeries::new(
            vec![
                (p1.translation.x, p1.translation.y),
                (p2.translation.x, p2.translation.y),
            ],
            GREY.stroke_width(2),
        ))?;
    }
    chart_rpe_top
        .draw_series(LineSeries::new(vec![], RED.stroke_width(3)))?
        .label("High Drift")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], RED.filled()));
    chart_rpe_top
        .draw_series(LineSeries::new(vec![], BLUE.stroke_width(3)))?
        .label("Low Drift")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], BLUE.filled()));
    chart_rpe_top
        .draw_series(LineSeries::new(vec![], GREY.stroke_width(3)))?
        .label("No RPE Data")
        .legend(|(x, y)| PathElement::new(vec![(x, y), (x + 20, y)], GREY.filled()));
    chart_rpe_top
        .configure_series_labels()
        .border_style(BLACK)
        .draw()?;

    root_ape_side.present()?;
    root_ape_top.present()?;
    root_rpe_side.present()?;
    root_rpe_top.present()?;

    Ok(())
}

fn run_trajectory_command(
    estimated_path: PathBuf,
    reference_path: PathBuf,
    json: bool,
    results_path: Option<PathBuf>,
    plots_dir: Option<PathBuf>,
    delta: usize,
    rerun_output: bool,
) -> anyhow::Result<()> {
    if let Some(ref path) = results_path {
        eprintln!("Exporting metrics JSON to {}", path.display());
    }

    if let Some(ref dir) = plots_dir {
        eprintln!("Exporting plots to {}", dir.display());
        std::fs::create_dir_all(dir).map_err(|e| {
            anyhow::anyhow!("Failed to create plot directory {}: {}", dir.display(), e)
        })?;
    }

    let estimated_trajectory_raw = read_trajectory(&estimated_path).map_err(|e| {
        anyhow::anyhow!(
            "Failed to read estimated trajectory from {}: {}",
            estimated_path.display(),
            e
        )
    })?;
    let reference_trajectory = read_trajectory(&reference_path).map_err(|e| {
        anyhow::anyhow!(
            "Failed to read reference trajectory from {}: {}",
            reference_path.display(),
            e
        )
    })?;

    let mut interpolated_estimated_trajectory = Vec::new();
    let mut synced_gt_poses = Vec::new();

    for gt_pose in &reference_trajectory {
        let mut p1_idx = 0;
        while p1_idx < estimated_trajectory_raw.len().saturating_sub(2)
            && estimated_trajectory_raw[p1_idx + 1].timestamp < gt_pose.timestamp
        {
            p1_idx += 1;
        }
        let p2_idx = p1_idx + 1;
        if p2_idx >= estimated_trajectory_raw.len() {
            continue;
        }

        let p1 = &estimated_trajectory_raw[p1_idx];
        let p2 = &estimated_trajectory_raw[p2_idx];
        let interpolated_pose = interpolate_pose(p1, p2, gt_pose.timestamp)?;
        interpolated_estimated_trajectory.push(interpolated_pose);
        synced_gt_poses.push(*gt_pose);
    }

    if synced_gt_poses.is_empty() {
        anyhow::bail!(
            "No overlapping timestamps found between {} and {}. \nEstimated time range: {} to {}\nReference time range: {} to {}",
            estimated_path.display(),
            reference_path.display(),
            estimated_trajectory_raw
                .first()
                .map_or(0.0, |p| p.timestamp),
            estimated_trajectory_raw.last().map_or(0.0, |p| p.timestamp),
            reference_trajectory.first().map_or(0.0, |p| p.timestamp),
            reference_trajectory.last().map_or(0.0, |p| p.timestamp)
        );
    }

    let corresponding_gt_trajectory = normalize_trajectory_to_first(&synced_gt_poses)?;
    let interpolated_estimated_trajectory =
        normalize_trajectory_to_first(&interpolated_estimated_trajectory)?;

    let (s, r, t) = umeyama_alignment(
        &corresponding_gt_trajectory,
        &interpolated_estimated_trajectory,
    )?;

    let aligned_est_trajectory: Vec<Pose> = interpolated_estimated_trajectory
        .iter()
        .map(|pose| {
            let aligned_trans = s * r * pose.translation.vector + t;
            let aligned_rot =
                UnitQuaternion::from_matrix(&(r * pose.rotation.to_rotation_matrix().matrix()));
            Pose {
                timestamp: pose.timestamp,
                translation: Translation3::from(aligned_trans),
                rotation: aligned_rot,
            }
        })
        .collect();

    let mut ape_trans_errors = Vec::new();
    let mut ape_rot_errors = Vec::new();
    let mut ape_data_for_plot = Vec::new();

    for (gt, est) in corresponding_gt_trajectory
        .iter()
        .zip(aligned_est_trajectory.iter())
    {
        let trans_error = (gt.translation.vector - est.translation.vector).norm();
        let rot_error = gt.rotation.angle_to(&est.rotation);
        ape_trans_errors.push(trans_error.powi(2));
        ape_rot_errors.push(rot_error.powi(2));
        ape_data_for_plot.push((gt.timestamp, trans_error));
    }

    if ape_trans_errors.is_empty() || ape_rot_errors.is_empty() {
        anyhow::bail!("Trajectory comparison produced no APE samples");
    }

    let ape_trans_rmse =
        (ape_trans_errors.iter().sum::<f64>() / ape_trans_errors.len() as f64).sqrt();
    let ape_rot_rmse = (ape_rot_errors.iter().sum::<f64>() / ape_rot_errors.len() as f64).sqrt();

    let mut rpe_trans_errors = Vec::new();
    let mut rpe_rot_errors = Vec::new();
    let mut rpe_data_for_plot = Vec::new();

    for i in 0..(corresponding_gt_trajectory.len().saturating_sub(delta)) {
        let gt_pose1 = pose_to_se3(&corresponding_gt_trajectory[i])?;
        let gt_pose2 = pose_to_se3(&corresponding_gt_trajectory[i + delta])?;
        let est_pose1 = pose_to_se3(&aligned_est_trajectory[i])?;
        let est_pose2 = pose_to_se3(&aligned_est_trajectory[i + delta])?;
        let rel_gt = gt_pose1.inverse().group_mul(gt_pose2);
        let rel_est = est_pose1.inverse().group_mul(est_pose2);

        let error_se3 = rel_gt.inverse().group_mul(rel_est);
        let trans_error = error_se3.translation().norm();
        let rot_error = error_se3.rotation().log().norm();

        rpe_trans_errors.push(trans_error.powi(2));
        rpe_rot_errors.push(rot_error.powi(2));
        rpe_data_for_plot.push((corresponding_gt_trajectory[i].timestamp, trans_error));
    }

    if rpe_trans_errors.is_empty() || rpe_rot_errors.is_empty() {
        anyhow::bail!(
            "RPE produced no samples; delta={} may be too large for {} synchronized poses",
            delta,
            corresponding_gt_trajectory.len()
        );
    }

    let rpe_trans_rmse =
        (rpe_trans_errors.iter().sum::<f64>() / rpe_trans_errors.len() as f64).sqrt();
    let rpe_rot_rmse = (rpe_rot_errors.iter().sum::<f64>() / rpe_rot_errors.len() as f64).sqrt();

    if let Some(ref out_dir) = plots_dir {
        create_error_plot(
            &out_dir.join("ape_plot.svg"),
            "Absolute Pose Error (Translation)",
            &ape_data_for_plot,
        )
        .map_err(|e| anyhow::anyhow!(e.to_string()))?;
        create_error_plot(
            &out_dir.join("rpe_plot.svg"),
            "Relative Pose Error (Translation)",
            &rpe_data_for_plot,
        )
        .map_err(|e| anyhow::anyhow!(e.to_string()))?;
        create_2d_plots(
            out_dir,
            &corresponding_gt_trajectory,
            &aligned_est_trajectory,
            &ape_data_for_plot,
            &rpe_data_for_plot,
        )
        .map_err(|e| anyhow::anyhow!(e.to_string()))?;
    }

    let results = EvalResults {
        umeyama_scale: s,
        umeyama_translation: [t.x, t.y, t.z],
        ape_trans_rmse,
        ape_rot_rmse_deg: ape_rot_rmse.to_degrees(),
        rpe_trans_rmse,
        rpe_rot_rmse_deg: rpe_rot_rmse.to_degrees(),
    };
    let results_json = serde_json::to_string_pretty(&results)?;

    if let Some(path) = results_path {
        if let Some(parent) = path.parent()
            && !parent.as_os_str().is_empty()
        {
            std::fs::create_dir_all(parent)?;
        }
        let mut file = File::create(&path)?;
        file.write_all(results_json.as_bytes())?;
    }

    if json {
        println!("{results_json}");
        let _ = std::io::stdout().flush();
    } else {
        let mut table = comfy_table::Table::new();
        table
            .load_preset(comfy_table::presets::UTF8_FULL)
            .set_header(vec![
                comfy_table::Cell::new("Metric")
                    .add_attribute(comfy_table::Attribute::Bold)
                    .fg(comfy_table::Color::Cyan),
                comfy_table::Cell::new("Value")
                    .add_attribute(comfy_table::Attribute::Bold)
                    .fg(comfy_table::Color::Cyan),
            ]);

        table.add_row(vec![
            comfy_table::Cell::new("Umeyama Scale").fg(comfy_table::Color::Yellow),
            comfy_table::Cell::new(format!("{s:.4}")),
        ]);
        table.add_row(vec![
            comfy_table::Cell::new("Umeyama Translation [x,y,z]").fg(comfy_table::Color::Yellow),
            comfy_table::Cell::new(format!("[{:.4}, {:.4}, {:.4}]", t.x, t.y, t.z)),
        ]);
        table.add_row(vec![
            comfy_table::Cell::new("APE Translational RMSE").fg(comfy_table::Color::Green),
            comfy_table::Cell::new(format!("{ape_trans_rmse:.4} m")),
        ]);
        table.add_row(vec![
            comfy_table::Cell::new("APE Rotational RMSE").fg(comfy_table::Color::Green),
            comfy_table::Cell::new(format!("{:.4} deg", ape_rot_rmse.to_degrees())),
        ]);
        table.add_row(vec![
            comfy_table::Cell::new(format!("RPE Translational RMSE (delta={delta})"))
                .fg(comfy_table::Color::Green),
            comfy_table::Cell::new(format!("{rpe_trans_rmse:.4} m")),
        ]);
        table.add_row(vec![
            comfy_table::Cell::new(format!("RPE Rotational RMSE (delta={delta})"))
                .fg(comfy_table::Color::Green),
            comfy_table::Cell::new(format!("{:.4} deg", rpe_rot_rmse.to_degrees())),
        ]);

        println!("\n{table}\n");
    }

    #[cfg(feature = "rerun")]
    if rerun_output {
        visualize_with_rerun(
            &corresponding_gt_trajectory,
            &aligned_est_trajectory,
            &ape_data_for_plot,
            &rpe_data_for_plot,
        )?;
    }
    #[cfg(not(feature = "rerun"))]
    let _ = rerun_output;

    Ok(())
}

fn main() -> anyhow::Result<()> {
    use clap::CommandFactory;
    let args = Args::parse();

    match args.command {
        Command::Completions { shell } => {
            let mut cmd = Args::command();
            clap_complete::generate(shell, &mut cmd, "sounding", &mut std::io::stdout());
            Ok(())
        }
        Command::Probe {
            target,
            interval,
            output_dir,
            json,
            #[cfg(feature = "rerun")]
            rerun,
            strict,
        } => {
            probe::run_probe(probe::ProbeConfig {
                target: &target,
                interval_secs: interval,
                out_dir: output_dir.as_path(),
                strict,
                json_output: json,
                #[cfg(feature = "rerun")]
                rerun_output: rerun,
                #[cfg(not(feature = "rerun"))]
                rerun_output: false,
            })?;
            Ok(())
        }
        Command::Trajectory {
            estimated_file,
            reference_file,
            json,
            export_results,
            export_plots,
            delta,
            #[cfg(feature = "rerun")]
            rerun,
        } => {
            let run_id = epoch_seconds();
            let results_path = resolve_optional_path(
                export_results,
                PathBuf::from(format!("trajectory_metrics_{run_id}.json")),
            );
            let plots_dir = resolve_optional_path(
                export_plots,
                PathBuf::from(format!("trajectory_plots_{run_id}")),
            );

            run_trajectory_command(
                estimated_file,
                reference_file,
                json,
                results_path,
                plots_dir,
                delta,
                #[cfg(feature = "rerun")]
                rerun,
                #[cfg(not(feature = "rerun"))]
                false,
            )
        }
    }
}

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

    #[test]
    fn resolve_optional_path_handles_missing_value() {
        let path = resolve_optional_path(Some(String::new()), PathBuf::from("default.json"));
        assert_eq!(path, Some(PathBuf::from("default.json")));
    }

    #[test]
    fn probe_output_dir_flag_parses() {
        let args = Args::try_parse_from([
            "sounding",
            "probe",
            "dlio",
            "--output-dir",
            "probe_bridge02_ouster_dlio_eval",
        ])
        .expect("probe output directory flag should parse");

        match args.command {
            Command::Probe { output_dir, .. } => {
                assert_eq!(output_dir, PathBuf::from("probe_bridge02_ouster_dlio_eval"));
            }
            _ => panic!("expected probe command"),
        }
    }

    #[test]
    fn parse_pose_parses_valid_tum_line() {
        let line = "1.0 2.0 3.0 4.0 0.0 0.0 0.0 1.0";
        let pose = parse_pose(line).expect("valid pose line");
        assert_eq!(pose.timestamp, 1.0);
        assert_eq!(pose.translation.x, 2.0);
        assert_eq!(pose.translation.y, 3.0);
        assert_eq!(pose.translation.z, 4.0);
    }

    #[test]
    fn parse_pose_converts_nanosecond_epoch_timestamps_to_seconds() {
        let line = "1690786976068341339 2.0 3.0 4.0 0.0 0.0 0.0 1.0";
        let pose = parse_pose(line).expect("valid pose line");
        assert!((pose.timestamp - 1690786976.0683413).abs() < 1e-6);
    }

    #[test]
    fn parse_pose_rejects_incomplete_line() {
        let line = "1.0 2.0 3.0";
        assert!(parse_pose(line).is_err());
    }

    #[test]
    fn normalize_trajectory_to_first_makes_poses_relative_to_start() {
        let first_rotation = UnitQuaternion::from_axis_angle(
            &nalgebra::Vector3::z_axis(),
            std::f64::consts::FRAC_PI_2,
        );
        let relative_rotation = UnitQuaternion::from_axis_angle(
            &nalgebra::Vector3::z_axis(),
            std::f64::consts::FRAC_PI_4,
        );
        let relative_translation = nalgebra::Vector3::new(1.0, 2.0, 0.0);
        let first_translation = nalgebra::Vector3::new(10.0, 0.0, 0.0);
        let second_translation = first_translation + first_rotation * relative_translation;

        let trajectory = vec![
            Pose {
                timestamp: 1.0,
                translation: Translation3::from(first_translation),
                rotation: first_rotation,
            },
            Pose {
                timestamp: 2.0,
                translation: Translation3::from(second_translation),
                rotation: first_rotation * relative_rotation,
            },
        ];

        let normalized =
            normalize_trajectory_to_first(&trajectory).expect("trajectory should normalize");

        assert!(normalized[0].translation.vector.norm() < 1e-12);
        assert!(normalized[0].rotation.angle() < 1e-12);
        assert!((normalized[1].translation.vector - relative_translation).norm() < 1e-12);
        assert!(normalized[1].rotation.angle_to(&relative_rotation) < 1e-12);
        assert_eq!(normalized[1].timestamp, 2.0);
    }
}
