fleascope_rs/
flea_scope.rs

1use crate::flea_connector::{FleaConnector, FleaConnectorError};
2use crate::serial_terminal::{BusyFleaTerminal, ConnectionLostError, IdleFleaTerminal};
3use crate::trigger_config::{DigitalTrigger, StringifiedTriggerConfig, TriggerConfig};
4use polars::prelude::*;
5use std::io::Read;
6use std::time::Duration;
7
8#[derive(Debug, Clone, Copy, PartialEq, Eq)]
9pub enum ProbeType {
10    X1,
11    X10,
12}
13
14impl ProbeType {
15    pub fn to_multiplier(&self) -> i32 {
16        match self {
17            ProbeType::X1 => 1,
18            ProbeType::X10 => 10,
19        }
20    }
21}
22
23#[derive(Debug, Clone, PartialEq, Copy)]
24pub enum Waveform {
25    Sine,
26    Square,
27    Triangle,
28    Ekg,
29}
30
31impl Waveform {
32    pub fn as_str(&self) -> &'static str {
33        match self {
34            Waveform::Sine => "sine",
35            Waveform::Square => "square",
36            Waveform::Triangle => "triangle",
37            Waveform::Ekg => "ekg",
38        }
39    }
40}
41
42#[derive(Debug, thiserror::Error)]
43pub enum CaptureConfigError {
44    #[error("Time frame too large (max 3.49 seconds)")]
45    TimeFrameTooLarge,
46
47    #[error("Time frame too small (min 111 microseconds)")]
48    TimeFrameTooSmall,
49
50    #[error("Delay too large (max 1 second)")]
51    DelayTooLarge,
52
53    #[error("Voltage out of range")]
54    VoltageOutOfRange,
55}
56
57#[derive(Debug, thiserror::Error)]
58pub enum CalibrationError {
59    #[error("No zero calibration available for this probe")]
60    NoZeroCalibrarion,
61
62    #[error("No calibrarion available for this probe")]
63    NoCalibrationPresent,
64
65    #[error("Signal to unstable")]
66    UnstableSignal,
67
68    #[error("Failure while processing calibration data")]
69    CalibrationDataError(#[from] PolarsError),
70}
71
72pub struct ScopeReading {
73    pub effective_msps: f64,
74    pub data: Vec<u8>,
75}
76
77pub const RAW_COLUMN_NAME: &str = "bnc_raw";
78pub const CALIBRATED_COLUMN_NAME: &str = "bnc_calibrated";
79pub const BITMAP_COLUMN_NAME: &str = "bitmap";
80pub const TIME_COLUMN_NAME: &str = "time";
81
82impl ScopeReading {
83    pub fn parse_csv(&self) -> Result<LazyFrame, PolarsError> {
84        #[cfg(feature = "puffin")]
85        puffin::profile_function!();
86
87        let df = CsvReadOptions::default()
88            .with_has_header(false)
89            .into_reader_with_file_handle(std::io::Cursor::new(&self.data))
90            .finish()?
91            .lazy()
92            .select([
93                col("column_1")
94                    .alias(RAW_COLUMN_NAME)
95                    .cast(DataType::Float64),
96                col("column_2").alias(BITMAP_COLUMN_NAME),
97            ])
98            .with_row_index("row_index", Some(0))
99            .with_columns([
100                // Create time column using row index - more efficient than separate vector creation
101                (col("row_index").cast(DataType::Float64)
102                    * lit(1.0 / (self.effective_msps * 1_000_000.0)))
103                .alias(TIME_COLUMN_NAME),
104            ])
105            .select([
106                col(TIME_COLUMN_NAME),
107                col(RAW_COLUMN_NAME),
108                col(BITMAP_COLUMN_NAME),
109            ]);
110
111        Ok(df)
112    }
113
114    /// Extract bits from bitmap column
115    pub fn extract_bits(mut df: &mut DataFrame) -> Result<&DataFrame, PolarsError> {
116        #[cfg(feature = "puffin")]
117        puffin::profile_function!();
118
119        let bitmap_column = df.column(BITMAP_COLUMN_NAME)?;
120        let bitmap_strings = bitmap_column.str()?;
121
122        // Parse hex strings and extract bits
123        let mut bit_columns: Vec<Vec<bool>> = vec![Vec::new(); 10];
124
125        for bitmap_opt in bitmap_strings.into_iter() {
126            if let Some(bitmap_str) = bitmap_opt {
127                let bitmap_str = bitmap_str.trim_start_matches("0x");
128                if let Ok(bitmap_val) = u32::from_str_radix(bitmap_str, 16) {
129                    for (bit, column) in bit_columns.iter_mut().enumerate().take(10) {
130                        column.push((bitmap_val >> bit) & 1 == 1);
131                    }
132                } else {
133                    // Default to false for unparseable values
134                    for column in bit_columns.iter_mut().take(10) {
135                        column.push(false);
136                    }
137                }
138            } else {
139                // Default to false for null values
140                for column in bit_columns.iter_mut().take(10) {
141                    column.push(false);
142                }
143            }
144        }
145
146        for (bit, values) in bit_columns.into_iter().enumerate() {
147            let column: Column = Series::new(format!("bit_{}", bit).into(), values).into();
148            df = df.with_column(column)?;
149        }
150
151        Ok(df)
152    }
153}
154
155pub struct ReadingFleaScope {
156    _ver: String,
157    hostname: String,
158    serial: BusyFleaTerminal,
159    effective_msps: f64,
160}
161
162impl ReadingFleaScope {
163    pub fn try_get_result(
164        mut self,
165    ) -> Result<Result<(IdleFleaScope, ScopeReading), ReadingFleaScope>, ConnectionLostError> {
166        #[cfg(feature = "puffin")]
167        puffin::profile_function!();
168
169        match self.serial.try_get_result() {
170            Ok(r) => match r {
171                Ok((data, idle_terminal)) => Ok(Ok((
172                    IdleFleaScope {
173                        serial: idle_terminal,
174                        _ver: self._ver,
175                        hostname: self.hostname,
176                    },
177                    ScopeReading {
178                        effective_msps: self.effective_msps,
179                        data,
180                    },
181                ))),
182                Err(busy_terminal) => {
183                    self.serial = busy_terminal;
184                    Ok(Err(self))
185                }
186            },
187            Err(e) => Err(e),
188        }
189    }
190    pub fn cancel(self) -> IdleFleaScope {
191        let idle_serial = self.serial.cancel();
192        IdleFleaScope {
193            serial: idle_serial,
194            _ver: self._ver,
195            hostname: self.hostname,
196        }
197    }
198}
199
200pub struct IdleFleaScope {
201    serial: IdleFleaTerminal,
202    _ver: String,
203    hostname: String,
204}
205
206impl IdleFleaScope {
207    // Constants
208    const MSPS: u32 = 18; // Million samples per second. target sample rate
209    const MCU_MHZ: f64 = 120.0; // MCU clock frequency in MHz, used for calculations
210    const INTERLEAVE: u32 = 5; // number of ADCs interleaved
211    const TOTAL_SAMPLES: u32 = 2000;
212
213    /// Connect to a FleaScope device
214    pub fn connect(
215        name: Option<&str>,
216        port: Option<&str>,
217        read_calibrations: bool,
218    ) -> Result<(Self, FleaProbe, FleaProbe), FleaConnectorError> {
219        let serial = FleaConnector::connect(name, port, true)?;
220        let mut x1 = FleaProbe::new(ProbeType::X1);
221        let mut x10 = FleaProbe::new(ProbeType::X10);
222
223        let mut scope = Self::new(serial);
224        if read_calibrations {
225            x1.read_calibration_from_flash(&mut scope.serial);
226            x10.read_calibration_from_flash(&mut scope.serial);
227        }
228        Ok((scope, x1, x10))
229    }
230
231    /// Create a new FleaScope from an existing terminal connection
232    pub fn new(mut serial: IdleFleaTerminal) -> Self {
233        log::debug!("Turning off echo");
234        serial.exec_sync("echo off", None);
235
236        let ver = String::from_utf8(serial.exec_sync("ver", None)).expect("Failed to read version");
237        log::debug!("FleaScope version: {}", ver);
238        // TODO: check if version is compatible
239
240        let hostname =
241            String::from_utf8(serial.exec_sync("hostname", None)).expect("Failed to read hostname");
242        log::debug!("FleaScope hostname: {}", hostname);
243        // TODO: check if hostname is correct
244
245        Self {
246            serial,
247            _ver: ver,
248            hostname,
249        }
250    }
251
252    /// Set the waveform generator
253    pub fn set_waveform(&mut self, waveform: Waveform, hz: i32) {
254        self.serial
255            .exec_sync(&format!("wave {} {}", waveform.as_str(), hz), None);
256    }
257
258    /// Convert number1 to prescaler value
259    fn number1_to_prescaler(number1: u32) -> Result<u32, CaptureConfigError> {
260        let ps = if number1 > 1000 { 16 } else { 1 };
261        let t =
262            ((Self::MCU_MHZ * (number1 * Self::INTERLEAVE) as f64 / ps as f64 / Self::MSPS as f64)
263                + 0.5) as u32;
264
265        if t == 0 {
266            return Err(CaptureConfigError::TimeFrameTooSmall);
267        }
268        if t > 65535 {
269            return Err(CaptureConfigError::TimeFrameTooLarge);
270        }
271
272        Ok(ps * t)
273    }
274
275    /// Convert prescaler to effective MSPS
276    fn prescaler_to_effective_msps(prescaler: u32) -> f64 {
277        Self::MCU_MHZ * Self::INTERLEAVE as f64 / prescaler as f64
278    }
279
280    fn prepare_read_command(
281        time_frame: Duration,
282        trigger_fields: StringifiedTriggerConfig,
283        delay: Option<Duration>,
284    ) -> Result<(f64, String), CaptureConfigError> {
285        #[cfg(feature = "puffin")]
286        puffin::profile_function!();
287
288        let delay = delay.unwrap_or(Duration::from_millis(0));
289
290        // Validate time frame
291        if time_frame.as_secs_f64() > 3.49 {
292            return Err(CaptureConfigError::TimeFrameTooLarge);
293        }
294        if time_frame.as_secs() == 0 && time_frame.as_micros() < 111 {
295            return Err(CaptureConfigError::TimeFrameTooSmall);
296        }
297
298        // Validate delay
299        if delay.as_secs_f64() > 1.0 {
300            return Err(CaptureConfigError::DelayTooLarge);
301        }
302
303        let number1 = Self::MSPS * (time_frame.as_micros() as u32) / Self::TOTAL_SAMPLES;
304        if number1 == 0 {
305            return Err(CaptureConfigError::TimeFrameTooSmall);
306        }
307
308        let prescaler = Self::number1_to_prescaler(number1)?;
309        let effective_msps = Self::prescaler_to_effective_msps(prescaler);
310
311        let delay_samples = (delay.as_micros() as f64 * effective_msps) as u32;
312        if delay_samples > 1_000_000 {
313            return Err(CaptureConfigError::DelayTooLarge);
314        }
315        Ok((
316            effective_msps,
317            format!(
318                "scope {} {} {}",
319                number1,
320                trigger_fields.into_string(),
321                delay_samples
322            ),
323        ))
324    }
325
326    /// Raw data read from the oscilloscope
327    pub fn read_async(
328        self,
329        time_frame: Duration,
330        trigger_fields: StringifiedTriggerConfig,
331        delay: Option<Duration>,
332    ) -> Result<ReadingFleaScope, (IdleFleaScope, CaptureConfigError)> {
333        #[cfg(feature = "puffin")]
334        puffin::profile_function!();
335
336        match Self::prepare_read_command(time_frame, trigger_fields, delay) {
337            Ok((effective_msps, command)) => {
338                let data = self.serial.exec_async(&command);
339                Ok(ReadingFleaScope {
340                    _ver: self._ver,
341                    hostname: self.hostname,
342                    serial: data,
343                    effective_msps,
344                })
345            }
346            Err(e) => Err((self, e)),
347        }
348    }
349
350    pub fn read_sync(
351        &mut self,
352        time_frame: Duration,
353        trigger_fields: StringifiedTriggerConfig,
354        delay: Option<Duration>,
355    ) -> Result<ScopeReading, CaptureConfigError> {
356        #[cfg(feature = "puffin")]
357        puffin::profile_function!();
358
359        let (effective_msps, command) =
360            Self::prepare_read_command(time_frame, trigger_fields, delay)?;
361
362        let data = self.serial.exec_sync(&command, None);
363        Ok(ScopeReading {
364            effective_msps,
365            data,
366        })
367    }
368
369    pub fn stream(self) -> StreamingScope {
370        StreamingScope {
371            _ver: self._ver,
372            hostname: self.hostname,
373            serial: self.serial.exec_async("stream"),
374        }
375    }
376
377    /// Set the hostname
378    pub fn set_hostname(&mut self, hostname: &str) {
379        self.serial
380            .exec_sync(&format!("hostname {}", hostname), None);
381        self.hostname = hostname.to_string();
382    }
383
384    pub fn teardown(mut self) {
385        let _ = self.serial.exec_sync("echo on", None);
386        let _ = self.serial.exec_sync("prompt on", None);
387    }
388}
389
390pub struct StreamingScope {
391    _ver: String,
392    hostname: String,
393    serial: BusyFleaTerminal,
394}
395
396impl StreamingScope {
397    pub fn stop(self) -> IdleFleaScope {
398        let serial = self.serial.cancel();
399        IdleFleaScope {
400            serial,
401            _ver: self._ver,
402            hostname: self.hostname,
403        }
404    }
405
406    pub fn read(&mut self, n: usize) -> Result<Vec<u16>, std::io::Error> {
407        #[cfg(feature = "puffin")]
408        puffin::profile_function!();
409        let mut buffer = vec![0u8; n * 2];
410        self.serial.read_exact(&mut buffer)?;
411        Ok(buffer
412            .chunks_exact(2)
413            .map(|chunk| u16::from_le_bytes([chunk[0], chunk[1]]))
414            .collect())
415    }
416}
417
418#[derive(Debug)]
419pub struct FleaProbe {
420    multiplier: ProbeType,
421    cal_zero: Option<f64>, // value for 0V
422    cal_3v3: Option<f64>,  // value-diff 0V - 3.3V
423}
424
425impl Clone for FleaProbe {
426    fn clone(&self) -> Self {
427        Self {
428            multiplier: self.multiplier,
429            cal_zero: self.cal_zero,
430            cal_3v3: self.cal_3v3,
431        }
432    }
433}
434
435impl FleaProbe {
436    /// Create a new probe with the given multiplier
437    pub fn new(multiplier: ProbeType) -> Self {
438        Self {
439            multiplier,
440            cal_zero: None,
441            cal_3v3: None,
442        }
443    }
444
445    pub fn read_calibration_from_flash(&mut self, serial: &mut IdleFleaTerminal) {
446        let dim_result = String::from_utf8(serial.exec_sync(
447            &format!(
448                "dim cal_zero_x{} as flash, cal_3v3_x{} as flash",
449                self.multiplier.to_multiplier(),
450                self.multiplier.to_multiplier()
451            ),
452            None,
453        ))
454        .expect("Failed to read calibration from flash");
455
456        let expected_response = format!(
457            "var 'cal_zero_x{}' already declared at this scope\r\nvar 'cal_3v3_x{}' already declared at this scope",
458            self.multiplier.to_multiplier(), self.multiplier.to_multiplier()
459        );
460
461        if dim_result == expected_response {
462            log::debug!("Variables for calibration already declared. Reading values.");
463        }
464
465        let cal_zero_raw: i32 = String::from_utf8(serial.exec_sync(
466            &format!("print cal_zero_x{}", self.multiplier.to_multiplier()),
467            None,
468        ))
469        .expect("Failed to read cal_zero_x value")
470        .trim()
471        .parse()
472        .expect("Failed to parse cal_zero_x value");
473        let cal_3v3_raw: i32 = String::from_utf8(serial.exec_sync(
474            &format!("print cal_3v3_x{}", self.multiplier.to_multiplier()),
475            None,
476        ))
477        .expect("Failed to read cal_3v3_x value")
478        .trim()
479        .parse()
480        .expect("Failed to parse cal_3v3_x value");
481
482        self.cal_zero = Some((cal_zero_raw - 1000) as f64 + 2048.0);
483        self.cal_3v3 = Some((cal_3v3_raw - 1000) as f64 / self.multiplier.to_multiplier() as f64);
484
485        log::debug!(
486            "Probe x{} calibration: cal_zero={:?}, cal_3v3={:?}",
487            self.multiplier.to_multiplier(),
488            self.cal_zero,
489            self.cal_3v3
490        );
491    }
492
493    /// Set calibration values manually
494    pub fn set_calibration(&mut self, offset_0: f64, offset_3v3: f64) {
495        self.cal_zero = Some(offset_0);
496        self.cal_3v3 = Some(offset_3v3);
497    }
498
499    /// Write calibration values to flash
500    pub fn write_calibration_to_flash(
501        &self,
502        scope: &mut IdleFleaScope,
503    ) -> Result<(), CalibrationError> {
504        let cal_zero = self
505            .cal_zero
506            .ok_or(CalibrationError::NoCalibrationPresent)?;
507        let cal_3v3 = self.cal_3v3.ok_or(CalibrationError::NoCalibrationPresent)?;
508
509        let zero_value = (cal_zero - 2048.0 + 1000.0 + 0.5) as i32;
510        let v3v3_value = (cal_3v3 * self.multiplier.to_multiplier() as f64 + 1000.0 + 0.5) as i32;
511
512        scope.serial.exec_sync(
513            &format!(
514                "cal_zero_x{} = {}",
515                self.multiplier.to_multiplier(),
516                zero_value
517            ),
518            None,
519        );
520        scope.serial.exec_sync(
521            &format!(
522                "cal_3v3_x{} = {}",
523                self.multiplier.to_multiplier(),
524                v3v3_value
525            ),
526            None,
527        );
528
529        Ok(())
530    }
531
532    pub fn apply_calibration(&self, df: LazyFrame) -> LazyFrame {
533        #[cfg(feature = "puffin")]
534        puffin::profile_function!();
535
536        df.with_column(
537            self.raw_to_voltage(col(RAW_COLUMN_NAME))
538                .alias(CALIBRATED_COLUMN_NAME),
539        )
540    }
541
542    /// Read a stable value for calibration purposes
543    pub fn read_stable_value_for_calibration(
544        &self,
545        scope: &mut IdleFleaScope,
546    ) -> Result<f64, CalibrationError> {
547        let trigger_fields = DigitalTrigger::start_capturing_when()
548            .is_matching()
549            .into_trigger_fields();
550
551        let reading = scope
552            .read_sync(Duration::from_millis(20), trigger_fields, None)
553            .expect("This should not fail, as we are reading a stable value for calibration");
554        let df = reading.parse_csv()?;
555
556        let relevant_data = df.select([col(RAW_COLUMN_NAME)]).collect()?;
557        let bnc_series = relevant_data.column(RAW_COLUMN_NAME)?;
558        let bnc_values: Vec<f64> = bnc_series.f64()?.into_no_null_iter().collect();
559
560        let min_val = bnc_values.iter().fold(f64::INFINITY, |a, &b| a.min(b));
561        let max_val = bnc_values.iter().fold(f64::NEG_INFINITY, |a, &b| a.max(b));
562
563        if max_val - min_val > 14.0 {
564            return Err(CalibrationError::UnstableSignal);
565        }
566
567        let mean = bnc_values.iter().sum::<f64>() / bnc_values.len() as f64;
568        Ok(mean)
569    }
570
571    /// Convert raw ADC value to voltage
572    pub fn raw_to_voltage(&self, raw_value: Expr) -> Expr {
573        let cal_zero = self.cal_zero.expect("Calibration for 0V is not set");
574        let cal_3v3 = self.cal_3v3.expect("Calibration for 3.3V is not set");
575
576        (raw_value - cal_zero.into()) / cal_3v3.into() * 3.3.into()
577    }
578
579    /// Convert voltage to raw ADC value
580    pub fn voltage_to_raw(&self, voltage: f64) -> f64 {
581        let cal_zero = self.cal_zero.expect("Calibration for 0V is not set");
582        let cal_3v3 = self.cal_3v3.expect("Calibration for 3.3V is not set");
583
584        (voltage / 3.3 * cal_3v3) + cal_zero
585    }
586
587    /// Calibrate for 0V
588    pub fn calibrate_0(&mut self, scope: &mut IdleFleaScope) -> Result<f64, CalibrationError> {
589        // Try to preserve existing 3.3V calibration if available
590        let raw_value_3v3 = if let (Some(_), Some(_)) = (self.cal_zero, self.cal_3v3) {
591            Some(self.voltage_to_raw(3.3))
592        } else {
593            None
594        };
595
596        self.cal_zero = Some(self.read_stable_value_for_calibration(scope)?);
597
598        if let Some(raw_3v3) = raw_value_3v3 {
599            self.cal_3v3 = Some(raw_3v3 - self.cal_zero.unwrap());
600        }
601
602        Ok(self.cal_zero.unwrap())
603    }
604
605    /// Calibrate for 3.3V
606    pub fn calibrate_3v3(&mut self, scope: &mut IdleFleaScope) -> Result<f64, CalibrationError> {
607        let cal_zero = self.cal_zero.ok_or(CalibrationError::NoZeroCalibrarion)?;
608
609        let raw_3v3 = self.read_stable_value_for_calibration(scope)?;
610        self.cal_3v3 = Some(raw_3v3 - cal_zero);
611
612        Ok(self.cal_3v3.unwrap())
613    }
614
615    /// Get calibration values
616    pub fn calibration(&self) -> (Option<f64>, Option<f64>) {
617        (self.cal_zero, self.cal_3v3)
618    }
619}
620
621#[cfg(test)]
622mod tests {
623    use super::*;
624
625    #[test]
626    fn test_waveform_as_str() {
627        assert_eq!(Waveform::Sine.as_str(), "sine");
628        assert_eq!(Waveform::Square.as_str(), "square");
629        assert_eq!(Waveform::Triangle.as_str(), "triangle");
630        assert_eq!(Waveform::Ekg.as_str(), "ekg");
631    }
632
633    #[test]
634    fn test_number1_to_prescaler() {
635        assert!(IdleFleaScope::number1_to_prescaler(100).is_ok());
636        assert!(IdleFleaScope::number1_to_prescaler(0).is_err());
637    }
638}