henad-app 0.3.0

//! The Model panel: model selection, and the declarations of the selected model.

use henad_compute::entry::ModelEntry;
use henad_compute::gpu::capacity::Demand;
use henad_core::helpers::fmt_bytes;
use henad_core::metadata::{LaneSpec, Structure};
use henad_core::params::ParamDescriptor;
use henad_core::topology::{NeighborhoodKind, TopologyHint};

use crate::icons::material_design_icons::{MDI_CHECK, MDI_CLOSE};
use crate::state::{AppState, OpeningRefusal};
use crate::ui::{KvGridRows, kv_grid};

/// Text the Model panel shows when no model of the set runs on this machine.
const NO_MODEL_RUNS: &str = "No model in this build runs on this device. GPU models need a GPU with compute support.";

/// Text the Model panel shows for a set without models.
const NO_MODELS: &str = "This build includes no models.";

/// Draws the opening that the app rejected and the reason, then `next_step`.
fn opening_refusal(ui: &mut egui::Ui, refusal: &OpeningRefusal, next_step: &str) {
    ui.colored_label(ui.visuals().warn_fg_color, refusal.lead);
    ui.label(format!("{}.", refusal.reason));
    ui.label(next_step);
}

pub fn model_ui(ui: &mut egui::Ui, app: &mut AppState) {
    let offered: Vec<(String, String)> = app
        .offered_models()
        .map(|entry| (entry.id().to_owned(), entry.name().to_owned()))
        .collect();
    if offered.is_empty() {
        // The refusal of a hidden GPU model already mentions the GPU, and an empty set has no GPU to blame.
        match &app.opening_refusal {
            Some(refusal) => opening_refusal(ui, refusal, "No model in this build runs on this device."),
            None if app.models.is_empty() => {
                ui.label(NO_MODELS);
            }
            None => {
                ui.label(NO_MODEL_RUNS);
            }
        }
        return;
    }
    let selected_name = app.selected_entry().map_or("None", ModelEntry::name).to_owned();
    let mut picked = None;

    egui::ComboBox::from_label("Select Model")
        .selected_text(selected_name)
        // As tall as the window allows. The model list outgrows egui's default of 200 points.
        .height(ui.ctx().content_rect().height())
        .show_ui(ui, |ui| {
            for (id, name) in &offered {
                let selected = app.selected_model.as_deref() == Some(id.as_str());
                if ui.selectable_label(selected, name).clicked() && !selected {
                    picked = Some(id.clone());
                }
            }
        });

    let changed_model = picked.is_some();
    if let Some(id) = picked {
        app.select_model(&id);
    }

    let Some(entry) = app.selected_entry() else {
        if let Some(refusal) = &app.opening_refusal {
            ui.separator();
            opening_refusal(ui, refusal, "Select a model to continue.");
        }
        return;
    };

    ui.separator();
    ui.label(entry.description());
    ui.separator();

    // Recomputed as the sliders move, so the footprint tracks the panel next door.
    let demand = entry.demand(&app.param_values, &app.runtime.granted);
    let storage_limit = app.runtime.granted.max_storage_buffers_per_shader_stage;

    let mut scroll = egui::ScrollArea::vertical();
    if changed_model {
        // Models differ in how many rows they declare, so a carried-over offset can open the
        // next model halfway down.
        scroll = scroll.vertical_scroll_offset(0.0);
    }

    scroll.show(ui, |ui| {
        section(ui, "Identity");
        kv_grid(ui, "model_identity_grid").show(ui, |ui, rows| {
            row(ui, rows, "Id", entry.id());
            row(ui, rows, "Backend", entry.metadata().backend.label());
            row(ui, rows, "Topology", topology_label(entry.topology_hint()));
        });

        ui.add_space(8.0);
        section(ui, "Structure");
        kv_grid(ui, "model_structure_grid").show(ui, |ui, rows| structure_rows(ui, rows, &entry.metadata().structure));

        ui.add_space(8.0);
        section(ui, "Interface");
        kv_grid(ui, "model_interface_grid").show(ui, |ui, rows| interface_rows(ui, rows, entry));

        // A CPU model allocates on the host, and the Performance tab already reports that.
        if let Some(demand) = &demand {
            ui.add_space(8.0);
            section(ui, "Footprint");
            kv_grid(ui, "model_footprint_grid").show(ui, |ui, rows| {
                footprint_rows(ui, rows, entry.id(), demand, storage_limit);
            });
        }
    });
}

fn section(ui: &mut egui::Ui, title: &str) {
    ui.strong(title);
    ui.add_space(2.0);
}

fn row(ui: &mut egui::Ui, rows: &mut KvGridRows, label: &str, value: impl Into<egui::WidgetText>) {
    ui.label(label);
    ui.label(value);
    rows.end_row(ui);
}

/// Adds a row with the count `value`, whose two cells both list `names` on hover.
///
/// A value as short as "7" is barely a pointer wide on its own.
fn row_listing(ui: &mut egui::Ui, rows: &mut KvGridRows, label: &str, value: String, names: &str) {
    ui.label(label).on_hover_text(names);
    ui.label(value).on_hover_text(names);
    rows.end_row(ui);
}

fn topology_label(hint: TopologyHint) -> &'static str {
    match (hint.grid, hint.agents, hint.edges) {
        (false, true, true) => "Network",
        (true, true, true) => "Network and Grid",
        (true, true, false) => "Agents and Grid",
        (true, false, false) => "Grid",
        (false, true, false) => "Agents",
        _ => "None",
    }
}

fn neighborhood_label(kind: NeighborhoodKind) -> &'static str {
    match kind {
        NeighborhoodKind::Moore => "Moore",
        NeighborhoodKind::VonNeumann => "Von Neumann",
    }
}

fn yes_no(present: bool) -> &'static str {
    if present { MDI_CHECK } else { MDI_CLOSE }
}

/// Returns the count `n`, followed by the count `flagged` labelled `flag` unless `flagged` is zero, as in
/// "7 (2 reload)".
fn count_of(n: usize, flagged: usize, flag: &str) -> String {
    if flagged == 0 {
        n.to_string()
    } else {
        format!("{n} ({flagged} {flag})")
    }
}

/// Shows the lane rows, shared by agent and network models.
fn lane_rows(ui: &mut egui::Ui, rows: &mut KvGridRows, lanes: &[LaneSpec]) {
    let doubled = lanes.iter().filter(|lane| lane.double_buffered).count();
    let names: Vec<String> = lanes
        .iter()
        .map(|lane| {
            let mark = if lane.double_buffered { " (dual)" } else { "" };
            format!("{}: {}{mark}", lane.name, lane.ty)
        })
        .collect();
    row_listing(
        ui,
        rows,
        "Lanes",
        count_of(lanes.len(), doubled, "double-buffered"),
        &names.join("\n"),
    );
}

fn structure_rows(ui: &mut egui::Ui, rows: &mut KvGridRows, structure: &Structure) {
    match structure {
        Structure::Grid { neighborhood } => {
            row(ui, rows, "Neighbourhood", neighborhood_label(*neighborhood));
        }
        Structure::Agents {
            chunk,
            lanes,
            index,
            field,
        } => {
            lane_rows(ui, rows, lanes);
            row(ui, rows, "Chunk size", format!("{chunk} agents"));
            row(ui, rows, "Neighbour index", *index);
            row(ui, rows, "Field layer", *field);
        }
        Structure::Network { chunk, lanes, .. } => {
            lane_rows(ui, rows, lanes);
            row(ui, rows, "Chunk size", format!("{chunk} nodes"));
        }
        Structure::GpuGrid { buffers, workgroup } => {
            row_listing(ui, rows, "Buffers", buffers.len().to_string(), &buffers.join("\n"));
            row(ui, rows, "Workgroup", format!("{workgroup} × {workgroup}"));
        }
        Structure::GpuAgents {
            buffers,
            passes,
            index,
            display,
            counters,
        } => {
            let doubled = buffers.iter().filter(|spec| spec.double_buffered).count();
            let buffer_names: Vec<&str> = buffers.iter().map(|spec| spec.label).collect();
            let pass_names: Vec<&str> = passes.iter().map(|pass| pass.label).collect();
            row_listing(
                ui,
                rows,
                "Buffers",
                count_of(buffers.len(), doubled, "double-buffered"),
                &buffer_names.join("\n"),
            );
            row_listing(
                ui,
                rows,
                "Step passes",
                passes.len().to_string(),
                &pass_names.join("\n"),
            );
            row(
                ui,
                rows,
                "Neighbour index",
                if *index { "Spatial hash" } else { "None" },
            );
            row(ui, rows, "Display pass", yes_no(*display));
            if *counters > 0 {
                row(ui, rows, "Counters", counters.to_string());
            }
        }
    }
}

fn interface_rows(ui: &mut egui::Ui, rows: &mut KvGridRows, entry: &ModelEntry) {
    let descs: &[ParamDescriptor] = entry.param_descriptors();
    let reload = descs.iter().filter(|desc| !desc.is_live()).count();
    row(ui, rows, "Parameters", count_of(descs.len(), reload, "reload"));

    ui.label("Statistics");
    ui.horizontal(|ui| {
        ui.label(entry.stat_descriptors().len().to_string());
        swatches(ui, entry.stat_descriptors().iter().map(|stat| stat.color));
    });
    rows.end_row(ui);

    ui.label("Palette");
    match entry.metadata().palette {
        Some(palette) => {
            ui.horizontal(|ui| {
                ui.label(palette.len().to_string());
                swatches(ui, palette.iter().copied());
            });
        }
        // A GPU agent model's shaders write RGBA directly and declare no palette.
        None => {
            ui.label("In shader").on_hover_text(
                "Palette colours are declared directly in the shaders and are not retrievable directly.",
            );
        }
    }
    rows.end_row(ui);

    if let Structure::Network { edge_palette, .. } = &entry.metadata().structure {
        ui.label("Edge palette");
        ui.horizontal(|ui| {
            ui.label(edge_palette.len().to_string());
            swatches(ui, edge_palette.iter().copied());
        });
        rows.end_row(ui);
    }
}

/// Adds the rows of the model's expected footprint on the GPU.
fn footprint_rows(ui: &mut egui::Ui, rows: &mut KvGridRows, id: &str, demand: &Demand, storage_limit: u32) {
    row(ui, rows, "Expected memory", fmt_bytes(demand.bytes()));

    if let Some(largest) = demand.buffers.iter().max_by_key(|alloc| alloc.bytes) {
        row_listing(
            ui,
            rows,
            "Largest buffer",
            format!("{}, {}", strip_id(&largest.label, id), fmt_bytes(largest.bytes)),
            &largest.label,
        );
    }

    if let Some((w, h)) = demand.texture {
        row(ui, rows, "Display texture", format!("{w} × {h}"));
    }

    if let Some(widest) = demand.passes.iter().max_by_key(|pass| pass.storage) {
        let hint = format!("Widest pass is '{}'", widest.label);
        ui.label("Storage bindings").on_hover_text(&hint);
        ui.horizontal(|ui| {
            ui.style_mut().spacing.item_spacing.x = 0.0;
            ui.label(widest.storage.to_string()).on_hover_text(&hint);
            ui.weak(format!(" of {storage_limit}")).on_hover_text(&hint);
        });
        rows.end_row(ui);
    }
}

/// Returns `label` without its `{id}_` prefix. The panel already shows the model's name.
fn strip_id<'a>(label: &'a str, id: &str) -> &'a str {
    label
        .strip_prefix(id)
        .and_then(|rest| rest.strip_prefix('_'))
        .unwrap_or(label)
}

/// Draws a row of colour chips.
fn swatches(ui: &mut egui::Ui, colors: impl Iterator<Item = [u8; 4]>) {
    ui.horizontal(|ui| {
        ui.spacing_mut().item_spacing.x = 2.0;
        let size = egui::vec2(9.0, 9.0);
        for [r, g, b, a] in colors {
            let (rect, _) = ui.allocate_exact_size(size, egui::Sense::hover());
            ui.painter().rect_filled(
                rect,
                egui::CornerRadius::ZERO,
                egui::Color32::from_rgba_unmultiplied(r, g, b, a),
            );
        }
    });
}

#[cfg(test)]
mod tests {
    use super::{count_of, strip_id, topology_label};
    use henad_core::topology::TopologyHint;

    #[test]
    fn topology_names_every_layer_combination() {
        assert_eq!(topology_label(TopologyHint::GRID), "Grid");
        assert_eq!(topology_label(TopologyHint::AGENTS), "Agents");
        assert_eq!(topology_label(TopologyHint::COMPOSITE), "Agents and Grid");
        assert_eq!(topology_label(TopologyHint::NETWORK), "Network");
        assert_eq!(topology_label(TopologyHint::NONE), "None");
    }

    /// A model with nothing flagged shows as a bare count, without "(0 reload)".
    #[test]
    fn a_count_drops_its_qualifier_when_nothing_carries_it() {
        assert_eq!(count_of(5, 0, "reload"), "5");
        assert_eq!(count_of(5, 2, "reload"), "5 (2 reload)");
    }

    /// A buffer whose label does not start with the id keeps its whole label.
    #[test]
    fn a_buffer_label_loses_only_the_model_id() {
        assert_eq!(strip_id("gpu_ants_accum_a", "gpu_ants"), "accum_a");
        assert_eq!(strip_id("gpu_ants_hash_counts", "gpu_ants"), "hash_counts");
        assert_eq!(strip_id("scratch", "gpu_ants"), "scratch");
        assert_eq!(strip_id("gpu_antsy", "gpu_ants"), "gpu_antsy");
    }
}