pico-strike / src / detector.rs

mod analysis;
mod data;

use core::cell::Cell;

use alloc::vec::Vec;
use sachy_fmt::{info, unwrap};

use embassy_futures::select::select3;
use embassy_rp::peripherals::DMA_CH1;
use embassy_sync::{
    blocking_mutex::raw::NoopRawMutex,
    zerocopy_channel::{Channel as ZChannel, Receiver, Sender},
};
use embassy_time::{Duration, Instant, Ticker, Timer};

use crate::{
    adc::AdcDriver,
    constants::{BLIP_SIZE, BLOCK_SIZE},
    pwm::PwmDriver,
    rtc::GlobalRtc,
    updates::UpdateConnection,
    utils::{static_alloc, try_buffer, try_static_timestamped_block_vecs},
};

#[embassy_executor::task]
pub async fn detector_task(
    adc: AdcDriver<'static, NoopRawMutex, DMA_CH1>,
    pwm: PwmDriver<'static>,
    rtc: GlobalRtc<'static>,
) {
    info!("Waiting for RTC to start running");
    while !rtc.is_ready().await {
        Timer::after_secs(2).await;
    }

    info!("Allocating detector resources");

    let blocks = unwrap!(
        try_static_timestamped_block_vecs(4, BLOCK_SIZE),
        "Couldn't allocate block buffers"
    );

    let buf_channel: &mut ZChannel<NoopRawMutex, (i64, Vec<u16>)> =
        static_alloc(ZChannel::new(blocks));

    let (mut sender, mut receiver) = buf_channel.split();

    let mut detector = Detector::new(adc, pwm);
    let mut samples = unwrap!(try_buffer(64), "Failed to allocate sample buffer");
    let mut peaks = Vec::new();

    unwrap!(
        peaks.try_reserve_exact(BLOCK_SIZE),
        "Failed to allocate peaks buffer"
    );

    loop {
        let average = detector.tune(samples.as_mut_slice()).await;

        select3(
            detector.sample(&mut sender, rtc),
            detector.analyse(&mut receiver, average, &mut peaks),
            detector.tick(rtc),
        )
        .await;
    }
}

struct Detector<'device> {
    adc: AdcDriver<'device, NoopRawMutex, DMA_CH1>,
    pwm: PwmDriver<'device>,
    state: DetectorState,
}

#[derive(Debug)]
struct DetectorState {
    strikes: Cell<u16>,
    warn_level: Cell<u16>,
}

impl Default for DetectorState {
    fn default() -> Self {
        Self {
            strikes: Cell::new(0),
            warn_level: Cell::new(255),
        }
    }
}

impl<'device> Detector<'device> {
    fn new(adc: AdcDriver<'device, NoopRawMutex, DMA_CH1>, mut pwm: PwmDriver<'device>) -> Self {
        pwm.enable();

        Self {
            adc,
            pwm,
            state: DetectorState::default(),
        }
    }
}

impl Detector<'_> {
    async fn tune(&mut self, samples: &mut [u16]) -> u16 {
        info!("Tuning Detector for correct voltage settings");
        let mut duty = 0;
        self.pwm.set_duty(duty);
        Timer::after_secs(2).await;
        let mut act_value = self.adc.sample_average(samples).await;

        info!("initial ACT: {}", act_value);

        while act_value < 1364 {
            duty += 2;
            if duty >= 256 {
                duty = 0;
                self.pwm.set_duty(duty);
                Timer::after_secs(2).await;
                act_value = self.adc.sample_average(samples).await;
                info!("Restarting tuning");
                continue;
            }
            self.pwm.set_duty(duty);
            Timer::after_millis(250).await;
            act_value = self.adc.sample_average(samples).await;
            info!("ACT: {}, Duty: {}", act_value, duty);
        }

        duty = (duty / 3) * 2;
        self.pwm.set_duty(duty);
        info!("Set detection duty to: {}", duty);
        // Allow voltage level to stabilize after tuning
        Timer::after_secs(2).await;
        info!("Ready for strike detection");

        self.adc.sample_average(samples).await
    }

    async fn sample(
        &self,
        data: &mut Sender<'static, NoopRawMutex, (i64, Vec<u16>)>,
        rtc: GlobalRtc<'static>,
    ) {
        let track_time = rtc.track_time().await;

        loop {
            let (timestamp, buf) = data.send().await;
            *timestamp = track_time.timestamp();
            self.adc.sample(buf).await;
            data.send_done();
        }
    }

    async fn analyse(
        &self,
        data: &mut Receiver<'static, NoopRawMutex, (i64, Vec<u16>)>,
        mut average: u16,
        peaks: &mut Vec<u16>,
    ) {
        loop {
            peaks.clear();

            let (timestamp, buf) = data.receive().await;

            let new_avg =
                analysis::analyse_buffer_by_stepped_windows(buf.as_slice(), average, peaks);

            let blips = peaks.len();

            if blips >= BLIP_SIZE {
                self.state
                    .strikes
                    .update(|strike| strike.saturating_add(32));

                if let Some(net_data) = UpdateConnection::can_update() {
                    data::transmit_strike(
                        *timestamp,
                        buf.as_slice(),
                        peaks.as_slice(),
                        average,
                        &net_data,
                    );
                }

                info!(
                    "Strikes detected at {}s! avg {}, signal strength: {}, blip {}, level {}",
                    Instant::now().as_secs(),
                    average,
                    blips,
                    peaks.iter().max(),
                    self.state.warn_level.get()
                );
            }

            data.receive_done();

            average = new_avg;
        }
    }

    async fn tick(&self, rtc: GlobalRtc<'static>) {
        let mut inactive: Option<Instant> = None;
        let mut interval = Ticker::every(Duration::from_secs(1));
        let track_time = rtc.track_time().await;

        loop {
            interval.next().await;
            let strikes = self.state.strikes.get();
            let mut warn_level = self.state.warn_level.get();

            if strikes > 32 {
                warn_level = warn_level.saturating_add(strikes);
            }

            let decay = warn_level >> 8;

            self.state.strikes.set(0);
            self.state.warn_level.set(warn_level - decay);

            match inactive {
                Some(_) if decay > 0 => {
                    inactive = None;
                }
                Some(val) if val.elapsed() >= Duration::from_secs(3600) => break,
                None if decay == 0 => {
                    inactive = Some(Instant::now());
                }
                None => {
                    if let Some(net_data) = UpdateConnection::can_update() {
                        let now = track_time.timestamp();
                        data::transmit_level_update(now, decay, &net_data);
                    }
                }
                _ => continue,
            }
        }
    }
}