ai-0/lab3/src/main.rs

#![cfg_attr(not(debug_assertions), windows_subsystem = "windows")] // hide console window on Windows in release

mod algo;

use crate::algo::{
    updateState, Ant, AntsSimulationConfig, AntsSimulationState, EdgeExtraData, EdgesVec,
    VerticesVec,
};
use eframe::egui;
use eframe::egui::{Frame, Ui};
use eframe::emath::Numeric;
use std::collections::HashSet;
use std::ops::RangeInclusive;
use bgtu_ai_utility::gui::lengths_table::{draw_lengths_table, UpdatePending};
use bgtu_ai_utility::gui::render::render_graph;

fn main() -> eframe::Result {
    let options = eframe::NativeOptions {
        viewport: egui::ViewportBuilder::default().with_inner_size([640.0, 400.0]),
        ..Default::default()
    };
    eframe::run_native(
        "Ants simulation",
        options,
        Box::new(|_cc| Ok(Box::<MyApp>::default())),
    )
}

enum GlobalState {
    Edit,
    Running,
}

struct MyApp {
    simulation: AntsSimulationState,
    config: AntsSimulationConfig,
    vertex_update: UpdatePending,
    state: GlobalState,
    vertex_locations: Vec<(f32, f32)>,
    ants_per_vertex: usize,
}

impl Default for MyApp {
    fn default() -> Self {
        return Self {
            simulation: AntsSimulationState {
                edges: EdgesVec::new(),
                vertices: VerticesVec::new(),
                ants: Vec::new(),
            },
            vertex_update: UpdatePending::NoChange,
            config: AntsSimulationConfig {
                ferment_weight: 0.5,
                heuristic_coefficient: 0.5,
                q: 1.0,
                r: 0.5,
            },
            state: GlobalState::Edit {},
            vertex_locations: Vec::new(),
            ants_per_vertex: 1,
        };
    }
}

fn _slider<T: Numeric>(
    ui: &mut Ui,
    name: &str,
    storage: &mut T,
    range: RangeInclusive<T>,
    step: f64,
) {
    let label = ui.label(name);

    ui.scope(|ui| {
        ui.spacing_mut().slider_width = ui.available_width()
            - ui.spacing().interact_size.x
            - ui.spacing().button_padding.x * 2.0;
        ui.add(egui::Slider::new(storage, range).step_by(step))
            .labelled_by(label.id);
    });
}

impl eframe::App for MyApp {
    fn update(&mut self, ui: &eframe::egui::Context, _frame: &mut eframe::Frame) {
        egui::CentralPanel::default().show(ui, |ui| match self.state {
            GlobalState::Edit {} => {
                match self.vertex_update {
                    UpdatePending::NoChange => {}
                    UpdatePending::Add => {
                        let new_vi = self.simulation.vertices.add(HashSet::new());
                        let mut newEdgesSet = HashSet::new();
                        for (vi, v) in self.simulation.vertices.iter_indexed_mut() {
                            if (vi == new_vi) {
                                continue;
                            }
                            let ei = self.simulation.edges.add(
                                new_vi,
                                vi,
                                1.0,
                                EdgeExtraData {
                                    ferment_intensity: 0.0,
                                },
                            );
                            newEdgesSet.insert(ei);
                            v.insert(ei);
                        }
                        self.simulation.vertices[new_vi] = newEdgesSet;
                        self.vertex_update = UpdatePending::NoChange;
                    }
                    UpdatePending::Remove(vi) => {
                        let mut eis = Vec::with_capacity(self.simulation.vertices[vi].len());
                        for ei in self.simulation.vertices[vi].iter() {
                            eis.push(*ei)
                        }
                        for ei in eis {
                            self.simulation.vertices[self.simulation.edges[ei].another(vi)]
                                .remove(&ei);
                            self.simulation.edges.remove(ei);
                        }
                        self.simulation.vertices.remove(vi);
                        self.vertex_update = UpdatePending::NoChange;
                    }
                }
                edit_panel(self, ui)
            }
            GlobalState::Running { .. } => {
                ui.add_enabled_ui(false, |ui| edit_panel(self, ui));
                ui.ctx().show_viewport_immediate(
                    egui::ViewportId::from_hash_of("Visualisation"),
                    egui::ViewportBuilder::default()
                        .with_title("Visualisation")
                        .with_inner_size([640.0, 480.0])
                        .with_resizable(false),
                    |ui, _| {
                        egui::CentralPanel::default()
                            .frame(Frame::default().inner_margin(0.0))
                            .show(ui, |ui| visualization_panel(self, ui));
                    },
                );
            }
        });
    }
}

fn edit_panel(data: &mut MyApp, ui: &mut Ui) {
    let mut run: bool = false;

    _slider(
        ui,
        "Ferment weight",
        &mut data.config.ferment_weight,
        0.0..=1.0,
        0.001,
    );
    ui.label("");
    _slider(
        ui,
        "Heuristic coefficient",
        &mut data.config.heuristic_coefficient,
        0.0..=1.0,
        0.001,
    );
    ui.label("");
    _slider(ui, "Q", &mut data.config.q, 0.0..=1.0, 0.001);
    ui.label("");
    _slider(ui, "r", &mut data.config.r, 0.0..=1.0, 0.001);
    ui.label("");
    _slider(
        ui,
        "Ants per vertex",
        &mut data.ants_per_vertex,
        1..=100,
        1.0,
    );
    ui.label("");
    ui.horizontal(|ui| {
        if ui.button("Add vertex").clicked() {
            data.vertex_update = UpdatePending::Add;
        }
        ui.separator();
        run = ui.button("Run").clicked();
    });

    draw_lengths_table(
        ui,
        &mut data.simulation.vertices,
        &mut data.simulation.edges,
        &mut data.vertex_update,
    );

    if run {
        let mut coords = vec![(0.0, 0.0); data.simulation.vertices.capacity()];

        for (i, _) in data.simulation.vertices.iter_indexed() {
            coords[i] = (
                rand::random::<f32>() * 0.8 + 0.1,
                rand::random::<f32>() * 0.8 + 0.1,
            )
        }

        data.state = GlobalState::Running;
        data.vertex_locations = coords;

        let allowed_locations = data
            .simulation
            .vertices
            .iter_indexed()
            .map(|(i, _)| i)
            .collect::<HashSet<usize>>();

        let mut ants = Vec::new();
        for (i, _) in data.simulation.vertices.iter_indexed() {
            for _ in 0..data.ants_per_vertex {
                ants.push(Ant::new(i, allowed_locations.clone()))
            }
        }
        data.simulation.ants = ants;
    }
}

fn visualization_panel(data: &mut MyApp, ui: &mut Ui) {
    ui.vertical(|ui| {
        ui.horizontal(|ui| {
            if ui.button("Exit").clicked() {
                data.state = GlobalState::Edit
            }

            if ui.button("Step").clicked() {
                updateState(&mut data.simulation, &mut data.config, &mut rand::rng());
            }
            ui.label("");
        });
        draw_ants(data, ui);
        if ui.input(|i| i.viewport().close_requested()) {
            data.state = GlobalState::Edit
        }
    });
}

fn draw_ants(data: &mut MyApp, ui: &mut Ui) {
    let mut cap: f64 = 0.000000000001;

    for w in data
        .simulation
        .edges
        .iter()
        .map(|e| e.extra.ferment_intensity)
    {
        if w > cap {
            cap = w
        }
    }

    render_graph(
        ui,
        &data.simulation.vertices,
        data.vertex_locations.as_mut_slice(),
        &data.simulation.edges,
        |e| e.extra.ferment_intensity / cap,
    )
}