bioz-host-rs/src/app.rs

use std::sync::{Arc, Mutex};
use std::sync::atomic::{AtomicBool, Ordering};

use atomic_float::AtomicF32;
use tokio::{sync::mpsc::{Sender}};

use eframe::egui::{self, Button, CollapsingHeader, Color32, DragValue, Key, Label, Layout, Modifiers};
use egui_plot::{Corner, Legend, Line, Plot, PlotPoints, Points, PlotBounds};
use egui_dock::{DockArea, DockState, Style};

use crate::plot::TimeSeriesPlot;

use crate::signals::SingleFrequencySignal;

pub struct App {
    tree: DockState<String>,
    tab_viewer: TabViewer,
    run_impedancemeter_tx: Sender<SingleFrequencySignal>,
    pub magnitude: Arc<Mutex<f32>>,
    pub phase: Arc<Mutex<f32>>,
    pub magnitude_series: Arc<Mutex<TimeSeriesPlot>>,
    pub phase_series: Arc<Mutex<TimeSeriesPlot>>,
    pub connected: Arc<AtomicBool>,
    pub on: Arc<Mutex<bool>>,
    pub data_frequency: Arc<AtomicF32>,
    pub single_frequency: Arc<Mutex<u32>>
}

struct TabViewer {
    magnitude: Arc<Mutex<f32>>,
    phase: Arc<Mutex<f32>>,
    magnitude_series: Arc<Mutex<TimeSeriesPlot>>,
    phase_series: Arc<Mutex<TimeSeriesPlot>>,
    on: Arc<Mutex<bool>>,
    single_frequency: Arc<Mutex<u32>>,
    show_settings: bool,
    show_settings_toggle: Option<bool>,
}

impl TabViewer {
    fn single_tab(&mut self, ui: &mut egui::Ui) {
        egui::Frame::default().inner_margin(5).show(ui, |ui| {
            let settings = CollapsingHeader::new("Settings")
            .open(self.show_settings_toggle)
            .show(ui, |ui| {
                if let Ok(on) = self.on.lock() {
                    ui.add_enabled_ui(!*on, |ui| {
                        ui.horizontal(|ui| {
                            ui.label("Single Frequency:");
                            if let Ok(mut freq) = self.single_frequency.lock() {
                                ui.add(DragValue::new(&mut *freq).speed(0.1));
                            }
                            ui.label("Hz");
                        });
                    });
                }
            });

            self.show_settings_toggle = None;
            self.show_settings = !settings.fully_closed();

            ui.separator();

            let available_height = ui.available_height();
            let half_height = available_height / 2.0-2.0;

            // Plot pressure
            ui.allocate_ui_with_layout(
                egui::vec2(ui.available_width(), half_height),
                Layout::top_down(egui::Align::Min),
                |ui| {
                    // Magnitude
                    let magnitude = self.magnitude_series.lock().unwrap();
                    Plot::new("magnitude")
                        .legend(Legend::default().position(Corner::LeftTop))
                        .y_axis_label("Magnitude at {} Hz[Ω]")
                        .y_axis_min_width(80.0)
                        .show(ui, |plot_ui| {
                            plot_ui.line(
                                Line::new(format!("Magnitude at {} Hz", self.single_frequency.lock().unwrap()), magnitude.plot_values())
                                    .color(Color32::BLUE)
                            );
                        });
                },
            );
            // Plot pressure
            ui.allocate_ui_with_layout(
                egui::vec2(ui.available_width(), half_height),
                Layout::top_down(egui::Align::Min),
                |ui| {
                    // Phase
                    let phase = self.phase_series.lock().unwrap();
                    Plot::new("phase")
                        .legend(Legend::default().position(Corner::LeftTop))
                        .y_axis_label("Phase [rad]")
                        .y_axis_min_width(80.0)
                        .show(ui, |plot_ui| {
                            plot_ui.line(
                                Line::new(format!("Phase at {} Hz", self.single_frequency.lock().unwrap()), phase.plot_values())
                                    .color(Color32::RED)
                            );
                        });
                },
            );
        });
    }

    fn multi_tab(&mut self, ui: &mut egui::Ui) {
        ui.heading("Bode plots!");
    }
}

impl egui_dock::TabViewer for TabViewer {
    type Tab = String;

    fn title(&mut self, tab: &mut Self::Tab) -> egui::WidgetText {
        (&*tab).into()
    }

    fn ui(&mut self, ui: &mut egui::Ui, tab: &mut Self::Tab) {
        match tab.as_str() {
            "Single" => self.single_tab(ui),
            "Multi" => self.multi_tab(ui),
            _ => {
                let _ = ui.label("Unknown tab");
            }
        }
    }
}

impl App {
    pub fn new(run_impedancemeter_tx: Sender<SingleFrequencySignal>) -> Self {
        // Step 1: Initialize shared fields first
        let magnitude = Arc::new(Mutex::new(0.0));
        let phase = Arc::new(Mutex::new(0.0));
        let magnitude_series = Arc::new(Mutex::new(TimeSeriesPlot::new()));
        let phase_series = Arc::new(Mutex::new(TimeSeriesPlot::new()));
        let single_frequency = Arc::new(Mutex::new(50000));
        let on = Arc::new(Mutex::new(true));

        // Step 2: Now we can initialize tab_viewer
        let tab_viewer = TabViewer {
            magnitude: magnitude.clone(),
            phase: phase.clone(),
            magnitude_series: magnitude_series.clone(),
            phase_series: phase_series.clone(),
            single_frequency: single_frequency.clone(),
            on: on.clone(),
            show_settings: false,
            show_settings_toggle: None,
        };

        // Step 3: Construct App
        let app = App {
            tree: DockState::new(vec!["Single".to_string(), "Multi".to_string()]),
            tab_viewer,
            run_impedancemeter_tx,
            magnitude,
            phase,
            magnitude_series,
            phase_series,
            connected: Arc::new(AtomicBool::new(false)),
            on,
            data_frequency: Arc::new(AtomicF32::new(0.0)),
            single_frequency,
        };

        app.update_start_stop();
        app
    }

    pub fn update_start_stop(&self) {
        match *self.on.lock().unwrap() {
            true => {
                if let Err(e) = self.run_impedancemeter_tx.try_send(SingleFrequencySignal::Start(*self.single_frequency.lock().unwrap())) {
                    eprintln!("Failed to send start command: {:?}", e);
                }
            },
            false => {
                if let Err(e) = self.run_impedancemeter_tx.try_send(SingleFrequencySignal::Stop) {
                    eprintln!("Failed to send stop command: {:?}", e);
                }
            },
        }
    }

    pub fn reset_view(&self) {
        self.magnitude_series.lock().unwrap().clear();
        self.phase_series.lock().unwrap().clear();
    }
}

impl eframe::App for App {
    fn update(&mut self, ctx: &egui::Context, _frame: &mut eframe::Frame) {
        // Egui add a top bar
        egui::TopBottomPanel::top("top_bar").show(ctx, |ui| {
            egui::MenuBar::new().ui(ui, |ui| {
                let connected = self.connected.load(Ordering::Relaxed);
                
                egui::widgets::global_theme_preference_switch(ui);
                ui.separator();

                ui.label(format!("Data rate: {} Hz", self.data_frequency.load(Ordering::Relaxed)));

                ui.separator();

                if ui.add_enabled(connected, toggle_start_stop(&mut *self.on.lock().unwrap())).changed() {
                    self.update_start_stop();
                };

                ui.separator();

                if ui.add_enabled(connected, Button::new("Reset view")).clicked() {
                    self.reset_view();
                }

                // Spacer to push the LED to the right
                ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
                    ui.scope(|ui| {
                    let is_connected = self.connected.load(Ordering::Relaxed);
                    let color = if is_connected { Color32::GREEN } else { Color32::RED };
                    let tooltip = if is_connected { "Connected" } else { "Disconnected" };

                    // Allocate a fixed-size rectangle for the LED
                    let led_size = egui::Vec2::splat(12.0);
                    let (rect, response) = ui.allocate_exact_size(led_size, egui::Sense::hover());

                    // Draw the circle
                    let center = rect.center();
                    let radius = 5.0;
                    ui.painter().circle_filled(center, radius, color);

                    // Tooltip
                    if response.hovered() {
                        response.on_hover_text(tooltip);
                    }
                    });
                });
            });
        });

        egui::CentralPanel::default().show(ctx, |ui| {
            DockArea::new(&mut self.tree)
            .style(Style::from_egui(ctx.style().as_ref()))
            .show_leaf_close_all_buttons(false)
            .show_leaf_collapse_buttons(false)
            .show_close_buttons(false)
            .show_inside(ui, &mut self.tab_viewer);
        });
        
        // CMD- or control-W to close window
        if ctx.input(|i| i.modifiers.cmd_ctrl_matches(Modifiers::COMMAND))
            && ctx.input(|i| i.key_pressed(Key::W))
        {
            ctx.send_viewport_cmd(egui::ViewportCommand::Close);
        }

        // CMD- or control-W to close window
        if ctx.input(|i| i.key_pressed(Key::Space))
        {
            let value = *self.on.lock().unwrap();
            *self.on.lock().unwrap() = !value;
            self.update_start_stop();
        }

        // Send stop command when the window is closed
        if ctx.input(|i| i.viewport().close_requested()) {
            *self.on.lock().unwrap() = false;
            self.update_start_stop();
        }

        // Reset view
        if ctx.input(|i| i.key_pressed(Key::C)) {
            self.reset_view();
        }

        // Toggle setttings view
        if ctx.input(|i| i.key_pressed(egui::Key::S)) {
            self.tab_viewer.show_settings = !self.tab_viewer.show_settings;
            if self.tab_viewer.show_settings {
                self.tab_viewer.show_settings_toggle = Some(true);
            } else {
                self.tab_viewer.show_settings_toggle = Some(false);
            }
        }

        ctx.request_repaint();
    }
}


fn toggle_ui_start_stop(ui: &mut egui::Ui, on: &mut bool) -> egui::Response {
    let desired_size = ui.spacing().interact_size.y * egui::vec2(2.0, 1.0);
    let (rect, mut response) = ui.allocate_exact_size(desired_size, egui::Sense::click());
    if response.clicked() {
        *on = !*on;
        response.mark_changed();
    }
    response.widget_info(|| {
        egui::WidgetInfo::selected(egui::WidgetType::Checkbox, ui.is_enabled(), *on, "")
    });

    if ui.is_rect_visible(rect) {
        let how_on = ui.ctx().animate_bool_responsive(response.id, *on);
        let visuals = ui.style().interact_selectable(&response, *on);
        let rect = rect.expand(visuals.expansion);
        let radius = 0.5 * rect.height();
        ui.painter().rect(
            rect,
            radius,
            visuals.bg_fill,
            visuals.bg_stroke,
            egui::StrokeKind::Inside,
        );
        let circle_x = egui::lerp((rect.left() + radius)..=(rect.right() - radius), how_on);
        let center = egui::pos2(circle_x, rect.center().y);
        ui.painter()
            .circle(center, 0.75 * radius, visuals.bg_fill, visuals.fg_stroke);
    }

    response
}

pub fn toggle_start_stop(on: &mut bool) -> impl egui::Widget + '_ {
    move |ui: &mut egui::Ui| toggle_ui_start_stop(ui, on)
}