brk/crates/brk_state/src/cohort.rs

use std::{cmp::Ordering, path::Path};

use brk_core::{CheckedSub, Dollars, Height, Result, Sats};

use crate::{PriceToAmount, UnrealizedState};

use super::{RealizedState, SupplyState};

#[derive(Clone)]
pub struct CohortState {
    pub supply: SupplyState,
    pub realized: Option<RealizedState>,
    pub satblocks_destroyed: Sats,
    pub satdays_destroyed: Sats,
    pub price_to_amount: PriceToAmount,
}

impl CohortState {
    pub fn default_and_import(path: &Path, name: &str, compute_dollars: bool) -> Result<Self> {
        Ok(Self {
            supply: SupplyState::default(),
            realized: compute_dollars.then_some(RealizedState::NAN),
            satblocks_destroyed: Sats::ZERO,
            satdays_destroyed: Sats::ZERO,
            price_to_amount: PriceToAmount::forced_import(path, name),
        })
    }

    pub fn reset_single_iteration_values(&mut self) {
        self.satdays_destroyed = Sats::ZERO;
        self.satblocks_destroyed = Sats::ZERO;
        if let Some(realized) = self.realized.as_mut() {
            realized.reset_single_iteration_values();
        }
    }

    pub fn increment(&mut self, supply_state: &SupplyState, price: Option<Dollars>) {
        self.supply += supply_state;

        if supply_state.value > Sats::ZERO {
            if let Some(realized) = self.realized.as_mut() {
                let price = price.unwrap();
                realized.increment(supply_state, price);
                *self.price_to_amount.entry(price).or_default() += supply_state.value;
            }
        }
    }

    pub fn decrement(&mut self, supply_state: &SupplyState, price: Option<Dollars>) {
        self.supply -= supply_state;

        if supply_state.value > Sats::ZERO {
            if let Some(realized) = self.realized.as_mut() {
                let price = price.unwrap();
                realized.decrement(supply_state, price);
                self.decrement_price_to_amount(supply_state, price);
            }
        }
    }

    fn decrement_price_to_amount(&mut self, supply_state: &SupplyState, price: Dollars) {
        let amount = self.price_to_amount.get_mut(&price).unwrap();
        *amount -= supply_state.value;
        if *amount == Sats::ZERO {
            self.price_to_amount.remove(&price);
        }
    }

    pub fn receive(&mut self, supply_state: &SupplyState, price: Option<Dollars>) {
        self.supply += supply_state;

        if supply_state.value > Sats::ZERO {
            if let Some(realized) = self.realized.as_mut() {
                let price = price.unwrap();
                realized.receive(supply_state, price);
                *self.price_to_amount.entry(price).or_default() += supply_state.value;
            }
        }
    }

    pub fn send(
        &mut self,
        supply_state: &SupplyState,
        current_price: Option<Dollars>,
        prev_price: Option<Dollars>,
        blocks_old: usize,
        days_old: f64,
        older_than_hour: bool,
    ) {
        self.supply -= supply_state;

        if supply_state.value > Sats::ZERO {
            self.satblocks_destroyed += supply_state.value * blocks_old;

            self.satdays_destroyed +=
                Sats::from((u64::from(supply_state.value) as f64 * days_old).floor() as u64);

            if let Some(realized) = self.realized.as_mut() {
                let current_price = current_price.unwrap();
                let prev_price = prev_price.unwrap();
                realized.send(supply_state, current_price, prev_price, older_than_hour);
                self.decrement_price_to_amount(supply_state, prev_price);
            }
        }
    }

    pub fn compute_unrealized_states(
        &self,
        height_price: Dollars,
        date_price: Option<Dollars>,
    ) -> (UnrealizedState, Option<UnrealizedState>) {
        if self.price_to_amount.is_empty() {
            return (
                UnrealizedState::NAN,
                date_price.map(|_| UnrealizedState::NAN),
            );
        }

        let mut height_unrealized_state = UnrealizedState::ZERO;
        let mut date_unrealized_state = date_price.map(|_| UnrealizedState::ZERO);

        let update_state =
            |price: Dollars, current_price: Dollars, sats: Sats, state: &mut UnrealizedState| {
                match price.cmp(&current_price) {
                    Ordering::Equal => {
                        state.supply_even += sats;
                    }
                    Ordering::Less => {
                        state.supply_in_profit += sats;
                        if price > Dollars::ZERO && current_price > Dollars::ZERO {
                            let diff = current_price.checked_sub(price).unwrap();
                            if diff <= Dollars::ZERO {
                                dbg!(price, current_price, diff, sats);
                                panic!();
                            }
                            state.unrealized_profit += diff * sats;
                        }
                    }
                    Ordering::Greater => {
                        state.supply_in_loss += sats;
                        if price > Dollars::ZERO && current_price > Dollars::ZERO {
                            let diff = price.checked_sub(current_price).unwrap();
                            if diff <= Dollars::ZERO {
                                dbg!(price, current_price, diff, sats);
                                panic!();
                            }
                            state.unrealized_loss += diff * sats;
                        }
                    }
                }
            };

        self.price_to_amount.iter().for_each(|(&price, &sats)| {
            update_state(price, height_price, sats, &mut height_unrealized_state);

            if let Some(date_price) = date_price {
                update_state(
                    price,
                    date_price,
                    sats,
                    date_unrealized_state.as_mut().unwrap(),
                )
            }
        });

        (height_unrealized_state, date_unrealized_state)
    }

    pub fn commit(&mut self, height: Height) -> Result<()> {
        self.price_to_amount.flush(height)
    }
}