ui.rs

use std::{collections::BTreeMap, fmt::Display};

use eframe::emath::Align2;
use egui::{epaint::TextShape, NumExt as _, Vec2};
use ndarray::Axis;

use re_log_types::component_types::{self, Tensor};
use re_renderer::Colormap;
use re_tensor_ops::dimension_mapping::{DimensionMapping, DimensionSelector};

use crate::ui::data_ui::image::tensor_summary_ui_grid_contents;

use super::dimension_mapping_ui;

// ---

/// How we slice a given tensor
#[derive(Clone, Debug, Hash, serde::Deserialize, serde::Serialize)]
pub struct SliceSelection {
    /// How we select which dimensions to project the tensor onto.
    pub dim_mapping: DimensionMapping,

    /// Selected value of every dimension (iff they are in [`DimensionMapping::selectors`]).
    pub selector_values: BTreeMap<usize, u64>,
}

#[derive(Clone, Debug, serde::Deserialize, serde::Serialize)]
pub struct ViewTensorState {
    /// What slice are we vieiwing?
    slice: SliceSelection,

    /// How we map values to colors.
    color_mapping: ColorMapping,

    /// Scaling, filtering, aspect ratio, etc for the rendered texture.
    texture_settings: TextureSettings,

    /// Last viewed tensor, copied each frame.
    /// Used for the selection view.
    #[serde(skip)]
    tensor: Option<Tensor>,
}

impl ViewTensorState {
    pub fn create(tensor: &Tensor) -> ViewTensorState {
        Self {
            slice: SliceSelection {
                dim_mapping: DimensionMapping::create(tensor.shape()),
                selector_values: Default::default(),
            },
            color_mapping: ColorMapping::default(),
            texture_settings: TextureSettings::default(),
            tensor: Some(tensor.clone()),
        }
    }

    pub fn slice(&self) -> &SliceSelection {
        &self.slice
    }

    pub fn color_mapping(&self) -> &ColorMapping {
        &self.color_mapping
    }

    pub(crate) fn ui(&mut self, ctx: &mut crate::misc::ViewerContext<'_>, ui: &mut egui::Ui) {
        let Some(tensor) = &self.tensor else {
            ui.label("No Tensor shown in this Space View.");
            return;
        };

        ctx.re_ui
            .selection_grid(ui, "tensor_selection_ui")
            .show(ui, |ui| {
                tensor_summary_ui_grid_contents(
                    ctx.re_ui,
                    ui,
                    tensor,
                    Some(ctx.cache.tensor_stats(tensor)),
                );
                self.texture_settings.ui(ctx.re_ui, ui);
                self.color_mapping.ui(ctx.render_ctx, ctx.re_ui, ui);
            });

        ui.separator();
        ui.strong("Dimension Mapping");
        dimension_mapping_ui(ctx.re_ui, ui, &mut self.slice.dim_mapping, tensor.shape());
        let default_mapping = DimensionMapping::create(tensor.shape());
        if ui
            .add_enabled(
                self.slice.dim_mapping != default_mapping,
                egui::Button::new("Reset mapping"),
            )
            .on_disabled_hover_text("The default is already set up.")
            .on_hover_text("Reset dimension mapping to the default.")
            .clicked()
        {
            self.slice.dim_mapping = DimensionMapping::create(tensor.shape());
        }
    }
}

pub(crate) fn view_tensor(
    ctx: &mut crate::misc::ViewerContext<'_>,
    ui: &mut egui::Ui,
    state: &mut ViewTensorState,
    tensor: &Tensor,
) {
    crate::profile_function!();

    state.tensor = Some(tensor.clone());

    if !state.slice.dim_mapping.is_valid(tensor.num_dim()) {
        state.slice.dim_mapping = DimensionMapping::create(tensor.shape());
    }

    let default_item_spacing = ui.spacing_mut().item_spacing;
    ui.spacing_mut().item_spacing.y = 0.0; // No extra spacing between sliders and tensor

    if state
        .slice
        .dim_mapping
        .selectors
        .iter()
        .any(|selector| selector.visible)
    {
        egui::Frame {
            inner_margin: egui::Margin::symmetric(16.0, 8.0),
            ..Default::default()
        }
        .show(ui, |ui| {
            ui.spacing_mut().item_spacing = default_item_spacing; // keep the default spacing between sliders
            selectors_ui(ui, state, tensor);
        });
    }

    let dimension_labels = {
        let dm = &state.slice.dim_mapping;
        [
            (
                dimension_name(&tensor.shape, dm.width.unwrap()),
                dm.invert_width,
            ),
            (
                dimension_name(&tensor.shape, dm.height.unwrap()),
                dm.invert_height,
            ),
        ]
    };

    egui::ScrollArea::both().show(ui, |ui| {
        if let Err(err) = tensor_slice_ui(ctx, ui, state, tensor, dimension_labels) {
            ui.label(ctx.re_ui.error_text(err.to_string()));
        }
    });
}

fn tensor_slice_ui(
    ctx: &mut crate::misc::ViewerContext<'_>,
    ui: &mut egui::Ui,
    state: &mut ViewTensorState,
    tensor: &Tensor,
    dimension_labels: [(String, bool); 2],
) -> anyhow::Result<()> {
    let (response, painter, image_rect) = paint_tensor_slice(ctx, ui, state, tensor)?;

    if !response.hovered() {
        let font_id = egui::TextStyle::Body.resolve(ui.style());
        paint_axis_names(ui, &painter, image_rect, font_id, dimension_labels);
    }

    Ok(())
}

fn paint_tensor_slice(
    ctx: &mut crate::misc::ViewerContext<'_>,
    ui: &mut egui::Ui,
    state: &mut ViewTensorState,
    tensor: &Tensor,
) -> anyhow::Result<(egui::Response, egui::Painter, egui::Rect)> {
    crate::profile_function!();

    let tensor_stats = ctx.cache.tensor_stats(tensor);
    let colormapped_texture = super::tensor_slice_to_gpu::colormapped_texture(
        ctx.render_ctx,
        tensor,
        tensor_stats,
        state,
    )?;
    let texture = ctx
        .render_ctx
        .texture_manager_2d
        .get(&colormapped_texture.texture)?;
    let size = texture.creation_desc.size;
    let width = size.width;
    let height = size.height;

    let img_size = egui::vec2(width as _, height as _);
    let img_size = Vec2::max(Vec2::splat(1.0), img_size); // better safe than sorry
    let desired_size = match state.texture_settings.scaling {
        TextureScaling::Original => img_size,
        TextureScaling::Fill => {
            let desired_size = ui.available_size();
            if state.texture_settings.keep_aspect_ratio {
                let scale = (desired_size / img_size).min_elem();
                img_size * scale
            } else {
                desired_size
            }
        }
    };

    let (response, painter) = ui.allocate_painter(desired_size, egui::Sense::hover());
    let rect = response.rect;
    let image_rect = egui::Rect::from_min_max(rect.min, rect.max);

    let debug_name = "tensor_slice";
    crate::gpu_bridge::render_image(
        ctx.render_ctx,
        &painter,
        image_rect,
        colormapped_texture,
        state.texture_settings.options,
        debug_name,
    )?;

    Ok((response, painter, image_rect))
}

// ----------------------------------------------------------------------------

/// How we map values to colors.
#[derive(Copy, Clone, Debug, serde::Deserialize, serde::Serialize)]
pub struct ColorMapping {
    pub map: Colormap,
    pub gamma: f32,
}

impl Default for ColorMapping {
    fn default() -> Self {
        Self {
            map: Colormap::Viridis,
            gamma: 1.0,
        }
    }
}

impl ColorMapping {
    fn ui(
        &mut self,
        render_ctx: &mut re_renderer::RenderContext,
        re_ui: &re_ui::ReUi,
        ui: &mut egui::Ui,
    ) {
        let ColorMapping { map, gamma } = self;

        re_ui.grid_left_hand_label(ui, "Color map");
        egui::ComboBox::from_id_source("color map select")
            .selected_text(map.to_string())
            .show_ui(ui, |ui| {
                ui.style_mut().wrap = Some(false);

                egui::Grid::new("colormap_selector")
                    .num_columns(2)
                    .show(ui, |ui| {
                        for option in Colormap::ALL {
                            ui.selectable_value(map, option, option.to_string());
                            colormap_preview_ui(render_ctx, ui, option);
                            ui.end_row();
                        }
                    });
            });
        ui.end_row();

        re_ui.grid_left_hand_label(ui, "Brightness");
        let mut brightness = 1.0 / *gamma;
        ui.add(egui::Slider::new(&mut brightness, 0.1..=10.0).logarithmic(true));
        *gamma = 1.0 / brightness;
        ui.end_row();
    }
}

/// Show the given colormap as a horizontal bar.
fn colormap_preview_ui(
    render_ctx: &mut re_renderer::RenderContext,
    ui: &mut egui::Ui,
    colormap: Colormap,
) -> egui::Response {
    crate::profile_function!();

    let desired_size = egui::vec2(128.0, 16.0);
    let (rect, response) = ui.allocate_exact_size(desired_size, egui::Sense::hover());

    if ui.is_rect_visible(rect) {
        if let Err(err) = paint_colormap_gradient(render_ctx, colormap, ui, rect) {
            re_log::error_once!("Failed to paint colormap preview: {err}");
        }
    }

    response
}

fn paint_colormap_gradient(
    render_ctx: &mut re_renderer::RenderContext,
    colormap: Colormap,
    ui: &mut egui::Ui,
    rect: egui::Rect,
) -> anyhow::Result<()> {
    let horizontal_gradient_id = egui::util::hash("horizontal_gradient");
    let horizontal_gradient =
        crate::gpu_bridge::get_or_create_texture(render_ctx, horizontal_gradient_id, || {
            let width = 256;
            let height = 1;
            let data: Vec<u8> = (0..width)
                .flat_map(|x| {
                    let t = x as f32 / (width as f32 - 1.0);
                    half::f16::from_f32(t).to_le_bytes()
                })
                .collect();

            re_renderer::resource_managers::Texture2DCreationDesc {
                label: "horizontal_gradient".into(),
                data: data.into(),
                format: wgpu::TextureFormat::R16Float,
                width,
                height,
            }
        });

    let colormapped_texture = re_renderer::renderer::ColormappedTexture {
        texture: horizontal_gradient,
        range: [0.0, 1.0],
        gamma: 1.0,
        color_mapper: Some(re_renderer::renderer::ColorMapper::Function(colormap)),
    };

    let debug_name = format!("colormap_{colormap}");
    crate::gpu_bridge::render_image(
        render_ctx,
        ui.painter(),
        rect,
        colormapped_texture,
        egui::TextureOptions::LINEAR,
        &debug_name,
    )
}

// ----------------------------------------------------------------------------

/// Should we scale the rendered texture, and if so, how?
#[derive(Copy, Clone, Debug, PartialEq, serde::Deserialize, serde::Serialize)]
enum TextureScaling {
    /// No scaling, texture size will match the tensor's width/height dimensions.
    Original,

    /// Scale the texture for the largest possible fit in the UI container.
    Fill,
}

impl Default for TextureScaling {
    fn default() -> Self {
        Self::Fill
    }
}

impl Display for TextureScaling {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            TextureScaling::Original => "Original".fmt(f),
            TextureScaling::Fill => "Fill".fmt(f),
        }
    }
}

/// Scaling, filtering, aspect ratio, etc for the rendered texture.
#[derive(Copy, Clone, Debug, PartialEq, serde::Deserialize, serde::Serialize)]
struct TextureSettings {
    /// Should the aspect ratio of the tensor be kept when scaling?
    keep_aspect_ratio: bool,

    /// Should we scale the texture when rendering?
    scaling: TextureScaling,

    /// Specifies the sampling filter used to render the texture.
    options: egui::TextureOptions,
}

impl Default for TextureSettings {
    fn default() -> Self {
        Self {
            keep_aspect_ratio: true,
            scaling: TextureScaling::default(),
            options: egui::TextureOptions {
                // This is best for low-res depth-images and the like
                magnification: egui::TextureFilter::Nearest,
                minification: egui::TextureFilter::Linear,
            },
        }
    }
}

// ui
impl TextureSettings {
    fn ui(&mut self, re_ui: &re_ui::ReUi, ui: &mut egui::Ui) {
        let TextureSettings {
            keep_aspect_ratio,
            scaling,
            options,
        } = self;

        re_ui.grid_left_hand_label(ui, "Scale");
        ui.vertical(|ui| {
            egui::ComboBox::from_id_source("texture_scaling")
                .selected_text(scaling.to_string())
                .show_ui(ui, |ui| {
                    ui.style_mut().wrap = Some(false);
                    ui.set_min_width(64.0);

                    let mut selectable_value =
                        |ui: &mut egui::Ui, e| ui.selectable_value(scaling, e, e.to_string());
                    selectable_value(ui, TextureScaling::Original);
                    selectable_value(ui, TextureScaling::Fill);
                });
            if *scaling == TextureScaling::Fill {
                ui.checkbox(keep_aspect_ratio, "Keep aspect ratio");
            }
        });
        ui.end_row();

        // TODO(#1612): support texture filtering again
        if false {
            re_ui
                .grid_left_hand_label(ui, "Filtering")
                .on_hover_text("Filtering to use when magnifying");

            fn tf_to_string(tf: egui::TextureFilter) -> &'static str {
                match tf {
                    egui::TextureFilter::Nearest => "Nearest",
                    egui::TextureFilter::Linear => "Linear",
                }
            }
            egui::ComboBox::from_id_source("texture_filter")
                .selected_text(tf_to_string(options.magnification))
                .show_ui(ui, |ui| {
                    ui.style_mut().wrap = Some(false);
                    ui.set_min_width(64.0);

                    let mut selectable_value = |ui: &mut egui::Ui, e| {
                        ui.selectable_value(&mut options.magnification, e, tf_to_string(e))
                    };
                    selectable_value(ui, egui::TextureFilter::Linear);
                    selectable_value(ui, egui::TextureFilter::Nearest);
                });
            ui.end_row();
        }
    }
}

// ----------------------------------------------------------------------------

pub fn selected_tensor_slice<'a, T: Copy>(
    slice_selection: &SliceSelection,
    tensor: &'a ndarray::ArrayViewD<'_, T>,
) -> ndarray::ArrayViewD<'a, T> {
    let SliceSelection {
        dim_mapping: dimension_mapping,
        selector_values,
    } = slice_selection;

    assert!(dimension_mapping.is_valid(tensor.ndim()));

    // TODO(andreas) - shouldn't just give up here
    if dimension_mapping.width.is_none() || dimension_mapping.height.is_none() {
        return tensor.view();
    }

    let axis = dimension_mapping
        .height
        .into_iter()
        .chain(dimension_mapping.width.into_iter())
        .chain(dimension_mapping.selectors.iter().map(|s| s.dim_idx))
        .collect::<Vec<_>>();
    let mut slice = tensor.view().permuted_axes(axis);

    for DimensionSelector { dim_idx, .. } in &dimension_mapping.selectors {
        let selector_value = selector_values.get(dim_idx).copied().unwrap_or_default() as usize;
        assert!(
            selector_value < slice.shape()[2],
            "Bad tensor slicing. Trying to select slice index {selector_value} of dim=2. tensor shape: {:?}, dim_mapping: {dimension_mapping:#?}",
            tensor.shape()
        );

        // 0 and 1 are width/height, the rest are rearranged by dimension_mapping.selectors
        // This call removes Axis(2), so the next iteration of the loop does the right thing again.
        slice.index_axis_inplace(Axis(2), selector_value);
    }
    if dimension_mapping.invert_height {
        slice.invert_axis(Axis(0));
    }
    if dimension_mapping.invert_width {
        slice.invert_axis(Axis(1));
    }

    slice
}

fn dimension_name(shape: &[component_types::TensorDimension], dim_idx: usize) -> String {
    let dim = &shape[dim_idx];
    dim.name.as_ref().map_or_else(
        || format!("Dimension {dim_idx} (size={})", dim.size),
        |name| format!("{name} (size={})", dim.size),
    )
}

fn paint_axis_names(
    ui: &mut egui::Ui,
    painter: &egui::Painter,
    rect: egui::Rect,
    font_id: egui::FontId,
    dimension_labels: [(String, bool); 2],
) {
    // Show axis names etc:
    let [(width_name, invert_width), (height_name, invert_height)] = dimension_labels;
    let text_color = ui.visuals().text_color();

    let rounding = re_ui::ReUi::normal_rounding();
    let inner_margin = rounding;
    let outer_margin = 8.0;

    let rect = rect.shrink(outer_margin + inner_margin);

    // We make sure that the label for the X axis is always at Y=0,
    // and that the label for the Y axis is always at X=0, no matter what inversions.
    //
    // For instance, with origin in the top right:
    //
    // foo ⬅
    // ..........
    // ..........
    // ..........
    // .......... ↓
    // .......... b
    // .......... a
    // .......... r

    // TODO(emilk): draw actual arrows behind the text instead of the ugly emoji arrows

    let paint_text_bg = |text_background, text_rect: egui::Rect| {
        painter.set(
            text_background,
            egui::Shape::rect_filled(
                text_rect.expand(inner_margin),
                rounding,
                ui.visuals().panel_fill,
            ),
        );
    };

    // Label for X axis:
    {
        let text_background = painter.add(egui::Shape::Noop);
        let text_rect = if invert_width {
            // On left, pointing left:
            let (pos, align) = if invert_height {
                (rect.left_bottom(), Align2::LEFT_BOTTOM)
            } else {
                (rect.left_top(), Align2::LEFT_TOP)
            };
            painter.text(
                pos,
                align,
                format!("{width_name} ⬅"),
                font_id.clone(),
                text_color,
            )
        } else {
            // On right, pointing right:
            let (pos, align) = if invert_height {
                (rect.right_bottom(), Align2::RIGHT_BOTTOM)
            } else {
                (rect.right_top(), Align2::RIGHT_TOP)
            };
            painter.text(
                pos,
                align,
                format!("➡ {width_name}"),
                font_id.clone(),
                text_color,
            )
        };
        paint_text_bg(text_background, text_rect);
    }

    // Label for Y axis:
    {
        let text_background = painter.add(egui::Shape::Noop);
        let text_rect = if invert_height {
            // On top, pointing up:
            let galley = painter.layout_no_wrap(format!("➡ {height_name}"), font_id, text_color);
            let galley_size = galley.size();
            let pos = if invert_width {
                rect.right_top() + egui::vec2(-galley_size.y, galley_size.x)
            } else {
                rect.left_top() + egui::vec2(0.0, galley_size.x)
            };
            painter.add(TextShape {
                pos,
                galley,
                angle: -std::f32::consts::TAU / 4.0,
                underline: Default::default(),
                override_text_color: None,
            });
            egui::Rect::from_min_size(
                pos - galley_size.x * egui::Vec2::Y,
                egui::vec2(galley_size.y, galley_size.x),
            )
        } else {
            // On bottom, pointing down:
            let galley = painter.layout_no_wrap(format!("{height_name} ⬅"), font_id, text_color);
            let galley_size = galley.size();
            let pos = if invert_width {
                rect.right_bottom() - egui::vec2(galley_size.y, 0.0)
            } else {
                rect.left_bottom()
            };
            painter.add(TextShape {
                pos,
                galley,
                angle: -std::f32::consts::TAU / 4.0,
                underline: Default::default(),
                override_text_color: None,
            });
            egui::Rect::from_min_size(
                pos - galley_size.x * egui::Vec2::Y,
                egui::vec2(galley_size.y, galley_size.x),
            )
        };
        paint_text_bg(text_background, text_rect);
    }
}

fn selectors_ui(ui: &mut egui::Ui, state: &mut ViewTensorState, tensor: &Tensor) {
    for selector in &state.slice.dim_mapping.selectors {
        if !selector.visible {
            continue;
        }

        let dim = &tensor.shape()[selector.dim_idx];
        let size = dim.size;

        let selector_value = state
            .slice
            .selector_values
            .entry(selector.dim_idx)
            .or_insert_with(|| size / 2); // start in the middle

        if size > 0 {
            *selector_value = selector_value.at_most(size - 1);
        }

        if size > 1 {
            ui.horizontal(|ui| {
                let name = dim
                    .name
                    .clone()
                    .unwrap_or_else(|| selector.dim_idx.to_string());

                let slider_tooltip = format!("Adjust the selected slice for the {name} dimension");
                ui.label(&name).on_hover_text(&slider_tooltip);

                // If the range is big (say, 2048) then we would need
                // a slider that is 2048 pixels wide to get the good precision.
                // So we add a high-precision drag-value instead:
                ui.add(
                    egui::DragValue::new(selector_value)
                        .clamp_range(0..=size - 1)
                        .speed(0.5),
                )
                .on_hover_text(format!(
                    "Drag to precisely control the slice index of the {name} dimension"
                ));

                // Make the slider as big as needed:
                const MIN_SLIDER_WIDTH: f32 = 64.0;
                if ui.available_width() >= MIN_SLIDER_WIDTH {
                    ui.spacing_mut().slider_width = (size as f32 * 4.0)
                        .at_least(MIN_SLIDER_WIDTH)
                        .at_most(ui.available_width());
                    ui.add(egui::Slider::new(selector_value, 0..=size - 1).show_value(false))
                        .on_hover_text(slider_tooltip);
                }
            });
        }
    }
}