brk_computer 0.11.0

use std::path::Path;

use brk_cohort::{Filter, Filtered};
use brk_error::Result;
use brk_indexer::Lengths;
use brk_traversable::Traversable;
use brk_types::{Cents, Height, PartsPerMillionSigned64, StoredI64, StoredU64, Version};
use rayon::prelude::*;
use vecdb::{AnyStoredVec, AnyVec, Exit, ReadableVec, Rw, StorageMode, WritableVec};

use crate::{distribution::state::AddrCohortState, internal::PerBlockWithDeltas, price};

use crate::distribution::metrics::{AllSupplyCache, ImportConfig};

use super::super::traits::{CohortVecs, DynCohortVecs};
use super::metrics::AddrCohortMetrics;

#[derive(Traversable)]
pub struct AddrCohortVecs<M: StorageMode = Rw> {
    starting_height: Option<Height>,

    #[traversable(skip)]
    pub state: Option<Box<AddrCohortState>>,

    #[traversable(flatten)]
    pub metrics: AddrCohortMetrics<M>,

    pub addr_count: PerBlockWithDeltas<StoredU64, StoredI64, PartsPerMillionSigned64, M>,
}

impl AddrCohortVecs {
    pub(crate) fn forced_import(
        cfg: &ImportConfig,
        states_path: Option<&Path>,
        all_supply: &AllSupplyCache,
    ) -> Result<Self> {
        let addr_count = PerBlockWithDeltas::forced_import(
            cfg.db,
            &cfg.name("addr_count"),
            cfg.version,
            Version::TWO,
            cfg.indexes,
            cfg.cached_starts,
        )?;

        Ok(Self {
            starting_height: None,
            state: states_path.map(|path| Box::new(AddrCohortState::new(path, cfg.full_name))),
            metrics: AddrCohortMetrics::forced_import(cfg, all_supply)?,
            addr_count,
        })
    }

    pub(crate) fn reset_starting_height(&mut self) {
        self.starting_height = Some(Height::ZERO);
    }

    pub(crate) fn par_iter_vecs_mut(
        &mut self,
    ) -> impl ParallelIterator<Item = &mut dyn AnyStoredVec> {
        let mut vecs: Vec<&mut dyn AnyStoredVec> = Vec::new();
        vecs.push(&mut self.addr_count.height as &mut dyn AnyStoredVec);
        vecs.extend(self.metrics.collect_all_vecs_mut());
        vecs.into_par_iter()
    }

    pub(crate) fn write_state(&mut self, height: Height, cleanup: bool) -> Result<()> {
        if let Some(state) = self.state.as_mut() {
            state.inner.write(height, cleanup)?;
        }
        Ok(())
    }
}

impl Filtered for AddrCohortVecs {
    fn filter(&self) -> &Filter {
        &self.metrics.filter
    }
}

impl DynCohortVecs for AddrCohortVecs {
    fn min_stateful_len(&self) -> usize {
        self.addr_count
            .height
            .len()
            .min(self.metrics.min_stateful_len())
    }

    fn reset_state_starting_height(&mut self) {
        self.reset_starting_height();
        if let Some(state) = self.state.as_mut() {
            state.reset();
        }
    }

    fn import_state(&mut self, starting_height: Height) -> Result<Height> {
        if let Some(state) = self.state.as_mut() {
            if let Some(mut prev_height) = starting_height.decremented() {
                prev_height = state.inner.import_at_or_before(prev_height)?;

                state.inner.supply.value = self
                    .metrics
                    .supply
                    .total
                    .sats
                    .height
                    .collect_one(prev_height)
                    .unwrap();
                state.inner.supply.utxo_count = *self
                    .metrics
                    .outputs
                    .unspent_count
                    .height
                    .collect_one(prev_height)
                    .unwrap();
                state.addr_count = *self.addr_count.height.collect_one(prev_height).unwrap();

                state.inner.restore_realized_cap();

                let result = prev_height.incremented();
                self.starting_height = Some(result);
                Ok(result)
            } else {
                self.starting_height = Some(Height::ZERO);
                Ok(Height::ZERO)
            }
        } else {
            self.starting_height = Some(starting_height);
            Ok(starting_height)
        }
    }

    fn validate_computed_versions(&mut self, base_version: Version) -> Result<()> {
        use vecdb::WritableVec;
        self.addr_count
            .height
            .validate_computed_version_or_reset(base_version)?;
        Ok(())
    }

    fn push_state(&mut self, height: Height) {
        if self.starting_height.is_some_and(|h| h > height) {
            return;
        }

        if let Some(state) = self.state.as_ref() {
            self.addr_count.height.push(state.addr_count.into());
            self.metrics.supply.push_state(&state.inner);
            self.metrics.outputs.push_state(&state.inner);
            self.metrics.activity.push_state(&state.inner);
            self.metrics.realized.push_state(&state.inner);
        }
    }

    fn push_unrealized_state(&mut self, _height_price: Cents) {}

    fn compute_rest_part1(
        &mut self,
        prices: &price::Vecs,
        starting_lengths: &Lengths,
        exit: &Exit,
    ) -> Result<()> {
        self.metrics
            .compute_rest_part1(prices, starting_lengths, exit)
    }

    fn write_state(&mut self, height: Height, cleanup: bool) -> Result<()> {
        if let Some(state) = self.state.as_mut() {
            state.inner.write(height, cleanup)?;
        }
        Ok(())
    }

    fn reset_cost_basis_data_if_needed(&mut self) -> Result<()> {
        if let Some(state) = self.state.as_mut() {
            state.inner.reset_cost_basis_data_if_needed()?;
        }
        Ok(())
    }

    fn reset_single_iteration_values(&mut self) {
        if let Some(state) = self.state.as_mut() {
            state.inner.reset_single_iteration_values();
        }
    }
}

impl CohortVecs for AddrCohortVecs {
    fn compute_from_stateful(
        &mut self,
        starting_lengths: &Lengths,
        others: &[&Self],
        exit: &Exit,
    ) -> Result<()> {
        self.addr_count.height.compute_sum_of_others(
            starting_lengths.height,
            others
                .iter()
                .map(|v| &v.addr_count.height)
                .collect::<Vec<_>>()
                .as_slice(),
            exit,
        )?;
        self.metrics.compute_from_sources(
            starting_lengths,
            &others.iter().map(|v| &v.metrics).collect::<Vec<_>>(),
            exit,
        )?;
        Ok(())
    }
}