use brk_error::Result;
use brk_traversable::Traversable;
use brk_types::{Cents, Height, Indexes, StoredF32, Version};
use vecdb::{
    AnyStoredVec, AnyVec, Database, EagerVec, Exit, PcoVec, ReadableVec, Rw, StorageMode, VecIndex,
    WritableVec,
};

use crate::{
    blocks, indexes,
    internal::{PerBlock, Price, PriceTimesRatioCents, per_block::stddev::period_suffix},
};

#[derive(Traversable)]
pub struct StdDevBand<M: StorageMode = Rw> {
    #[traversable(flatten)]
    pub ratio: PerBlock<StoredF32, M>,
    pub price: Price<PerBlock<Cents, M>>,
}

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

    pub _0sd: Price<PerBlock<Cents, M>>,
    pub p0_5sd: StdDevBand<M>,
    pub p1sd: StdDevBand<M>,
    pub p1_5sd: StdDevBand<M>,
    pub p2sd: StdDevBand<M>,
    pub p2_5sd: StdDevBand<M>,
    pub p3sd: StdDevBand<M>,
    pub m0_5sd: StdDevBand<M>,
    pub m1sd: StdDevBand<M>,
    pub m1_5sd: StdDevBand<M>,
    pub m2sd: StdDevBand<M>,
    pub m2_5sd: StdDevBand<M>,
    pub m3sd: StdDevBand<M>,
}

impl StdDevPerBlockExtended {
    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);

        macro_rules! import {
            ($suffix:expr) => {
                PerBlock::forced_import(db, &format!("{name}_{}{p}", $suffix), version, indexes)?
            };
        }

        macro_rules! import_price {
            ($suffix:expr) => {
                Price::forced_import(db, &format!("{name}_{}{p}", $suffix), version, indexes)?
            };
        }

        macro_rules! import_band {
            ($suffix:expr) => {{
                StdDevBand {
                    ratio: import!(concat!("ratio_", $suffix)),
                    price: import_price!($suffix),
                }
            }};
        }

        Ok(Self {
            days,
            sd: import!("ratio_sd"),
            zscore: import!("ratio_zscore"),
            _0sd: import_price!("0sd"),
            p0_5sd: import_band!("p0_5sd"),
            p1sd: import_band!("p1sd"),
            p1_5sd: import_band!("p1_5sd"),
            p2sd: import_band!("p2sd"),
            p2_5sd: import_band!("p2_5sd"),
            p3sd: import_band!("p3sd"),
            m0_5sd: import_band!("m0_5sd"),
            m1sd: import_band!("m1sd"),
            m1_5sd: import_band!("m1_5sd"),
            m2sd: import_band!("m2sd"),
            m2_5sd: import_band!("m2_5sd"),
            m3sd: import_band!("m3sd"),
        })
    }

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

        self.compute_bands(starting_indexes, exit, sma, source)
    }

    fn compute_bands(
        &mut self,
        starting_indexes: &Indexes,
        exit: &Exit,
        sma: &impl ReadableVec<Height, StoredF32>,
        source: &impl ReadableVec<Height, StoredF32>,
    ) -> Result<()> {
        let source_version = source.version();

        self.mut_band_height_vecs()
            .try_for_each(|v| -> Result<()> {
                v.validate_computed_version_or_reset(source_version)?;
                Ok(())
            })?;

        let starting_height = self
            .mut_band_height_vecs()
            .map(|v| Height::from(v.len()))
            .min()
            .unwrap()
            .min(starting_indexes.height);

        let start = starting_height.to_usize();

        let source_len = source.len();
        let source_data = source.collect_range_at(start, source_len);

        let sma_data = sma.collect_range_at(start, sma.len());
        let sd_data = self.sd.height.collect_range_at(start, self.sd.height.len());

        const MULTIPLIERS: [f32; 12] = [
            0.5, 1.0, 1.5, 2.0, 2.5, 3.0, -0.5, -1.0, -1.5, -2.0, -2.5, -3.0,
        ];
        for (vec, mult) in self.mut_band_height_vecs().zip(MULTIPLIERS) {
            vec.truncate_if_needed_at(start)?;
            for (offset, _) in source_data.iter().enumerate() {
                let average = sma_data[offset];
                let sd = sd_data[offset];
                vec.push(average + StoredF32::from(mult * *sd));
            }
        }

        {
            let _lock = exit.lock();
            self.mut_band_height_vecs()
                .try_for_each(|v| v.write().map(|_| ()))?;
        }

        self.zscore.height.compute_zscore(
            starting_indexes.height,
            source,
            sma,
            &self.sd.height,
            exit,
        )?;

        Ok(())
    }

    pub(crate) fn compute_cents_bands(
        &mut self,
        starting_indexes: &Indexes,
        series_price: &impl ReadableVec<Height, Cents>,
        sma: &impl ReadableVec<Height, StoredF32>,
        exit: &Exit,
    ) -> Result<()> {
        macro_rules! compute_band_price {
            ($price:expr, $band_source:expr) => {
                $price
                    .cents
                    .compute_binary::<Cents, StoredF32, PriceTimesRatioCents>(
                        starting_indexes.height,
                        series_price,
                        $band_source,
                        exit,
                    )?;
            };
        }

        compute_band_price!(&mut self._0sd, sma);
        compute_band_price!(&mut self.p0_5sd.price, &self.p0_5sd.ratio.height);
        compute_band_price!(&mut self.p1sd.price, &self.p1sd.ratio.height);
        compute_band_price!(&mut self.p1_5sd.price, &self.p1_5sd.ratio.height);
        compute_band_price!(&mut self.p2sd.price, &self.p2sd.ratio.height);
        compute_band_price!(&mut self.p2_5sd.price, &self.p2_5sd.ratio.height);
        compute_band_price!(&mut self.p3sd.price, &self.p3sd.ratio.height);
        compute_band_price!(&mut self.m0_5sd.price, &self.m0_5sd.ratio.height);
        compute_band_price!(&mut self.m1sd.price, &self.m1sd.ratio.height);
        compute_band_price!(&mut self.m1_5sd.price, &self.m1_5sd.ratio.height);
        compute_band_price!(&mut self.m2sd.price, &self.m2sd.ratio.height);
        compute_band_price!(&mut self.m2_5sd.price, &self.m2_5sd.ratio.height);
        compute_band_price!(&mut self.m3sd.price, &self.m3sd.ratio.height);

        Ok(())
    }

    fn mut_band_height_vecs(
        &mut self,
    ) -> impl Iterator<Item = &mut EagerVec<PcoVec<Height, StoredF32>>> {
        [
            &mut self.p0_5sd.ratio.height,
            &mut self.p1sd.ratio.height,
            &mut self.p1_5sd.ratio.height,
            &mut self.p2sd.ratio.height,
            &mut self.p2_5sd.ratio.height,
            &mut self.p3sd.ratio.height,
            &mut self.m0_5sd.ratio.height,
            &mut self.m1sd.ratio.height,
            &mut self.m1_5sd.ratio.height,
            &mut self.m2sd.ratio.height,
            &mut self.m2_5sd.ratio.height,
            &mut self.m3sd.ratio.height,
        ]
        .into_iter()
    }
}
