brk_computer 0.2.5

use brk_error::Result;
use brk_traversable::Traversable;
use brk_types::{Height, Indexes, StoredF32, Version};
use vecdb::{Database, Exit, ReadableVec, Rw, StorageMode};

use crate::{blocks, indexes, internal::PerBlock};

use super::period_suffix;

#[derive(Traversable)]
pub struct StdDevPerBlock<M: StorageMode = Rw> {
    days: usize,
    pub sma: PerBlock<StoredF32, M>,
    pub sd: PerBlock<StoredF32, M>,
}

impl StdDevPerBlock {
    pub(crate) fn forced_import(
        db: &Database,
        name: &str,
        period: &str,
        days: usize,
        parent_version: Version,
        indexes: &indexes::Vecs,
    ) -> Result<Self> {
        let version = parent_version + Version::TWO;
        let p = period_suffix(period);

        let sma = PerBlock::forced_import(db, &format!("{name}_sma{p}"), version, indexes)?;
        let sd = PerBlock::forced_import(db, &format!("{name}_sd{p}"), version, indexes)?;

        Ok(Self { days, sma, sd })
    }

    pub(crate) fn compute_all(
        &mut self,
        blocks: &blocks::Vecs,
        starting_indexes: &Indexes,
        exit: &Exit,
        source: &impl ReadableVec<Height, StoredF32>,
    ) -> Result<()> {
        if self.days == usize::MAX {
            self.sma.height.compute_sma_(
                starting_indexes.height,
                source,
                usize::MAX,
                exit,
                None,
            )?;
            self.sd.height.compute_expanding_sd(
                starting_indexes.height,
                source,
                &self.sma.height,
                exit,
            )?;
            return Ok(());
        }

        let window_starts = blocks.lookback.start_vec(self.days);

        self.sma.height.compute_rolling_average(
            starting_indexes.height,
            window_starts,
            source,
            exit,
        )?;

        self.sd.height.compute_rolling_sd(
            starting_indexes.height,
            window_starts,
            source,
            &self.sma.height,
            exit,
        )?;

        Ok(())
    }
}