bioz-host-rs/src/app.rs

use std::fmt::Debug;

use log::{info, error};
use tokio::sync;

use core::f32;

// use std::f32::consts::PI;
use std::ops::RangeInclusive;
use std::sync::{Arc, Mutex};
use std::sync::atomic::Ordering;

use chrono::Local;

use tokio::sync::{mpsc, watch};

use eframe::egui::{self, Button, CollapsingHeader, Color32, ComboBox, DragValue, Id, Key, Label, Layout, Modal, Modifiers, RichText, TextEdit, Widget};
use egui_plot::{Corner, GridInput, GridMark, Legend, Line, Plot, PlotPoint, Points};
use egui_dock::{DockArea, DockState, Style};
use egui_extras::{TableBuilder, Column};

use crate::state::{AppState, ControlCommand, HardwareConnected, HardwareState, MeasurementDataState, Mode};

use crate::logging::LoggingStates;
use crate::plot::{TimeSeriesPlot, BodePlot};

use crate::signals::{LoggingSignal, StartStopSignal};

use crate::icd::{BioImpedanceLeadMode, IcdDftNum, SweepPoints, ElectrodeConfiguration, ElectrodeOptionsWithMultiplexer};

const LEAD_MODES: [BioImpedanceLeadMode; 2] = [
    BioImpedanceLeadMode::TwoLead,
    BioImpedanceLeadMode::FourLead,
];

const DFTNUM_VARIANTS: [IcdDftNum; 13] = [
    IcdDftNum::Num4,    IcdDftNum::Num8,    IcdDftNum::Num16,   IcdDftNum::Num32,
    IcdDftNum::Num64,   IcdDftNum::Num128,  IcdDftNum::Num256,  IcdDftNum::Num512,
    IcdDftNum::Num1024, IcdDftNum::Num2048, IcdDftNum::Num4096,
    IcdDftNum::Num8192, IcdDftNum::Num16384,
];

const SWEEP_POINTS_VARIANTS: [SweepPoints; 2] = [
    SweepPoints::Eight,
    SweepPoints::Eighteen,
];

#[derive(Clone, Copy,Debug, PartialEq, Eq)]
enum TabActive {
    Single,
    Sweep,
    Shortcuts,
}

pub struct App {
    tree: DockState<String>,
    tab_viewer: TabViewer,
    tab_active: TabActive,
    log_tx: mpsc::Sender<LoggingSignal>,
    pub gui_logging_state: Arc<Mutex<LoggingStates>>,
    log_filename: String,
    log_marker_modal: bool,
    log_marker: String,
    measurement_data: Arc<MeasurementDataState>,
    control_tx: mpsc::Sender<ControlCommand>,
    app_state_rx: sync::watch::Receiver<AppState>,
    hardware_state_rx: sync::watch::Receiver<HardwareState>
}

struct TabViewer {
    show_settings: bool,
    show_settings_toggle: Option<bool>,
    measurement_data: Arc<MeasurementDataState>,
    control_tx: mpsc::Sender<ControlCommand>,
    app_state_rx: sync::watch::Receiver<AppState>,
    hardware_state_rx: sync::watch::Receiver<HardwareState>,
}

trait ElectrodeOption: Copy + Debug + PartialEq + 'static {
    const ALL: &'static [Self];
}

impl ElectrodeOption for ElectrodeOptionsWithMultiplexer {
    const ALL: &'static [Self] = &ElectrodeOptionsWithMultiplexer::ALL;
}

fn electrode_combo<T: ElectrodeOption>(
    ui: &mut egui::Ui,
    id: &str,
    value: &mut T,
) -> bool {
    let mut changed = false;
    egui::ComboBox::from_id_salt(id)
        .selected_text(format!("{:?}", value))
        .width(60.0)
        .show_ui(ui, |ui| {
            for &option in T::ALL {
                let response = ui.selectable_value(value, option, format!("{:?}", option));
                if response.changed() {
                    changed = true;
                }
            }
        });
    changed
}

fn render_two_lead<T: ElectrodeOption>(
    ui: &mut egui::Ui,
    values: &mut [T; 2],
) -> bool {
    let mut changed = false;
    ui.horizontal(|ui| {
        ui.label("Drive/sense: Electrode+ (");
        changed |= electrode_combo(ui, "e1_select", &mut values[0]);
        ui.label("), Electrode- (");
        changed |= electrode_combo(ui, "e2_select", &mut values[1]);
        ui.label(")");
    });
    changed
}

fn render_four_lead<T: ElectrodeOption>(
    ui: &mut egui::Ui,
    values: &mut [T; 4],
) -> bool {
    let mut changed = false;
    ui.horizontal(|ui| {
        ui.label("Drive: I+ (");
        changed |= electrode_combo(ui, "i1_select", &mut values[0]);
        ui.label("), I- (");
        changed |= electrode_combo(ui, "i2_select", &mut values[1]);

        ui.label(") | Sense: V+ (");
        changed |= electrode_combo(ui, "v1_select", &mut values[2]);
        ui.label("), V- (");
        changed |= electrode_combo(ui, "v2_select", &mut values[3]);
        ui.label(")");
    });
    changed
}

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 running = self.hardware_state_rx.borrow().running {
                    ui.add_enabled_ui(!running, |ui| {
                        ui.horizontal(|ui| {
                            ui.label("Lead Mode:");

                            let mut lead_mode = self.app_state_rx.borrow().lead_mode;

                            // Map current lead mode to index
                            let mut index = LEAD_MODES
                                .iter()
                                .position(|&m| m == lead_mode)
                                .unwrap_or(0);

                            ComboBox::from_id_salt("LeadMode")
                                .width(60.0)
                                .show_index(ui, &mut index, LEAD_MODES.len(), |i| {
                                    match LEAD_MODES[i] {
                                        BioImpedanceLeadMode::TwoLead => "2-Lead",
                                        BioImpedanceLeadMode::FourLead => "4-Lead",
                                    }
                                    .to_string()
                                });

                            // Update lead mode if changed
                            if lead_mode != LEAD_MODES[index] {
                                lead_mode = LEAD_MODES[index];
                                self.control_tx.try_send(ControlCommand::ChangeLeadMode(lead_mode)).unwrap();
                            }
                        });
                        ui.horizontal(|ui| {
                            // Show lead configuration
                            ui.label("Lead Configuration:");
                            let hardware_connected = self.hardware_state_rx.borrow().hardware_connected;
                            let lead_mode = self.app_state_rx.borrow().lead_mode;

                            let mut settings = self.app_state_rx.borrow().electrode_settings;
                        
                            match (hardware_connected, lead_mode) {
                                (HardwareConnected::WithoutMultiplexer, BioImpedanceLeadMode::TwoLead) => {
                                    ui.label(format!("Drive/sense: Electrode+ (CE0), Electrode- (AIN1)"));
                                }
                                (HardwareConnected::WithoutMultiplexer, BioImpedanceLeadMode::FourLead) => {
                                    ui.label(format!("Drive: I+ (CE0), I- (AIN1) | Sense: V+ (AIN2), V- (AIN3)"));
                                }
                                (HardwareConnected::WithMultiplexer, BioImpedanceLeadMode::TwoLead) => {
                                    if render_two_lead(ui, &mut settings.with_multiplexer_2_lead) {
                                        self.control_tx.try_send(ControlCommand::ChangeElectrodeSettings(settings)).unwrap();
                                    }
                                }
                                (HardwareConnected::WithMultiplexer, BioImpedanceLeadMode::FourLead) => {
                                    if render_four_lead(ui, &mut settings.with_multiplexer_4_lead) {
                                        self.control_tx.try_send(ControlCommand::ChangeElectrodeSettings(settings)).unwrap();
                                    }
                                }
                                (HardwareConnected::None, _) => {}
                            }
                        });
                    });
                    ui.add_enabled_ui(!running, |ui| {
                        ui.horizontal(|ui| {
                            ui.label("Single Frequency:");
                            let mut freq = self.app_state_rx.borrow().single_frequency;
                            let response = ui.add(DragValue::new(&mut freq).speed(0.1));
                            if response.changed() {
                                self.control_tx.try_send(ControlCommand::SetFrequency(freq)).unwrap();
                            }
                            ui.label("Hz");
                        });
                        ui.horizontal(|ui| {
                            ui.label("ADC samples per DFT:");
                            let mut dft_num = self.app_state_rx.borrow().dft_num;
                            let mut index = DFTNUM_VARIANTS.iter().position(|&x| x == dft_num).unwrap_or(0);
                            ComboBox::from_id_salt("Dftnum")
                                .width(75.0)
                                .show_index(ui, &mut index, DFTNUM_VARIANTS.len(), |i| {
                                    format!("{}", 1 << (2 + i)) // 2^2 = 4, 2^3 = 8, ..., 2^14 = 16384
                                });
                            let new_value = DFTNUM_VARIANTS[index];
                            if dft_num != new_value {
                                dft_num = new_value;
                                self.control_tx.try_send(ControlCommand::ChangeDftNum(dft_num)).unwrap();
                            };
                        });
                    });
                    ui.add_enabled_ui(running, |ui| {
                        ui.horizontal(|ui| {
                            ui.label("Periods per DFT:");
                            match (running, self.hardware_state_rx.borrow().periods_per_dft) {
                                (true, Some(periods)) => {
                                    ui.add(Label::new(format!("{:.2}", periods)));
                                },
                                (true, None) => {
                                    ui.add(Label::new("N/A"));
                                },
                                (false, _) => {
                                    ui.add(Label::new("Start to determine!"));
                                }
                            }
                        });
                    });
                }
            });

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


    fn sweep_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 running = self.hardware_state_rx.borrow().running {
                    ui.add_enabled_ui(!running, |ui| {
                        ui.horizontal(|ui| {
                            ui.label("Lead Mode:");

                            let mut lead_mode = self.app_state_rx.borrow().lead_mode;

                            // Map current lead mode to index
                            let mut index = LEAD_MODES
                                .iter()
                                .position(|&m| m == lead_mode)
                                .unwrap_or(0);

                            ComboBox::from_id_salt("LeadMode")
                                .width(60.0)
                                .show_index(ui, &mut index, LEAD_MODES.len(), |i| {
                                    match LEAD_MODES[i] {
                                        BioImpedanceLeadMode::TwoLead => "2-Lead",
                                        BioImpedanceLeadMode::FourLead => "4-Lead",
                                    }
                                    .to_string()
                                });

                            // Update lead mode if changed
                            if lead_mode != LEAD_MODES[index] {
                                lead_mode = LEAD_MODES[index];
                                self.control_tx.try_send(ControlCommand::ChangeLeadMode(lead_mode)).unwrap();
                            }
                        });
                        ui.horizontal(|ui| {
                            // Show lead configuration
                            ui.label("Lead Configuration:");
                            let hardware_connected = self.hardware_state_rx.borrow().hardware_connected;
                            let lead_mode = self.app_state_rx.borrow().lead_mode;

                            let mut settings = self.app_state_rx.borrow().electrode_settings;
                            match (hardware_connected, lead_mode) {
                                (HardwareConnected::WithoutMultiplexer, BioImpedanceLeadMode::TwoLead) => {
                                    ui.label(format!("Drive/sense: Electrode+ (CE0), Electrode- (AIN1)"));
                                }
                                (HardwareConnected::WithoutMultiplexer, BioImpedanceLeadMode::FourLead) => {
                                    ui.label("Drive: I+ (CE0), I- (AIN1) | Sense: V+ (AIN2), V- (AIN3)");
                                }
                                (HardwareConnected::WithMultiplexer, BioImpedanceLeadMode::TwoLead) => {
                                    if render_two_lead(ui, &mut settings.with_multiplexer_2_lead) {
                                        self.control_tx.try_send(ControlCommand::ChangeElectrodeSettings(settings)).unwrap();
                                    }
                                }
                                (HardwareConnected::WithMultiplexer, BioImpedanceLeadMode::FourLead) => {
                                    if render_four_lead(ui, &mut settings.with_multiplexer_4_lead) {
                                        self.control_tx.try_send(ControlCommand::ChangeElectrodeSettings(settings)).unwrap();
                                    }
                                }
                                (HardwareConnected::None, _) => {}
                            }
                        });
                    });
                    ui.add_enabled_ui(!running, |ui| {
                        ui.horizontal(|ui| {
                            ui.label("Measurement Points:");
                            let mut measurement_points = self.app_state_rx.borrow().sweep_points;
                            let mut index = SWEEP_POINTS_VARIANTS.iter().position(|&x| x == measurement_points).unwrap_or(0);
                            ComboBox::from_id_salt("MeasurementPoints")
                                .width(75.0)
                                .show_index(ui, &mut index, SWEEP_POINTS_VARIANTS.len(), |i| {
                                    format!("{:?}", SWEEP_POINTS_VARIANTS[i].len())
                                });
                            let new_value = SWEEP_POINTS_VARIANTS[index];
                            if measurement_points != new_value {
                                measurement_points = new_value;
                                self.control_tx.try_send(ControlCommand::ChangeSweepPoints(measurement_points)).unwrap();
                                info!("Sweep Points setting changed!");
                            }
                        });
                    });
                    ui.add_enabled_ui(running, |ui| {
                        let (freq, periods_per_dft_vec) = self.hardware_state_rx.borrow().periods_per_dft_sweep.clone();

                        fn format_frequency(freq: u32) -> String {
                            if freq >= 1_000 {
                                // kHz
                                if freq % 1_000 == 0 {
                                    format!("{}k", (freq / 1_000) as u64)
                                } else {
                                    format!("{:.1}k", freq as f32 / 1_000.0)
                                }
                            } else {
                                // Hz
                                if freq % 1 == 0 {
                                    format!("{}", freq as u64)
                                } else {
                                    format!("{:.1}", freq)
                                }
                            }
                        }

                        TableBuilder::new(ui)
                            .column(Column::auto())
                            .columns(Column::remainder().at_most(30.0), freq.len())
                            .header(10.0, |mut header| {
                                header.col(|ui| {
                                    ui.label("Frequency [Hz]:");
                                });
                                for freq in &freq {
                                    header.col(|ui| {
                                        ui.label(format_frequency(*freq));
                                    });
                                }
                            })
                            .body(|mut body| {
                                body.row(5.0, |mut row| {
                                    row.col(|ui| {
                                        ui.label("Periods per DFT:");
                                    });
                                    if let Some(periods) = periods_per_dft_vec {
                                        for period in &periods {
                                            row.col(|ui| {
                                                ui.label(format!("{:.1}", period));
                                            });
                                        }
                                    } else {
                                        for _ in &freq {
                                            row.col(|ui| {
                                                ui.label("N/A");
                                            });
                                        }
                                    }
                                });
                            });
                    });
                }
            });

            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;

            // Bodeplot magnitude
            ui.allocate_ui_with_layout(
                egui::vec2(ui.available_width(), half_height),
                Layout::top_down(egui::Align::Min),
                |ui| {
                    // Magnitude
                    let bode_plot = self.measurement_data.bode_plot.lock().unwrap();
                    Plot::new("bode_mag")
                        .legend(Legend::default().position(Corner::LeftTop))
                        .y_axis_label("Magnitude [Ω]")
                        .x_grid_spacer(log10_grid_spacer)
                        .x_axis_formatter(log10_axis_formatter)
                        // .auto_bounds([true; 2].into())
                        .y_axis_min_width(80.0)
                        .default_x_bounds(0.0, 200_000_f64.log10())
                        .label_formatter(log10x_formatter)
                        .show(ui, |plot_ui| {
                            plot_ui.line(
                                Line::new("Magnitude", bode_plot.plot_magnitudes())
                                    .color(Color32::LIGHT_BLUE)
                            );
                            plot_ui.points(
                                Points::new("Magnitude measurement", bode_plot.plot_magnitudes())
                                    .color(Color32::BLUE)
                                    .radius(1.5)
                            );
                        });
                },
            );

            // Bodeplot phase
            ui.allocate_ui_with_layout(
                egui::vec2(ui.available_width(), half_height),
                Layout::top_down(egui::Align::Min),
                |ui| {
                    // Phase
                    let bode_plot = self.measurement_data.bode_plot.lock().unwrap();
                    Plot::new("bode_phase")
                        .legend(Legend::default().position(Corner::LeftTop))
                        .y_axis_label("Phase [rad]")
                        .x_grid_spacer(log10_grid_spacer)
                        .x_axis_formatter(log10_axis_formatter)
                        // .auto_bounds([true; 2].into())
                        .y_axis_min_width(80.0)
                        .default_x_bounds(0.0, 200_000_f64.log10())
                        .label_formatter(log10x_formatter)
                        .show(ui, |plot_ui| {
                            plot_ui.line(
                                Line::new("Phase", bode_plot.plot_phases())
                                    .color(Color32::LIGHT_RED)
                            );
                            plot_ui.points(
                                Points::new("Phase measurement", bode_plot.plot_phases())
                                    .color(Color32::RED)
                                    .radius(1.5)
                            );
                        });
                },
            );
        });
    }

    fn shortcuts(&mut self, ui: &mut egui::Ui) {
        ui.heading("Shortcuts");
        ui.label("Space: Start/Stop measurement");
        ui.label("R: Reset view/plots");
        ui.label("S: Toggle settings");
        ui.label("L: Toggle logging");
        ui.label("A: Add marker (when logging)");
        ui.label("CMD-W: Close window");
    }
}

fn log10_grid_spacer(input: GridInput) -> Vec<GridMark> {
    let base = 10u32;
    assert!(base >= 2);
    let basef = base as f64;

    let step_size = basef.powi(input.base_step_size.abs().log(basef).ceil() as i32);
    let (min, max) = input.bounds;
    let mut steps = vec![];

    for step_size in [step_size, step_size * basef, step_size * basef * basef] {
        if step_size == basef.powi(-1) {
            // FIXME: float comparison
            let first = ((min / (step_size * basef)).floor() as i64) * base as i64;
            let last = (max / step_size).ceil() as i64;

            let mut logs = vec![0.0; base as usize - 2];
            for (i, j) in logs.iter_mut().enumerate() {
                *j = basef * ((i + 2) as f64).log(basef);
            }

            steps.extend((first..last).step_by(base as usize).flat_map(|j| {
                logs.iter().map(move |i| GridMark {
                    value: (j as f64 + i) * step_size,
                    step_size,
                })
            }));
        } else {
            let first = (min / step_size).ceil() as i64;
            let last = (max / step_size).ceil() as i64;

            steps.extend((first..last).map(move |i| GridMark {
                value: (i as f64) * step_size,
                step_size,
            }));
        }
    }
    steps
}

fn log10_axis_formatter(mark: GridMark, _range: &RangeInclusive<f64>) -> String {
    let base = 10u32;
    let basef = base as f64;
    let prec = (-mark.step_size.log10().round()).max(0.0) as usize;
    format!("{:.*e}", prec, basef.powf(mark.value))
}

fn log10x_formatter(name: &str, value: &PlotPoint) -> String {
    let base = 10u32;
    let basef = base as f64;
    format!(
        "{}\nx: {:.d$e}\ny: {:.d$e}",
        name,
        basef.powf(value.x),
        value.y,
        d = 3
    )
}

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

    fn title(&mut self, tab: &mut Self::Tab) -> eframe::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),
            "Sweep" => self.sweep_tab(ui),
            "Shortcuts" => self.shortcuts(ui),
            _ => {
                let _ = ui.label("Unknown tab");
            }
        }
    }
}

impl App {
    pub fn new(log_tx: mpsc::Sender<LoggingSignal>, 
        measurement_data: Arc<MeasurementDataState>, 
        control_tx: mpsc::Sender<ControlCommand>, 
        app_state_rx: watch::Receiver<AppState>, 
        hardware_state_rx: watch::Receiver<HardwareState>) -> Self {
        // Step 1: Initialize shared fields first 
        let tab_active = TabActive::Single;

        // Step 2: Now we can initialize tab_viewer
        let control_tx_clone = control_tx.clone();
        let app_state_clone = app_state_rx.clone();
        let hardware_state_clone = hardware_state_rx.clone();
        let tab_viewer = TabViewer {            measurement_data: measurement_data.clone(),
            show_settings: false,
            show_settings_toggle: None,
            control_tx: control_tx_clone,
            app_state_rx: app_state_clone,
            hardware_state_rx: hardware_state_clone
        };

        // Step 3: Construct App
        let app = App {
            tree: DockState::new(vec!["Single".to_string(), "Sweep".to_string(), "Shortcuts".to_string()]),
            tab_viewer,
            tab_active,
            log_tx,
            gui_logging_state: Arc::new(Mutex::new(LoggingStates::Idle)),
            log_filename: format!("log_{}_single.csv", Local::now().format("%Y%m%d")),
            log_marker_modal: false,
            log_marker: String::new(),
            measurement_data,
            control_tx,
            app_state_rx,
            hardware_state_rx,
        };

        // For testing purposes, populate the Bode plot with a sample low-pass filter response
        // let fc = 1000.0; // cutoff frequency in Hz

        // let freqs = MeasurementPointSet::Eighteen.values().to_vec();
        // let magnitudes = freqs.iter()
        //     .map(|&f| {
        //         1.0 / (1.0 + (f / fc).powi(2)).sqrt()
        //     })
        //     .collect::<Vec<f32>>();
        // let phases = freqs.iter()
        //     .map(|&f| {
        //         -(f / fc).atan() * 180.0 / PI
        //     })
        //     .collect::<Vec<f32>>();

        // app.bode_plot.lock().unwrap().update_magnitudes(MeasurementPointSet::Eighteen, magnitudes);
        // app.bode_plot.lock().unwrap().update_phases(MeasurementPointSet::Eighteen, phases);

        app.toggle_start_stop();
        app
    }

    pub fn toggle_start_stop(&self) {
        match (self.tab_active, self.hardware_state_rx.borrow().running) {
            (TabActive::Single, false) => {
                self.control_tx.try_send(ControlCommand::Start(Mode::Single)).unwrap();
            },
            (TabActive::Sweep, false) => {
                self.control_tx.try_send(ControlCommand::Start(Mode::Sweep)).unwrap();
            },
            (_, true) => {
                self.control_tx.try_send(ControlCommand::Stop).unwrap();
            },
            (_, _) => {}
        }
    }

    pub fn reset_view(&self) {
        self.measurement_data.magnitude_series.lock().unwrap().clear();
        self.measurement_data.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 is_connected = self.hardware_state_rx.borrow().hardware_connected != HardwareConnected::None;
                
                egui::widgets::global_theme_preference_switch(ui);
                ui.separator();

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

                ui.separator();

                let running = self.hardware_state_rx.borrow().running;

                let (color, text) = match running {
                    true => (Color32::DARK_RED,  "Stop"),
                    false => (Color32::DARK_GREEN, "Start"),
                };

                let start_stop_button = Button::new(
                    RichText::new(text)
                        .strong()
                        .color(Color32::WHITE)
                ).fill(color)
                .corner_radius(5.0)
                .min_size(egui::vec2(45.0, 0.0))
                .frame(true);

                if ui.add_enabled(is_connected, start_stop_button).clicked() {
                    self.toggle_start_stop();
                }

                ui.separator();

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

                ui.separator();

                ui.add_enabled(is_connected, Label::new("Logging:"));

                let gui_logging_state = *self.gui_logging_state.lock().unwrap();

                let mut logging_enabled = !matches!(gui_logging_state, LoggingStates::Idle);

                if ui.add_enabled(is_connected, custom_toggle_widget(&mut logging_enabled)).changed() {
                    let signal = match gui_logging_state {
                        LoggingStates::Idle => LoggingSignal::StartFileLogging(self.log_filename.clone()),
                        LoggingStates::Starting | LoggingStates::Logging => LoggingSignal::StopFileLogging,
                    };

                    self.log_tx.try_send(signal).unwrap_or_else(|e| {
                        error!("Failed to send logging toggle signal: {:?}", e);
                    });
                };

                ui.separator();

                ui.add_enabled_ui(is_connected && logging_enabled, |ui| {
                    if Button::new("Add marker").corner_radius(5.0).min_size(egui::vec2(20.0, 0.0)).ui(ui).on_hover_text("Add a marker to the current time series plots").clicked() {
                        self.log_marker_modal = true;
                    }
                });

                if self.log_marker_modal {
                    if Modal::new(Id::new("modal_marker"))
                        .show(ctx, |ui| {
                            ui.vertical_centered(|ui| {
                                ui.heading("Add marker");

                                ui.add_space(16.0);

                                // Input field for marker
                                let text_edit_response = TextEdit::singleline(&mut self.log_marker)
                                    .desired_width(100.0)
                                    .id(Id::new("marker_field"))
                                    .hint_text("Marker name")
                                    .ui(ui);

                                ui.add_space(16.0);

                                // Centered Add button
                                let add_clicked = Button::new("Add")
                                    .corner_radius(5.0)
                                    .min_size(egui::vec2(80.0, 30.0))
                                    .ui(ui)
                                    .clicked();

                                // Check for Enter key in the TextEdit
                                let enter_pressed = text_edit_response.lost_focus() && ui.input(|i| i.key_pressed(egui::Key::Enter));

                                if add_clicked || enter_pressed {
                                    info!("Adding marker: {}", self.log_marker);
                                    self.log_tx.try_send(LoggingSignal::AddMarker(self.log_marker.clone())).unwrap_or_else(|e| {
                                        error!("Failed to send logging marker signal: {:?}", e);
                                    });
                                    self.log_marker = String::new();
                                    ui.close();
                                }

                                // Request focus on the text edit when the modal opens
                                text_edit_response.request_focus();
                            });
                        })
                        .should_close()
                    {
                        self.log_marker_modal = false;
                        self.log_marker = String::new();
                    }
                }

                ui.separator();

                ui.add_enabled_ui(gui_logging_state == LoggingStates::Idle, |ui| {
                    ui.label("Log filename:");
                    TextEdit::singleline(&mut self.log_filename).desired_width(150.0).id(Id::new("file_name_field")).ui(ui);
                });

                // Spacer to push the LED to the right
                ui.with_layout(egui::Layout::right_to_left(egui::Align::Center), |ui| {
                    ui.scope(|ui| {
                    let (color, tooltip) = match self.hardware_state_rx.borrow().hardware_connected {
                        HardwareConnected::None => {
                            (Color32::DARK_RED, "Disconnected")
                        },
                        HardwareConnected::WithoutMultiplexer => {
                            (Color32::DARK_GREEN, "Connected without multiplexer")
                        },
                        HardwareConnected::WithMultiplexer => {
                            (Color32::DARK_GREEN, "Connected with multiplexer")
                        },
                    };

                    // 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);
        });

        // Do something when a tab is changed
        if let Some((_, tab)) = self.tree.find_active_focused() {
            match tab.as_str() {
                "Single" => {
                    if self.tab_active != TabActive::Single {
                        self.tab_active = TabActive::Single;
                        self.control_tx.try_send(ControlCommand::Stop).unwrap();
                        if *self.gui_logging_state.lock().unwrap() == LoggingStates::Logging {
                            self.log_tx.try_send(LoggingSignal::StopFileLogging).unwrap_or_else(|e| {
                                error!("Failed to send logging logging signal: {:?}", e);
                            });
                        }
                        self.log_filename = format!("log_{}_single.csv", Local::now().format("%Y%m%d"));
                        info!("Switched to Single tab");
                    }
                }
                "Sweep" => {
                    if self.tab_active != TabActive::Sweep {
                        self.tab_active = TabActive::Sweep;
                        self.control_tx.try_send(ControlCommand::Stop).unwrap();
                        if *self.gui_logging_state.lock().unwrap() == LoggingStates::Logging {
                            self.log_tx.try_send(LoggingSignal::StopFileLogging).unwrap_or_else(|e| {
                                error!("Failed to send logging logging signal: {:?}", e);
                            });
                        }
                        self.log_filename = format!("log_{}_sweep.csv", Local::now().format("%Y%m%d"));
                        info!("Switched to Sweep tab");
                    }
                }
                "Shortcuts" => {
                    // if self.tab_active != TabActive::Shortcuts {
                        // self.tab_active = TabActive::Shortcuts;
                        // *self.on.lock().unwrap() = false;
                        // self.update_start_stop();
                        // info!("Switched to Shortcuts tab");
                    // }
                }
                _ => {
                }
            }
        }

        // 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);
        }
        
        // Check if the file name field is focused
        // If it is, we don't want to trigger shortcuts
        let file_name_field_is_focused = ctx.memory(|memory| memory.has_focus(Id::new("file_name_field")));
        let marker_field_is_focused = ctx.memory(|memory| memory.has_focus(Id::new("marker_field")));
        if !file_name_field_is_focused && !marker_field_is_focused {

            // Space to start/stop measurement
            if ctx.input(|i| i.key_pressed(Key::Space))
            {
                self.toggle_start_stop();
            }

            // Send stop command when the window is closed
            if ctx.input(|i| i.viewport().close_requested()) {
                self.toggle_start_stop();
            }

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

            // Toggle marker modal
            if ctx.input(|i| i.key_pressed(egui::Key::A)) && *self.gui_logging_state.lock().unwrap() == LoggingStates::Logging {
                self.log_marker_modal = !self.log_marker_modal;
            }

            // Enable/disable GUI logging
            if ctx.input(|i| i.key_pressed(egui::Key::L)) {
                let gui_logging_enabled = *self.gui_logging_state.lock().unwrap();
                match gui_logging_enabled {
                    LoggingStates::Starting => {
                        // If currently starting, do nothing
                        return;
                    },
                    LoggingStates::Logging => {
                        if let Err(e) = self.log_tx.try_send(LoggingSignal::StopFileLogging) {
                            error!("Failed to send logging signal: {:?}", e);
                        }
                    },
                    LoggingStates::Idle => {
                        if let Err(e) = self.log_tx.try_send(LoggingSignal::StartFileLogging(self.log_filename.clone())) {
                            error!("Failed to send logging signal: {:?}", e);
                        }
                    },
                }
            }

            // 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 custom_toggle_widget(on: &mut bool) -> impl egui::Widget + '_ {
    move |ui: &mut egui::Ui| toggle_ui_start_stop(ui, on)
}