use align3d::{
    bilateral::BilateralFilter,
    icp::{multiscale::MultiscaleAlign, MsIcpParams},
    io::dataset::{IndoorLidarDataset, RgbdDataset, SubsetDataset},
    metrics::TransformMetrics,
    range_image::RangeImageBuilder,
    trajectory_builder::TrajectoryBuilder,
    viz::rgbd_dataset_viewer::RgbdDatasetViewer,
};

fn main() -> Result<(), Box<dyn std::error::Error + 'static>> {
    // Load the dataset
    let dataset = Box::new(SubsetDataset::new(
        Box::new(IndoorLidarDataset::load("tests/data/indoor_lidar/bedroom")?),
        (0..20).collect(),
    ));

    // RangeImageBuilder composes the processing steps when loading RGB-D frames (or `RangeImage`).
    let range_image_transform = RangeImageBuilder::default()
        .with_intensity(true) // Use intensity besides RGB
        .with_normals(true) // Compute the normals
        .with_bilateral_filter(Some(BilateralFilter::default())) // Apply bilateral filter
        .pyramid_levels(3); // Compute 3-level Gaussian pyramid.

    // Default ICP parameters
    let icp_params = MsIcpParams::default();

    // TrajectoryBuilder accumulates the per-frame alignment to form the odometry of the camera poses.
    let mut traj_builder = TrajectoryBuilder::default();

    // Use the `.build()` method to create a RangeImage pyramid.
    let mut prev_frame = range_image_transform.build(dataset.get(0).unwrap());

    // Iterate over the dataset
    for i in 1..dataset.len() {
        let current_frame = range_image_transform.build(dataset.get(i).unwrap());
        
        // Perform ICP alignment
        let icp = MultiscaleAlign::new(icp_params.clone(), &prev_frame).unwrap();
        let transform = icp.align(&current_frame);

        // Accumulate transformations for obtaining odometry
        traj_builder.accumulate(&transform, Some(i as f32));
        prev_frame = current_frame;
    }

    // Compute metrics in relation to the ground truth
    // Get the predicted trajectory
    let pred_trajectory = traj_builder.build();
    
    // Get the ground truth trajectory
    let gt_trajectory = &dataset
        .trajectory()
        .expect("Dataset has no trajectory")
        .first_frame_at_origin();

    // Compute the metrics
    let metrics = TransformMetrics::mean_trajectory_error(&pred_trajectory, &gt_trajectory)?;
    println!("Mean trajectory error: {}", metrics);

    // Visualization part
    RgbdDatasetViewer::new(dataset)
        .with_trajectory(pred_trajectory.clone())
        .run();

    Ok(())
}