egui_plot 0.35.0

use std::string::String;

use egui::Align;
use egui::Color32;
use egui::Direction;
use egui::Frame;
use egui::Id;
use egui::Layout;
use egui::PointerButton;
use egui::Rect;
use egui::Response;
use egui::Sense;
use egui::Shadow;
use egui::Shape;
use egui::TextStyle;
use egui::Ui;
use egui::Widget;
use egui::WidgetInfo;
use egui::WidgetType;
use egui::epaint::CircleShape;
use egui::pos2;
use egui::vec2;

use crate::items::PlotItem;
use crate::placement::Corner;

/// How to handle multiple conflicting color for a legend item.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub enum ColorConflictHandling {
    PickFirst,
    PickLast,
    RemoveColor,
}

/// How to group legend entries.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub enum LegendGrouping {
    /// Items with the same name share a single legend entry (default).
    #[default]
    ByName,

    /// Each item gets its own legend entry, keyed by its unique [`Id`].
    ById,
}

/// The configuration for a plot legend.
#[derive(Clone, PartialEq)]
#[cfg_attr(feature = "serde", derive(serde::Deserialize, serde::Serialize))]
pub struct Legend {
    pub text_style: TextStyle,
    pub background_alpha: f32,
    pub position: Corner,
    pub title: Option<String>,

    follow_insertion_order: bool,
    grouping: LegendGrouping,
    color_conflict_handling: ColorConflictHandling,

    /// Used for overriding the `hidden_items` set in [`LegendWidget`].
    hidden_items: Option<ahash::HashSet<Id>>,
}

impl Default for Legend {
    fn default() -> Self {
        Self {
            text_style: TextStyle::Body,
            background_alpha: 0.75,
            position: Corner::RightTop,
            title: None,
            follow_insertion_order: false,
            grouping: LegendGrouping::default(),
            color_conflict_handling: ColorConflictHandling::RemoveColor,
            hidden_items: None,
        }
    }
}

impl Legend {
    /// Which text style to use for the legend. Default: `TextStyle::Body`.
    #[inline]
    pub fn text_style(mut self, style: TextStyle) -> Self {
        self.text_style = style;
        self
    }

    /// The alpha of the legend background. Default: `0.75`.
    #[inline]
    pub fn background_alpha(mut self, alpha: f32) -> Self {
        self.background_alpha = alpha;
        self
    }

    /// In which corner to place the legend. Default: `Corner::RightTop`.
    #[inline]
    pub fn position(mut self, corner: Corner) -> Self {
        self.position = corner;
        self
    }

    /// Set the title of the legend. Default: `None`.
    #[inline]
    pub fn title(mut self, title: &str) -> Self {
        self.title = Some(title.to_owned());
        self
    }

    /// Specifies hidden items in the legend configuration to override the
    /// existing ones. This allows the legend traces' visibility to be
    /// controlled from the application code.
    #[inline]
    pub fn hidden_items<I>(mut self, hidden_items: I) -> Self
    where
        I: IntoIterator<Item = Id>,
    {
        self.hidden_items = Some(hidden_items.into_iter().collect());
        self
    }

    /// Specifies if the legend item order should be the inserted order.
    /// Default: `false`.
    /// If `true`, the order of the legend items will be the same as the order
    /// as they were added. By default it will be sorted alphabetically.
    #[inline]
    pub fn follow_insertion_order(mut self, follow: bool) -> Self {
        self.follow_insertion_order = follow;
        self
    }

    /// Specifies how to handle conflicting colors for an item.
    #[inline]
    pub fn color_conflict_handling(mut self, color_conflict_handling: ColorConflictHandling) -> Self {
        self.color_conflict_handling = color_conflict_handling;
        self
    }

    /// Specifies how legend entries are grouped. Default: [`LegendGrouping::ByName`].
    ///
    /// With [`LegendGrouping::ByName`], items sharing the same name are
    /// merged into a single legend entry. With [`LegendGrouping::ById`],
    /// each item gets its own entry keyed by its unique [`Id`].
    #[inline]
    pub fn grouping(mut self, grouping: LegendGrouping) -> Self {
        self.grouping = grouping;
        self
    }
}

#[derive(Clone)]
struct LegendEntry {
    id: Id,
    name: String,
    color: Color32,
    checked: bool,
    hovered: bool,
}

impl LegendEntry {
    fn new(id: Id, name: String, color: Color32, checked: bool) -> Self {
        Self {
            id,
            name,
            color,
            checked,
            hovered: false,
        }
    }

    fn ui(&self, ui: &mut Ui, text_style: &TextStyle) -> Response {
        let Self {
            id: _,
            name,
            color,
            checked,
            hovered: _,
        } = self;

        let font_id = text_style.resolve(ui.style());

        let galley = ui.fonts_mut(|f| f.layout_delayed_color(name.clone(), font_id, f32::INFINITY));

        let icon_size = galley.size().y;
        let icon_spacing = icon_size / 5.0;
        let total_extra = vec2(icon_size + icon_spacing, 0.0);

        let desired_size = total_extra + galley.size();
        let (rect, response) = ui.allocate_exact_size(desired_size, Sense::click());

        response.widget_info(|| WidgetInfo::selected(WidgetType::Checkbox, ui.is_enabled(), *checked, galley.text()));

        let visuals = ui.style().interact(&response);
        let label_on_the_left = ui.layout().horizontal_placement() == Align::RIGHT;

        let icon_position_x = if label_on_the_left {
            rect.right() - icon_size / 2.0
        } else {
            rect.left() + icon_size / 2.0
        };
        let icon_position = pos2(icon_position_x, rect.center().y);
        let icon_rect = Rect::from_center_size(icon_position, vec2(icon_size, icon_size));

        let painter = ui.painter();

        // Gray background, for interaction effects, and to sow something if we're
        // disabled:
        painter.add(CircleShape {
            center: icon_rect.center(),
            radius: icon_size * 0.35,
            fill: visuals.bg_fill,
            stroke: visuals.bg_stroke,
        });

        if *checked {
            let fill = if *color == Color32::TRANSPARENT {
                ui.visuals().noninteractive().fg_stroke.color
            } else {
                *color
            };
            painter.add(Shape::circle_filled(icon_rect.center(), icon_size * 0.25, fill));
        }

        let text_position_x = if label_on_the_left {
            rect.right() - icon_size - icon_spacing - galley.size().x
        } else {
            rect.left() + icon_size + icon_spacing
        };

        let text_position = pos2(text_position_x, rect.center().y - 0.5 * galley.size().y);
        painter.galley(text_position, galley, visuals.text_color());

        response
    }
}

#[derive(Clone)]
pub struct LegendWidget {
    rect: Rect,
    entries: Vec<LegendEntry>,
    config: Legend,
}

impl LegendWidget {
    /// Create a new legend from items, the names of items that are hidden and
    /// the style of the text. Returns `None` if the legend has no entries.
    pub(crate) fn try_new<'a>(
        rect: Rect,
        config: Legend,
        items: &[Box<dyn PlotItem + 'a>],
        hidden_items: &ahash::HashSet<Id>, // Existing hidden items in the plot memory.
    ) -> Option<Self> {
        // If `config.hidden_items` is not `None`, it is used.
        let hidden_items = config.hidden_items.as_ref().unwrap_or(hidden_items);

        // Collect the legend entries. With `ByName` grouping, items sharing the
        // same name are merged into a single checkbox. With `ById` grouping,
        // items sharing the same `Id` are merged instead. When colors conflict
        // within a merged entry, `color_conflict_handling` decides which color
        // to show.
        let mut entries: Vec<LegendEntry> = Vec::new();
        let mut seen: ahash::HashMap<Id, usize> = ahash::HashMap::default();
        for item in items.iter().filter(|item| !item.name().is_empty()) {
            let dedup_key = match config.grouping {
                LegendGrouping::ByName => Id::new(item.name()),
                LegendGrouping::ById => item.id(),
            };

            if let Some(&idx) = seen.get(&dedup_key) {
                let entry = &mut entries[idx];
                if entry.color != item.color() {
                    match config.color_conflict_handling {
                        ColorConflictHandling::PickFirst => (),
                        ColorConflictHandling::PickLast => entry.color = item.color(),
                        ColorConflictHandling::RemoveColor => {
                            entry.color = Color32::TRANSPARENT;
                        }
                    }
                }
            } else {
                seen.insert(dedup_key, entries.len());
                let color = item.color();
                let checked = !hidden_items.contains(&item.id());
                entries.push(LegendEntry::new(item.id(), item.name().to_owned(), color, checked));
            }
        }

        if !config.follow_insertion_order {
            entries.sort_by(|a, b| a.name.cmp(&b.name));
        }

        (!entries.is_empty()).then_some(Self { rect, entries, config })
    }

    // Get the names of the hidden items.
    pub fn hidden_items(&self) -> ahash::HashSet<Id> {
        self.entries
            .iter()
            .filter_map(|entry| (!entry.checked).then_some(entry.id))
            .collect()
    }

    // Get the name of the hovered items.
    pub fn hovered_item(&self) -> Option<Id> {
        self.entries.iter().find_map(|entry| entry.hovered.then_some(entry.id))
    }
}

impl Widget for &mut LegendWidget {
    fn ui(self, ui: &mut Ui) -> Response {
        let LegendWidget { rect, entries, config } = self;

        let main_dir = match config.position {
            Corner::LeftTop | Corner::RightTop => Direction::TopDown,
            Corner::LeftBottom | Corner::RightBottom => Direction::BottomUp,
        };
        let cross_align = match config.position {
            Corner::LeftTop | Corner::LeftBottom => Align::LEFT,
            Corner::RightTop | Corner::RightBottom => Align::RIGHT,
        };
        let layout = Layout::from_main_dir_and_cross_align(main_dir, cross_align);
        let legend_pad = 4.0;
        let legend_rect = rect.shrink(legend_pad);
        let mut legend_ui = ui.new_child(egui::UiBuilder::new().max_rect(legend_rect).layout(layout));
        legend_ui
            .scope(|ui| {
                let background_frame = Frame {
                    inner_margin: vec2(8.0, 4.0).into(),
                    corner_radius: ui.style().visuals.window_corner_radius,
                    shadow: Shadow::NONE,
                    fill: ui.style().visuals.extreme_bg_color,
                    stroke: ui.style().visuals.window_stroke(),
                    ..Default::default()
                }
                .multiply_with_opacity(config.background_alpha);
                background_frame
                    .show(ui, |ui| {
                        // always show on top of the legend - so we need to use a new scope
                        if main_dir == Direction::TopDown
                            && let Some(title) = &config.title
                        {
                            ui.heading(title);
                        }
                        let mut focus_on_item = None;

                        #[expect(
                            clippy::expect_used,
                            reason = "we checked that entries is not empty when creating the legend"
                        )]
                        let response_union = entries
                            .iter_mut()
                            .map(|entry| {
                                let response = entry.ui(ui, &config.text_style);

                                // Handle interactions. Alt-clicking must be deferred to end of loop
                                // since it may affect all entries.
                                handle_interaction_on_legend_item(&response, entry);
                                if response.clicked() && ui.input(|r| r.modifiers.alt) {
                                    focus_on_item = Some(entry.id);
                                }

                                response
                            })
                            .reduce(|r1, r2| r1.union(r2))
                            .expect("No entries in the legend");

                        if main_dir == Direction::BottomUp
                            && let Some(title) = &config.title
                        {
                            ui.heading(title);
                        }

                        if let Some(focus_on_item) = focus_on_item {
                            handle_focus_on_legend_item(&focus_on_item, entries);
                        }

                        response_union
                    })
                    .inner
            })
            .inner
    }
}

/// Handle per-entry interactions.
fn handle_interaction_on_legend_item(response: &Response, entry: &mut LegendEntry) {
    entry.checked ^= response.clicked_by(PointerButton::Primary);
    entry.hovered = response.hovered();
}

/// Handle alt-click interaction (which may affect all entries).
fn handle_focus_on_legend_item(clicked_entry: &Id, entries: &mut [LegendEntry]) {
    // if all other items are already hidden, we show everything
    let is_focus_item_only_visible = entries
        .iter()
        .all(|entry| !entry.checked || (clicked_entry == &entry.id));

    // either show everything or show only the focus item
    for entry in entries.iter_mut() {
        entry.checked = is_focus_item_only_visible || clicked_entry == &entry.id;
    }
}