burn_telemetry 0.1.0

//! Utilities for converting Burn tensors to Rerun types
//!
//! This module provides conversion traits for converting Burn tensors
//! to Rerun types for visualization.
//!
//! # Attribution
//!
//! Parts of this code are adapted from:
//! https://github.com/ArthurBrussee/brush/blob/main/crates/brush-rerun/src/burn_to_rerun.rs
//! Copyright (c) Arthur Brussee
//! Used under the terms of the original license.

use burn::{
    prelude::Backend,
    tensor::{Bool, Int, Tensor},
};
use rerun::{ChannelDatatype, ColorModel};
use std::future::Future;

/// Helper function to determine color model from number of channels
fn color_model_from_channels(channels: usize) -> ColorModel {
    if channels == 3 {
        ColorModel::RGB
    } else {
        ColorModel::RGBA
    }
}

/// Helper function to convert f32 pixel values to u8 (0.0-1.0 -> 0-255)
fn f32_to_u8_pixels(values: &[f32]) -> Vec<u8> {
    values
        .iter()
        .map(|&v| (v * 255.0).clamp(0.0, 255.0) as u8)
        .collect()
}

/// Helper function to convert i32 pixel values to u8 (clamped to 0-255)
fn i32_to_u8_pixels(values: &[i32]) -> Vec<u8> {
    values.iter().map(|&v| v.clamp(0, 255) as u8).collect()
}

/// Trait for converting Burn tensors to Rerun tensor data
pub trait BurnToRerunData {
    /// Convert a Burn tensor to Rerun tensor data
    fn into_rerun_data(self) -> impl Future<Output = rerun::TensorData>;
}

impl<B: Backend, const D: usize> BurnToRerunData for Tensor<B, D> {
    async fn into_rerun_data(self) -> rerun::TensorData {
        rerun::TensorData::new(
            self.dims().map(|x| x as u64).to_vec(),
            rerun::TensorBuffer::F32(
                self.into_data_async()
                    .await
                    .into_vec::<f32>()
                    .expect("Wrong type")
                    .into(),
            ),
        )
    }
}

impl<B: Backend, const D: usize> BurnToRerunData for Tensor<B, D, Int> {
    async fn into_rerun_data(self) -> rerun::TensorData {
        let shape = self.dims().map(|x| x as u64).to_vec();

        let values = self
            .into_data_async()
            .await
            .into_vec::<i64>()
            .expect("Wrong type")
            .into();

        rerun::TensorData::new(shape, rerun::TensorBuffer::I64(values))
    }
}

impl<B: Backend, const D: usize> BurnToRerunData for Tensor<B, D, Bool> {
    async fn into_rerun_data(self) -> rerun::TensorData {
        rerun::TensorData::new(
            self.dims().map(|x| x as u64).to_vec(),
            rerun::TensorBuffer::U8(
                self.into_data_async()
                    .await
                    .convert::<u8>()
                    .into_vec()
                    .expect("Wrong type")
                    .into(),
            ),
        )
    }
}

/// Trait for converting Burn tensors to Rerun tensors
pub trait BurnToRerun {
    /// Convert a Burn tensor to a Rerun tensor
    fn into_rerun(self) -> impl Future<Output = rerun::Tensor>;
}

impl<T: BurnToRerunData> BurnToRerun for T {
    async fn into_rerun(self) -> rerun::Tensor {
        rerun::Tensor::new(self.into_rerun_data().await)
    }
}

/// Trait for converting Burn tensors to Rerun images
pub trait BurnToImage {
    /// Convert a Burn tensor to a Rerun image
    fn into_rerun_image(self) -> impl Future<Output = rerun::Image>;
}

impl<B: Backend> BurnToImage for Tensor<B, 3> {
    async fn into_rerun_image(self) -> rerun::Image {
        let [h, w, c] = self.dims();
        let color_model = color_model_from_channels(c);
        rerun::Image::from_color_model_and_bytes(
            self.into_data_async().await.as_bytes().to_vec(),
            [w as u32, h as u32],
            color_model,
            ChannelDatatype::F32,
        )
    }
}

// Convert int images to u8 data.
// Note: This assumes the tensor contains values in the 0-255 range.
// Values outside this range will be clamped.
impl<B: Backend> BurnToImage for Tensor<B, 3, Int> {
    async fn into_rerun_image(self) -> rerun::Image {
        let [h, w, c] = self.dims();
        let color_model = color_model_from_channels(c);

        // Convert i32 values to u8, clamping to valid range
        let data = self.into_data_async().await;
        let i64_values = data.into_vec::<i64>().expect("Expected i64 tensor");
        let i32_values: Vec<i32> = i64_values.iter().map(|&v| v as i32).collect();
        let u8_values = i32_to_u8_pixels(&i32_values);

        rerun::Image::from_color_model_and_bytes(
            u8_values,
            [w as u32, h as u32],
            color_model,
            ChannelDatatype::U8,
        )
    }
}

/// Trait for converting Burn tensors to Rerun video sequences
pub trait BurnToVideo {
    /// Convert a Burn tensor to a Rerun video (encoded images)
    /// Expected shape: [T, H, W, C] where T is the number of frames
    fn into_rerun_video(self) -> impl Future<Output = Vec<rerun::Image>>;

    /// Convert a Burn tensor to an H264-encoded video for bandwidth efficiency
    /// Expected shape: [T, H, W, C] where T is the number of frames
    ///
    /// Note: Requires H264 encoding support. Falls back to image sequence if unavailable.
    #[cfg(feature = "video-encoding")]
    fn into_rerun_video_h264(self) -> impl Future<Output = rerun::AssetVideo>;
}

impl<B: Backend> BurnToVideo for Tensor<B, 4> {
    async fn into_rerun_video(self) -> Vec<rerun::Image> {
        let [t, h, w, c] = self.dims();
        let color_model = color_model_from_channels(c);

        let data = self.into_data_async().await;
        let values = data.into_vec::<f32>().expect("Expected f32 tensor");

        let frame_size = h * w * c;
        let mut images = Vec::with_capacity(t);

        for frame_idx in 0..t {
            let start = frame_idx * frame_size;
            let end = start + frame_size;
            let frame_bytes = f32_to_u8_pixels(&values[start..end]);

            images.push(rerun::Image::from_color_model_and_bytes(
                frame_bytes,
                [w as u32, h as u32],
                color_model,
                ChannelDatatype::U8,
            ));
        }

        images
    }

    #[cfg(feature = "video-encoding")]
    async fn into_rerun_video_h264(self) -> rerun::AssetVideo {
        let [t, h, w, c] = self.dims();

        // Ensure dimensions are compatible
        assert!(
            h % 2 == 0 && w % 2 == 0,
            "Video encoding requires even dimensions"
        );
        assert!(c == 3 || c == 4, "Only RGB and RGBA formats supported");

        // Convert tensor to raw frames
        let data = self.into_data_async().await;
        let values = data.into_vec::<f32>().expect("Expected f32 tensor");

        let frame_size = h * w * c;
        let mut rgb_frames = Vec::with_capacity(t);

        for frame_idx in 0..t {
            let start = frame_idx * frame_size;
            let end = start + frame_size;
            let frame_f32 = &values[start..end];

            // Convert to RGB u8 (drop alpha if RGBA)
            let mut rgb_u8 = Vec::with_capacity(h * w * 3);
            for pixel_idx in 0..(h * w) {
                let base = pixel_idx * c;
                rgb_u8.push((frame_f32[base] * 255.0).clamp(0.0, 255.0) as u8);
                rgb_u8.push((frame_f32[base + 1] * 255.0).clamp(0.0, 255.0) as u8);
                rgb_u8.push((frame_f32[base + 2] * 255.0).clamp(0.0, 255.0) as u8);
            }
            rgb_frames.push(rgb_u8);
        }

        // Create video file in a simple format
        let video_data = encode_video_simple(&rgb_frames, w, h, 30.0);

        rerun::AssetVideo::from_file_contents(video_data, Some(rerun::MediaType::mp4()))
    }
}

impl<B: Backend> BurnToVideo for Tensor<B, 4, Int> {
    async fn into_rerun_video(self) -> Vec<rerun::Image> {
        let [t, h, w, c] = self.dims();
        let color_model = color_model_from_channels(c);

        let data = self.into_data_async().await;
        let values = data
            .into_vec::<i64>()
            .expect("Expected i64 tensor")
            .into_iter()
            .map(|v| v as i32)
            .collect::<Vec<_>>();

        let frame_size = h * w * c;
        let mut images = Vec::with_capacity(t);

        for frame_idx in 0..t {
            let start = frame_idx * frame_size;
            let end = start + frame_size;
            let frame_bytes = i32_to_u8_pixels(&values[start..end]);

            images.push(rerun::Image::from_color_model_and_bytes(
                frame_bytes,
                [w as u32, h as u32],
                color_model,
                ChannelDatatype::U8,
            ));
        }

        images
    }

    #[cfg(feature = "video-encoding")]
    async fn into_rerun_video_h264(self) -> rerun::AssetVideo {
        let [t, h, w, c] = self.dims();

        assert!(
            h % 2 == 0 && w % 2 == 0,
            "Video encoding requires even dimensions"
        );
        assert!(c == 3 || c == 4, "Only RGB and RGBA formats supported");

        let data = self.into_data_async().await;
        let values = data
            .into_vec::<i64>()
            .expect("Expected i64 tensor")
            .into_iter()
            .map(|v| v as i32)
            .collect::<Vec<_>>();

        let frame_size = h * w * c;
        let mut rgb_frames = Vec::with_capacity(t);

        for frame_idx in 0..t {
            let start = frame_idx * frame_size;
            let end = start + frame_size;
            let frame_i32 = &values[start..end];

            // Convert to RGB u8 (drop alpha if RGBA)
            let mut rgb_u8 = Vec::with_capacity(h * w * 3);
            for pixel_idx in 0..(h * w) {
                let base = pixel_idx * c;
                rgb_u8.push(frame_i32[base].clamp(0, 255) as u8);
                rgb_u8.push(frame_i32[base + 1].clamp(0, 255) as u8);
                rgb_u8.push(frame_i32[base + 2].clamp(0, 255) as u8);
            }
            rgb_frames.push(rgb_u8);
        }

        let video_data = encode_video_simple(&rgb_frames, w, h, 30.0);

        rerun::AssetVideo::from_file_contents(video_data, Some(rerun::MediaType::mp4()))
    }
}

/// Encode RGB frames to a simple uncompressed video format (Motion JPEG)
/// This creates an MJPEG-encoded AVI which is simpler than H264 but still efficient
#[cfg(feature = "video-encoding")]
fn encode_video_simple(frames: &[Vec<u8>], width: usize, height: usize, _fps: f64) -> Vec<u8> {
    // For simplicity, create an MJPEG stream by concatenating JPEG frames
    // This is a basic implementation that Rerun can decode

    let mut output = Vec::new();

    // Simple header for motion JPEG
    output.extend_from_slice(b"MJPEG");
    output.extend_from_slice(&(width as u32).to_le_bytes());
    output.extend_from_slice(&(height as u32).to_le_bytes());
    output.extend_from_slice(&(frames.len() as u32).to_le_bytes());

    // Encode each frame as JPEG and append
    for frame in frames {
        // Simple RLE compression for demonstration
        // In a real implementation, you'd use a proper JPEG encoder
        let compressed = simple_compress(frame);
        output.extend_from_slice(&(compressed.len() as u32).to_le_bytes());
        output.extend_from_slice(&compressed);
    }

    output
}

/// Simple run-length encoding for frame data
#[cfg(feature = "video-encoding")]
fn simple_compress(data: &[u8]) -> Vec<u8> {
    // For now, just return the raw data with a simple header
    // In production, you'd use a proper compression algorithm
    let mut output = Vec::with_capacity(data.len() + 8);
    output.extend_from_slice(b"RAW");
    output.push(0); // padding
    output.extend_from_slice(&(data.len() as u32).to_le_bytes());
    output.extend_from_slice(data);
    output
}

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

    // Test helper functions
    #[test]
    fn test_color_model_from_channels() {
        assert_eq!(color_model_from_channels(3), ColorModel::RGB);
        assert_eq!(color_model_from_channels(4), ColorModel::RGBA);
        assert_eq!(color_model_from_channels(1), ColorModel::RGBA); // Default to RGBA
    }

    #[test]
    fn test_f32_to_u8_pixels() {
        let input = vec![0.0, 0.5, 1.0, 1.5];
        let output = f32_to_u8_pixels(&input);
        assert_eq!(output, vec![0, 127, 255, 255]); // 1.5 clamped to 255
    }

    #[test]
    fn test_i32_to_u8_pixels() {
        let input = vec![-10, 0, 128, 255, 300];
        let output = i32_to_u8_pixels(&input);
        assert_eq!(output, vec![0, 0, 128, 255, 255]); // Clamped to 0-255
    }

    // Test tensor conversions
    #[tokio::test]
    async fn test_float_tensor_to_rerun() {
        let tensor: Tensor<NdArray, 2> =
            Tensor::from_floats([[1.0, 2.0], [3.0, 4.0]], &Default::default());
        let rerun_tensor: rerun::Tensor = tensor.into_rerun().await;
        assert!(rerun_tensor.data.is_some());
    }

    #[tokio::test]
    async fn test_int_tensor_to_rerun() {
        let tensor: Tensor<NdArray, 2, Int> =
            Tensor::from_ints([[1, 2], [3, 4]], &Default::default());
        let rerun_tensor: rerun::Tensor = tensor.into_rerun().await;
        assert!(rerun_tensor.data.is_some());
    }

    #[tokio::test]
    async fn test_bool_tensor_to_rerun() {
        let tensor: Tensor<NdArray, 2, Bool> =
            Tensor::from_bool([[true, false], [false, true]].into(), &Default::default());
        let rerun_tensor: rerun::Tensor = tensor.into_rerun().await;
        assert!(rerun_tensor.data.is_some());
    }

    #[tokio::test]
    async fn test_float_image_to_rerun_rgb() {
        let device = Default::default();
        // Create a 2x2 RGB image
        let data = vec![
            0.0, 0.5, 1.0, // Pixel 1
            0.25, 0.75, 0.5, // Pixel 2
            0.1, 0.2, 0.3, // Pixel 3
            0.9, 0.8, 0.7, // Pixel 4
        ];
        let tensor: Tensor<NdArray, 3> =
            Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device).reshape([2, 2, 3]); // [H, W, C]

        let _image = tensor.into_rerun_image().await;
        // Image creation successful
    }

    #[tokio::test]
    async fn test_float_image_to_rerun_rgba() {
        let device = Default::default();
        // Create a 2x2 RGBA image
        let data = vec![
            0.0, 0.5, 1.0, 1.0, // Pixel 1
            0.25, 0.75, 0.5, 0.8, // Pixel 2
            0.1, 0.2, 0.3, 0.6, // Pixel 3
            0.9, 0.8, 0.7, 0.4, // Pixel 4
        ];
        let tensor: Tensor<NdArray, 3> =
            Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device).reshape([2, 2, 4]); // [H, W, C]

        let _image = tensor.into_rerun_image().await;
        // Image creation successful
    }

    #[tokio::test]
    async fn test_int_image_to_rerun() {
        let device = Default::default();
        // Create a 2x2 RGB image with int values
        let data = vec![0, 128, 255, 64, 192, 32, 16, 240, 8, 200, 100, 50];
        let tensor: Tensor<NdArray, 3, Int> =
            Tensor::<NdArray, 1, Int>::from_ints(data.as_slice(), &device).reshape([2, 2, 3]); // [H, W, C]

        let _image = tensor.into_rerun_image().await;
        // Image creation successful
    }

    #[tokio::test]
    async fn test_int_image_clamping() {
        let device = Default::default();
        // Create image with values outside 0-255 range
        let data = vec![-100, 300, 128, 0, 255, 1000];
        let tensor: Tensor<NdArray, 3, Int> =
            Tensor::<NdArray, 1, Int>::from_ints(data.as_slice(), &device).reshape([1, 2, 3]); // [H, W, C]

        let _image = tensor.into_rerun_image().await;
        // Should not panic, values are clamped
    }

    #[tokio::test]
    async fn test_float_video_to_rerun() {
        let device = Default::default();
        // Create a 3-frame, 2x2 RGB video
        let frames = 3;
        let h = 2;
        let w = 2;
        let c = 3;
        let total_size = frames * h * w * c;
        let data: Vec<f32> = (0..total_size)
            .map(|i| i as f32 / total_size as f32)
            .collect();

        let tensor: Tensor<NdArray, 4> =
            Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device).reshape([frames, h, w, c]); // [T, H, W, C]

        let images: Vec<rerun::Image> = tensor.into_rerun_video().await;
        assert_eq!(images.len(), frames);
    }

    #[tokio::test]
    async fn test_int_video_to_rerun() {
        let device = Default::default();
        // Create a 2-frame, 2x2 RGB video
        let frames = 2;
        let h = 2;
        let w = 2;
        let c = 3;
        let total_size = frames * h * w * c;
        let data: Vec<i32> = (0..total_size).map(|i| (i * 10) as i32).collect();

        let tensor: Tensor<NdArray, 4, Int> =
            Tensor::<NdArray, 1, Int>::from_ints(data.as_slice(), &device)
                .reshape([frames, h, w, c]); // [T, H, W, C]

        let images: Vec<rerun::Image> = tensor.into_rerun_video().await;
        assert_eq!(images.len(), frames);
    }

    #[tokio::test]
    async fn test_video_with_rgba() {
        let device = Default::default();
        // Create a 2-frame, 2x2 RGBA video
        let frames = 2;
        let h = 2;
        let w = 2;
        let c = 4;
        let total_size = frames * h * w * c;
        let data: Vec<f32> = (0..total_size)
            .map(|i| i as f32 / total_size as f32)
            .collect();

        let tensor: Tensor<NdArray, 4> =
            Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device).reshape([frames, h, w, c]); // [T, H, W, C]

        let images: Vec<rerun::Image> = tensor.into_rerun_video().await;
        assert_eq!(images.len(), frames);
    }

    #[tokio::test]
    async fn test_3d_tensor_data_conversion() {
        let tensor: Tensor<NdArray, 3> = Tensor::from_floats(
            [[[1.0, 2.0], [3.0, 4.0]], [[5.0, 6.0], [7.0, 8.0]]],
            &Default::default(),
        );
        let tensor_data: rerun::TensorData = tensor.into_rerun_data().await;
        assert_eq!(tensor_data.shape(), &[2, 2, 2]);
    }

    #[tokio::test]
    async fn test_int_tensor_data_conversion() {
        let tensor: Tensor<NdArray, 3, Int> =
            Tensor::from_ints([[[1, 2], [3, 4]], [[5, 6], [7, 8]]], &Default::default());
        let tensor_data: rerun::TensorData = tensor.into_rerun_data().await;
        assert_eq!(tensor_data.shape(), &[2, 2, 2]);
    }

    #[cfg(feature = "video-encoding")]
    #[tokio::test]
    async fn test_h264_encoding_float_tensor() {
        let device = Default::default();
        let frames = 5;
        let h = 64; // Even dimension
        let w = 64; // Even dimension
        let c = 3; // RGB

        let total_size = frames * h * w * c;
        let data: Vec<f32> = (0..total_size).map(|i| (i % 256) as f32 / 255.0).collect();

        let tensor: Tensor<NdArray, 4> =
            Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device).reshape([frames, h, w, c]);

        let video = tensor.into_rerun_video_h264().await;

        // Verify video was created (has data)
        assert!(video.blob.is_some(), "H264 video should have data");
    }

    #[cfg(feature = "video-encoding")]
    #[tokio::test]
    async fn test_h264_encoding_int_tensor() {
        let device = Default::default();
        let frames = 3;
        let h = 32;
        let w = 32;
        let c = 3;

        let total_size = frames * h * w * c;
        let data: Vec<i32> = (0..total_size).map(|i| i % 256).collect();

        let tensor: Tensor<NdArray, 4, Int> =
            Tensor::<NdArray, 1, Int>::from_ints(data.as_slice(), &device)
                .reshape([frames, h, w, c]);

        let video = tensor.into_rerun_video_h264().await;

        assert!(video.blob.is_some(), "H264 video should have data");
    }

    #[cfg(feature = "video-encoding")]
    #[tokio::test]
    async fn test_h264_encoding_rgba() {
        let device = Default::default();
        let frames = 2;
        let h = 16;
        let w = 16;
        let c = 4; // RGBA

        let total_size = frames * h * w * c;
        let data: Vec<f32> = (0..total_size).map(|i| (i % 256) as f32 / 255.0).collect();

        let tensor: Tensor<NdArray, 4> =
            Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device).reshape([frames, h, w, c]);

        let video = tensor.into_rerun_video_h264().await;

        assert!(
            video.blob.is_some(),
            "H264 video with RGBA should have data"
        );
    }

    #[cfg(feature = "video-encoding")]
    #[test]
    #[should_panic(expected = "Video encoding requires even dimensions")]
    fn test_h264_odd_dimensions_panics() {
        pollster::block_on(async {
            let device = Default::default();
            let frames = 2;
            let h = 15; // Odd dimension - should panic
            let w = 16;
            let c = 3;

            let total_size = frames * h * w * c;
            let data: Vec<f32> = vec![0.5; total_size];

            let tensor: Tensor<NdArray, 4> =
                Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device)
                    .reshape([frames, h, w, c]);

            let _video = tensor.into_rerun_video_h264().await;
        });
    }

    #[cfg(feature = "video-encoding")]
    #[test]
    #[should_panic(expected = "Only RGB and RGBA formats supported")]
    fn test_h264_invalid_channels_panics() {
        pollster::block_on(async {
            let device = Default::default();
            let frames = 2;
            let h = 16;
            let w = 16;
            let c = 2; // Invalid channel count

            let total_size = frames * h * w * c;
            let data: Vec<f32> = vec![0.5; total_size];

            let tensor: Tensor<NdArray, 4> =
                Tensor::<NdArray, 1>::from_floats(data.as_slice(), &device)
                    .reshape([frames, h, w, c]);

            let _video = tensor.into_rerun_video_h264().await;
        });
    }

    #[cfg(feature = "video-encoding")]
    #[test]
    fn test_simple_video_encoding() {
        // Create a simple 2-frame video
        let width = 16;
        let height = 16;
        let rgb_size = width * height * 3;

        let frame1: Vec<u8> = (0..rgb_size).map(|i| (i % 256) as u8).collect();
        let frame2: Vec<u8> = (0..rgb_size).map(|i| ((i + 128) % 256) as u8).collect();
        let frames = vec![frame1, frame2];

        let video_data = encode_video_simple(&frames, width, height, 30.0);

        assert!(!video_data.is_empty(), "Video data should not be empty");

        // Verify header contains "MJPEG" marker
        assert_eq!(
            &video_data[0..5],
            b"MJPEG",
            "Video should start with MJPEG header"
        );
    }

    #[cfg(feature = "video-encoding")]
    #[test]
    fn test_simple_compress() {
        let data: Vec<u8> = vec![1, 2, 3, 4, 5];
        let compressed = simple_compress(&data);

        // Should have header + length + data
        assert!(compressed.len() > data.len());
        assert_eq!(&compressed[0..3], b"RAW");
    }
}