re_view_spatial 0.31.4

use nohash_hasher::IntMap;

use re_log_types::EntityPathHash;
use re_sdk_types::{
    Archetype as _,
    archetypes::EncodedDepthImage,
    components::{Blob, Colormap, MagnificationFilter, MediaType},
};
use re_viewer_context::{
    IdentifiedViewSystem, ImageDecodeCache, ViewContext, ViewContextCollection, ViewQuery,
    ViewSystemExecutionError, VisualizerExecutionOutput, VisualizerQueryInfo,
    VisualizerReportSeverity, VisualizerSystem,
};

use super::entity_iterator::process_archetype;
use super::{
    SpatialViewVisualizerData,
    depth_images::{DepthImageComponentData, process_depth_image_data},
};
use crate::{
    contexts::TransformTreeContext,
    view_kind::SpatialViewKind,
    visualizers::{
        depth_images::{DepthImageProcessResult, populate_depth_visualizer_execution_result},
        first_copied,
    },
};
use re_sdk_types::reflection::Enum as _;

pub struct EncodedDepthImageVisualizer {
    pub data: SpatialViewVisualizerData,

    /// Expose image infos for depth clouds - we need this for picking interaction.
    pub depth_cloud_entities: IntMap<EntityPathHash, DepthImageProcessResult>,
}

impl Default for EncodedDepthImageVisualizer {
    fn default() -> Self {
        Self {
            data: SpatialViewVisualizerData::new(Some(SpatialViewKind::TwoD)),
            depth_cloud_entities: Default::default(),
        }
    }
}

impl IdentifiedViewSystem for EncodedDepthImageVisualizer {
    fn identifier() -> re_viewer_context::ViewSystemIdentifier {
        "EncodedDepthImage".into()
    }
}

impl VisualizerSystem for EncodedDepthImageVisualizer {
    fn visualizer_query_info(
        &self,
        _app_options: &re_viewer_context::AppOptions,
    ) -> VisualizerQueryInfo {
        VisualizerQueryInfo::single_required_component::<Blob>(
            &EncodedDepthImage::descriptor_blob(),
            &EncodedDepthImage::all_components(),
        )
    }

    fn execute(
        &mut self,
        ctx: &ViewContext<'_>,
        view_query: &ViewQuery<'_>,
        context_systems: &ViewContextCollection,
    ) -> Result<VisualizerExecutionOutput, ViewSystemExecutionError> {
        let preferred_view_kind = self.data.preferred_view_kind;
        let output = VisualizerExecutionOutput::default();
        let mut depth_clouds = Vec::new();

        let transforms = context_systems.get::<TransformTreeContext>(&output)?;

        process_archetype::<Self, EncodedDepthImage, _>(
            ctx,
            view_query,
            context_systems,
            &output,
            preferred_view_kind,
            |ctx, spatial_ctx, results| {
                let all_blobs =
                    results.iter_required(EncodedDepthImage::descriptor_blob().component);
                if all_blobs.is_empty() {
                    return Ok(());
                }
                let all_media_types =
                    results.iter_optional(EncodedDepthImage::descriptor_media_type().component);
                let all_colormaps =
                    results.iter_optional(EncodedDepthImage::descriptor_colormap().component);
                let all_value_ranges =
                    results.iter_optional(EncodedDepthImage::descriptor_depth_range().component);
                let all_depth_meters =
                    results.iter_optional(EncodedDepthImage::descriptor_meter().component);
                let all_fill_ratios = results
                    .iter_optional(EncodedDepthImage::descriptor_point_fill_ratio().component);
                let all_magnification_filters = results
                    .iter_optional(EncodedDepthImage::descriptor_magnification_filter().component);

                for (
                    (_time, row_id),
                    blobs,
                    media_type,
                    colormap,
                    value_range,
                    depth_meter,
                    fill_ratio,
                    magnification_filter,
                ) in re_query::range_zip_1x6(
                    all_blobs.slice::<&[u8]>(),
                    all_media_types.slice::<String>(),
                    all_colormaps.slice::<u8>(),
                    all_value_ranges.slice::<[f64; 2]>(),
                    all_depth_meters.slice::<f32>(),
                    all_fill_ratios.slice::<f32>(),
                    all_magnification_filters.slice::<u8>(),
                ) {
                    let Some(blob) = blobs.first() else {
                        // If missing we already reported an error.
                        continue;
                    };

                    let media_type = media_type
                        .and_then(|types| types.first().cloned())
                        .map(|mt| MediaType(mt.into()));

                    let image = match ctx.store_ctx().memoizer(|c: &mut ImageDecodeCache| {
                        c.entry_encoded_depth(
                            row_id,
                            EncodedDepthImage::descriptor_blob().component,
                            blob,
                            media_type.as_ref(),
                        )
                    }) {
                        Ok(image) => image,
                        Err(err) => {
                            results.report_for_component(
                                EncodedDepthImage::descriptor_blob().component,
                                VisualizerReportSeverity::Error,
                                format!("Failed to decode EncodedDepthImage blob: {err}"),
                            );
                            continue;
                        }
                    };

                    let data = DepthImageComponentData {
                        image,
                        depth_meter: first_copied(depth_meter).map(Into::into),
                        fill_ratio: first_copied(fill_ratio).map(Into::into),
                        colormap: colormap.and_then(|s| Colormap::from_integer_slice(s).next()?),
                        value_range: first_copied(value_range),
                        magnification_filter: first_copied(magnification_filter)
                            .and_then(MagnificationFilter::from_u8)
                            .unwrap_or_default(),
                    };

                    let mut report_error = |error: String| {
                        results.report_unspecified_source(VisualizerReportSeverity::Error, error);
                    };

                    process_depth_image_data(
                        ctx,
                        spatial_ctx,
                        &mut self.data,
                        &mut self.depth_cloud_entities,
                        &mut depth_clouds,
                        transforms,
                        data,
                        EncodedDepthImage::name(),
                        EncodedDepthImage::descriptor_meter().component,
                        EncodedDepthImage::descriptor_colormap().component,
                        &mut report_error,
                    );
                }

                Ok(())
            },
        )?;

        populate_depth_visualizer_execution_result(ctx, &self.data, depth_clouds, output)
    }

    fn data(&self) -> Option<&dyn std::any::Any> {
        Some(self.data.as_any())
    }
}