bioz-host-rs/src/communication.rs

use log::{error, info};

use tokio::select;
use tokio::sync::mpsc::{Sender, Receiver};

use std::sync::atomic::{AtomicBool, AtomicU32, Ordering};
use atomic_float::AtomicF32;

use std::sync::{Arc, Mutex};

use crate::icd;
use crate::client::WorkbookClient;

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

use crate::signals::StartStopSignal;

pub async fn communicate_with_hardware(
    mut run_impedancemeter_rx: Receiver<StartStopSignal>,
    run_impedancemeter_tx: Sender<StartStopSignal>,
    magnitude: Arc<Mutex<f32>>,
    phase: Arc<Mutex<f32>>,
    magnitude_series: Arc<Mutex<TimeSeriesPlot>>,
    phase_series: Arc<Mutex<TimeSeriesPlot>>,
    bode_series: Arc<Mutex<BodePlot>>,
    connected: Arc<AtomicBool>,
    data_frequency: Arc<AtomicF32>,
) {
    let data_counter = Arc::new(AtomicU32::new(0));
    let data_counter_clone = data_counter.clone();
    tokio::spawn(async move {
        loop {
            tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
            data_frequency.store(data_counter.load(Ordering::Relaxed) as f32, Ordering::Relaxed);
            data_counter.store(0, Ordering::Relaxed);
        }
    });

    #[derive(Default)]
    struct Settings {
        frequency: Option<StartStopSignal>,
    }
    let mut settings = Settings::default();
    
    loop {
        let workbook_client = match WorkbookClient::new() {
            Ok(client) => {
                info!("Connected to hardware successfully.");
                if let Some(frequency) = settings.frequency {
                    run_impedancemeter_tx.send(frequency).await.unwrap();
                }
                connected.store(true, Ordering::Relaxed);
                client
            },
            Err(e) => {
                error!("Failed to connect to hardware: {:?}", e);
                tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                continue;
            }
        };

        // Subscribe to SingleImpedanceOutputTopic
        let mut single_impedance_sub = workbook_client
            .client
            .subscribe_multi::<icd::SingleImpedanceOutputTopic>(8)
            .await
            .unwrap();

        let data = (magnitude_series.clone(), phase_series.clone(), magnitude.clone(), phase.clone());
        let data_counter_clone_single = data_counter_clone.clone();

        tokio::spawn(async move {
            while let Ok(val) = single_impedance_sub.recv().await {
                let mut mag_plot = data.0.lock().unwrap();
                let mut phase_plot = data.1.lock().unwrap(); 
                let mut mag_val = data.2.lock().unwrap();
                let mut phase_val = data.3.lock().unwrap();

                *mag_val = val.magnitude;
                *phase_val = val.phase;
    
                mag_plot.add(val.magnitude as f64);
                phase_plot.add(val.phase as f64);
                
                data_counter_clone_single.fetch_add(1, Ordering::Relaxed);
            }
            info!("SingleImpedanceOutputTopic subscription ended.");
        });

        // Subscribe to MultiImpedanceOutputTopic28
        let mut multi_impedance_28_sub = workbook_client
            .client
            .subscribe_multi::<icd::MultiImpedanceOutputTopic28>(8)
            .await
            .unwrap();

        let data = bode_series.clone();
        let data_counter_clone_multi = data_counter_clone.clone();

        tokio::spawn(async move {
            while let Ok(val) = multi_impedance_28_sub.recv().await {
                let mut bode_plot = data.lock().unwrap();

                let magnitudes = val.magnitudes.into_iter().collect::<Vec<f32>>();
                bode_plot.update_magnitudes(magnitudes);
                let phases = val.phases.into_iter().collect::<Vec<f32>>();
                bode_plot.update_phases(phases);

                data_counter_clone_multi.fetch_add(1, Ordering::Relaxed);   
            }
        });
        
        loop {
            select! {
                Some(frequency) = run_impedancemeter_rx.recv() => {
                    match frequency {
                        StartStopSignal::StartSingle(freq, dft_num) => {
                            if let Err(e) = workbook_client.start_impedancemeter_single(freq, dft_num).await {
                                error!("Failed to start impedancemeter: {:?}", e);
                            } else {
                                settings.frequency = Some(StartStopSignal::StartSingle(freq, dft_num));
                                info!("Impedancemeter started at frequency: {}", freq);
                            }
                        },
                        StartStopSignal::StartMulti(dft_num, num_points) => {
                            if let Err(e) = workbook_client.start_impedancemeter_multi(dft_num, num_points).await {
                                error!("Failed to start multi-point impedancemeter: {:?}", e);
                            } else {
                                settings.frequency = Some(StartStopSignal::StartMulti(dft_num, num_points));
                                info!("Multi-point Impedancemeter started.");
                            }
                        },
                        StartStopSignal::Stop => {
                            if let Err(e) = workbook_client.stop_impedancemeter().await {
                                error!("Failed to stop impedancemeter: {:?}", e);
                            } else {
                                settings.frequency = Some(StartStopSignal::Stop);
                                info!("Impedancemeter stopped.");
                            }
                        },
                    }
                }
                _ = workbook_client.wait_closed() => {
                    // Handle client closure
                    info!("Client connection closed.");
                    break;
                }

                else => {
                    // All channels closed
                    break;
                }

            }
        }
        info!("Communication with hardware ended.");
        connected.store(false, Ordering::Relaxed);
        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
    }
}