henad-app 0.3.0

use std::sync::Arc;

use crate::state::PointRenderMode;
use crate::ui::agent_layer::{AgentLayer, padded_palette};
use crate::ui::painted::{Painted, painted};
use crate::{icons::material_design_icons::MDI_CUBE_OFF_OUTLINE, state::AppState};
use eframe::egui_wgpu;
use egui::{ColorImage, RichText, TextureOptions};
use egui_wgpu::{CallbackResources, CallbackTrait};
use henad_compute::display_scale::{display_dims, source_row};
use henad_compute::gpu::GpuDisplay;
use henad_compute::snapshot::{CpuLayers, EdgeSnapshot, GpuSnapshot, GridSnapshot, PointSnapshot, SnapshotView};

/// Paints a GPU model's display texture straight into the viewport.
///
/// A GPU model's cells never reach the CPU, so there is no `ColorImage` to upload. The sim thread
/// has already rendered the grid into a texture, and all that is left is to sample it.
///
/// The callback carries its own `Arc<GpuDisplay>` instead of a lookup in egui's type-keyed `CallbackResources`. A
/// model switch mid-frame then leaves the pipeline and texture alive until the render pass is done with them.
struct GpuViewportPaint {
    display: Painted<Arc<GpuDisplay>>,
}

impl CallbackTrait for GpuViewportPaint {
    fn paint(
        &self,
        _info: egui::PaintCallbackInfo,
        render_pass: &mut wgpu::RenderPass<'static>,
        _callback_resources: &CallbackResources,
    ) {
        render_pass.set_pipeline(&self.display.render_pipeline);
        render_pass.set_bind_group(0, &self.display.render_bind_group, &[]);
        render_pass.draw(0..3, 0..1);
    }
}

/// Composites whichever GPU layers the model published into one rect, field first and agents over
/// the top, in the order the CPU path draws them.
///
/// Takes `app` because the agent pipeline is built lazily on first use.
fn paint_gpu_view(ui: &mut egui::Ui, app: &mut AppState, gpu: &GpuSnapshot) {
    // A field, when present, fixes the pixel shape, as `layer_extent` does for CPU models.
    let extent = gpu.display.as_ref().map_or_else(
        || gpu.agents.as_ref().map(|a| (a.world_w, a.world_h)),
        |d| Some((d.width as f32, d.height as f32)),
    );
    let Some((world_w, world_h)) = extent else {
        return;
    };
    if gpu.agents.is_some() {
        ensure_agent_layer(app);
    }

    let size = fit_aspect(ui.available_size(), world_w, world_h);
    ui.centered_and_justified(|ui| {
        let (rect, _response) = ui.allocate_exact_size(size, egui::Sense::hover());

        if let Some(display) = &gpu.display {
            ui.painter().add(egui_wgpu::Callback::new_paint_callback(
                rect,
                GpuViewportPaint {
                    display: painted(Arc::clone(display)),
                },
            ));
        }
        if let (Some(agents), Some(layer)) = (&gpu.agents, &app.agent_layer) {
            layer.paint_gpu(ui, rect, agents);
        }
    });
}

/// Returns the largest size with the aspect ratio of `width` by `height` that fits inside `available`.
fn fit_aspect(available: egui::Vec2, width: f32, height: f32) -> egui::Vec2 {
    let tex_aspect = width / height;
    let panel_aspect = available.x / available.y;
    if tex_aspect > panel_aspect {
        egui::Vec2::new(available.x, available.x / tex_aspect)
    } else {
        egui::Vec2::new(available.y * tex_aspect, available.y)
    }
}

/// Draws the viewport and records its cost.
pub fn viewport_ui(ui: &mut egui::Ui, app: &mut AppState) {
    let render_start = web_time::Instant::now();

    let mode_before = app.point_render_mode;
    ui.horizontal(|ui| {
        ui.checkbox(&mut app.rendering_enabled, "Rendering");

        if app.rendering_enabled && has_points(app) {
            ui.separator();
            ui.label("Agents:");
            ui.selectable_value(&mut app.point_render_mode, PointRenderMode::Agents, "Sprites");
            ui.selectable_value(&mut app.point_render_mode, PointRenderMode::Density, "Density");

            // Edges are only drawn along with sprites, and only on a GPU that can build the edge layer.
            if let Some(directed) = edges_directed(app)
                && app.point_render_mode == PointRenderMode::Agents
                && app.runtime.vertex_storage
            {
                ui.separator();
                ui.checkbox(&mut app.show_edges, "Edges");
                if app.show_edges {
                    ui.add_enabled(directed, egui::Checkbox::new(&mut app.edge_arrows, "Arrows"))
                        .on_hover_text("Show direction with arrowheads")
                        .on_disabled_hover_text("Only directed edges have arrows");
                }
            }
        }
    });
    if app.point_render_mode != mode_before {
        // The snapshot is unchanged, but the mode drawing it has changed.
        app.last_rendered_serial = None;
    }

    draw_view(ui, app);

    app.timings.frame_render_ms = render_start.elapsed().as_secs_f64() * 1000.0;
}

fn draw_view(ui: &mut egui::Ui, app: &mut AppState) {
    let ctx = ui.ctx().clone();

    if app.snapshot.is_none() {
        ui.centered_and_justified(|ui| {
            ui.heading("Load a model to start simulation");
        });
        return;
    }

    if !app.rendering_enabled {
        ui.vertical_centered(|ui| {
            ui.label(RichText::new(MDI_CUBE_OFF_OUTLINE).size(64.0));
            ui.heading("Rendering disabled");
        });
        return;
    }

    // A GPU snapshot has nothing to convert or upload. The branch reads the snapshot variant, since a GPU Game of Life
    // is still `TopologyHint::GRID`.
    //
    // The snapshot is taken out for the duration. `paint_gpu_view` needs `app` for the agent pipeline.
    if matches!(app.snapshot.as_ref().map(|s| &s.view), Some(SnapshotView::Gpu(_))) {
        let Some(snapshot) = app.snapshot.take() else {
            return;
        };
        if let SnapshotView::Gpu(gpu) = &snapshot.view {
            paint_gpu_view(ui, app, gpu);
        }
        app.snapshot = Some(snapshot);
        return;
    }

    // Keyed on the serial. A layout relaxing while the simulation is paused moves positions without advancing the tick.
    let current_serial = app.snapshot.as_ref().map_or(0, |s| s.serial);
    let needs_update = app.last_rendered_serial != Some(current_serial);

    // The snapshot is taken out for the duration, to keep it clear of borrows of the rest of `app`.
    let Some(snapshot) = app.snapshot.take() else {
        return;
    };
    // The GPU arm returned above, before any CPU-side pixel work.
    let SnapshotView::Cpu(layers) = &snapshot.view else {
        app.snapshot = Some(snapshot);
        return;
    };

    if layers.is_empty() {
        app.snapshot = Some(snapshot);
        ui.label("No view available");
        return;
    }

    if needs_update {
        if let Some(grid) = &layers.grid {
            upload_grid(&ctx, app, grid);
        }
        if let Some(points) = &layers.points {
            match app.point_render_mode {
                PointRenderMode::Agents => upload_agents(app, points, layers.edges.as_deref()),
                PointRenderMode::Density => {
                    render_density_heatmap(&ctx, app, &points.pos_x, &points.pos_y, points.world_w, points.world_h);
                }
            }
        }
        app.last_rendered_serial = Some(current_serial);
    }

    // Both layers share one rect so they stay registered. See `PointView`'s docs.
    let Some((world_w, world_h)) = layer_extent(layers) else {
        app.snapshot = Some(snapshot);
        return;
    };
    let display_size = fit_aspect(ui.available_size(), world_w, world_h);

    ui.centered_and_justified(|ui| {
        let (rect, _response) = ui.allocate_exact_size(display_size, egui::Sense::hover());

        // Painter order is the compositing order, field first and agents over the top.
        if let Some(tex) = &app.grid_texture {
            ui.painter().image(tex.id(), rect, UV_FULL, egui::Color32::WHITE);
        }
        if layers.points.is_some() {
            match app.point_render_mode {
                PointRenderMode::Agents => {
                    if let Some(layer) = &app.agent_layer {
                        layer.paint(ui, rect, world_w, world_h, app.edge_style());
                    }
                }
                PointRenderMode::Density => {
                    if let Some(tex) = &app.density_texture {
                        ui.painter().image(tex.id(), rect, UV_FULL, egui::Color32::WHITE);
                    }
                }
            }
        }
    });

    app.snapshot = Some(snapshot);
}

const UV_FULL: egui::Rect = egui::Rect::from_min_max(egui::pos2(0.0, 0.0), egui::pos2(1.0, 1.0));

/// Returns the extent the shared rect is fitted to. The agent shader takes it as the world size.
///
/// A field takes precedence when present, since it is the layer with a fixed pixel shape.
fn layer_extent(layers: &CpuLayers) -> Option<(f32, f32)> {
    if let Some(grid) = &layers.grid {
        return Some((grid.width as f32, grid.height as f32));
    }
    let points = layers.points.as_ref()?;
    Some((points.world_w, points.world_h))
}

/// Returns whether the snapshot's edges are directed, or `None` for a model without edges.
fn edges_directed(app: &AppState) -> Option<bool> {
    match app.snapshot.as_ref().map(|s| &s.view) {
        Some(SnapshotView::Cpu(layers)) => layers.edges.as_ref().map(|edges| edges.directed),
        _ => None,
    }
}

/// Returns whether the snapshot has points to draw, the one case the mode selector applies to.
fn has_points(app: &AppState) -> bool {
    matches!(
        app.snapshot.as_ref().map(|s| &s.view),
        Some(SnapshotView::Cpu(layers)) if layers.points.is_some()
    )
}

/// Returns the grid's colours at one texel per cell.
fn expand_grid(grid: &GridSnapshot) -> Vec<egui::Color32> {
    use rayon::prelude::*;
    let colors = grid_colors(grid.palette);
    grid.cells.par_iter().map(|&cell| colors[usize::from(cell)]).collect()
}

/// Returns the colour of one representative cell per texel, for a grid too large to upload whole.
fn sample_grid(grid: &GridSnapshot, tex_w: u32, tex_h: u32) -> Vec<egui::Color32> {
    let width = grid.width as usize;
    let colors = grid_colors(grid.palette);
    let row = |ty: u32| {
        let sy = source_row(ty, grid.height, tex_h) as usize;
        let cells = &grid.cells[sy * width..(sy + 1) * width];
        (0..tex_w as usize)
            .map(|tx| colors[usize::from(cells[tx * width / tex_w as usize])])
            .collect::<Vec<_>>()
    };

    use rayon::prelude::*;
    (0..tex_h).into_par_iter().flat_map_iter(row).collect()
}

fn upload_grid(ctx: &egui::Context, app: &mut AppState, grid: &GridSnapshot) {
    // A CPU grid meets the texture limit too, through egui's upload.
    let device_max = app.render_ctx.device.limits().max_texture_dimension_2d;
    let (tex_w, tex_h) = display_dims(grid.width, grid.height, device_max);

    let pixels = if (tex_w, tex_h) == (grid.width, grid.height) {
        expand_grid(grid)
    } else {
        sample_grid(grid, tex_w, tex_h)
    };

    let image = ColorImage::new([tex_w as usize, tex_h as usize], pixels);

    match &mut app.grid_texture {
        Some(tex) => tex.set(image, TextureOptions::NEAREST),
        None => app.grid_texture = Some(ctx.load_texture("grid", image, TextureOptions::NEAREST)),
    }
}

/// Builds the agent layer on first use. Both backends share it.
fn ensure_agent_layer(app: &mut AppState) {
    if app.agent_layer.is_none() {
        app.agent_layer = Some(AgentLayer::new(
            &app.render_ctx.device,
            &app.render_ctx.queue,
            app.render_ctx.target_format,
            app.runtime.vertex_storage,
        ));
    }
}

fn upload_agents(app: &mut AppState, points: &PointSnapshot, edges: Option<&EdgeSnapshot>) {
    ensure_agent_layer(app);
    if let Some(layer) = &mut app.agent_layer {
        layer.upload(points);
        layer.upload_edges(edges);
    }
}

const DENSITY_W: usize = 512;
const DENSITY_H: usize = 512;

/// Returns the colour of every cell value under `palette`, a value past its end taking the first colour.
///
/// Note that a cell value is the model's to choose. Indexing `palette` itself would panic on the UI thread for a
/// value past its end.
fn grid_colors(palette: &[[u8; 4]]) -> [egui::Color32; 256] {
    padded_palette(palette).map(|[r, g, b, a]| egui::Color32::from_rgba_unmultiplied(r, g, b, a))
}

/// Returns the colour at `t` of a five-stop piecewise linear approximation of the Inferno colour map.
#[inline]
fn inferno(t: f32) -> egui::Color32 {
    const STOPS: [[u8; 3]; 5] = [[0, 0, 4], [64, 4, 104], [183, 55, 121], [251, 136, 97], [252, 255, 164]];

    let t = t.clamp(0.0, 1.0);
    let scaled = t * 4.0;
    let seg = (scaled as usize).min(3);
    let frac = scaled - seg as f32;

    let a = STOPS[seg];
    let b = STOPS[seg + 1];

    let r = a[0] as f32 + (b[0] as f32 - a[0] as f32) * frac;
    let g = a[1] as f32 + (b[1] as f32 - a[1] as f32) * frac;
    let bl = a[2] as f32 + (b[2] as f32 - a[2] as f32) * frac;

    egui::Color32::from_rgb(r as u8, g as u8, bl as u8)
}

fn render_density_heatmap(
    ctx: &egui::Context,
    app: &mut AppState,
    pos_x: &[f32],
    pos_y: &[f32],
    world_w: f32,
    world_h: f32,
) {
    let num_pixels = DENSITY_W * DENSITY_H;
    let inv_world_w = DENSITY_W as f32 / world_w;
    let inv_world_h = DENSITY_H as f32 / world_h;

    let density = {
        use rayon::prelude::*;

        let n_threads = rayon::current_num_threads().max(1);
        let chunk_size = pos_x.len().div_ceil(n_threads);

        let partial: Vec<Vec<u32>> = pos_x
            .par_chunks(chunk_size)
            .zip(pos_y.par_chunks(chunk_size))
            .map(|(xs, ys)| {
                let mut buf = vec![0u32; num_pixels];
                for (&x, &y) in xs.iter().zip(ys.iter()) {
                    if !(x.is_finite() && y.is_finite()) {
                        continue;
                    }
                    let px = ((x * inv_world_w) as usize).min(DENSITY_W - 1);
                    let py = ((y * inv_world_h) as usize).min(DENSITY_H - 1);
                    buf[py * DENSITY_W + px] += 1;
                }
                buf
            })
            .collect();

        let mut density = vec![0u32; num_pixels];
        for buf in &partial {
            for (d, &p) in density.iter_mut().zip(buf.iter()) {
                *d += p;
            }
        }
        density
    };

    let max_density = app.density_max;
    // An empty pixel is transparent, so a field underneath still shows through.
    let pixels: Vec<egui::Color32> = density
        .iter()
        .map(|&d| {
            if d == 0 {
                egui::Color32::TRANSPARENT
            } else {
                inferno(d as f32 / max_density)
            }
        })
        .collect();
    let image = ColorImage::new([DENSITY_W, DENSITY_H], pixels);

    match &mut app.density_texture {
        Some(tex) => {
            tex.set(image, TextureOptions::LINEAR);
        }
        None => {
            app.density_texture = Some(ctx.load_texture("density_heatmap", image, TextureOptions::LINEAR));
        }
    }
}

#[cfg(test)]
mod tests {
    use henad_compute::snapshot::GridSnapshot;

    use super::{expand_grid, sample_grid};

    /// A model's cell value can run past its palette. Indexing past its end would panic on the UI thread and end the
    /// app.
    #[test]
    fn a_cell_past_the_palette_takes_the_first_colour() {
        let grid = GridSnapshot {
            width: 4,
            height: 2,
            cells: vec![0, 1, 2, 255, 1, 1, 7, 0],
            palette: &[[10, 20, 30, 255], [40, 50, 60, 255]],
        };
        let first = egui::Color32::from_rgb(10, 20, 30);
        let second = egui::Color32::from_rgb(40, 50, 60);
        assert_eq!(
            expand_grid(&grid),
            [first, second, first, first, second, second, first, first]
        );
        assert_eq!(sample_grid(&grid, 2, 1), [first, first]);
    }
}