apple-quant-algorithmic 0.3.0

use apple_quant_core::AddUnchecked;
use tracing::instrument;

use crate::{
	price::{AskPriceVolumeLevel, BidPriceVolumeLevel},
	aggregation::TradeTradeTimestamp, instrument::InstrumentSpec,
};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
struct Levels<IS: InstrumentSpec> {
	lowest_bid_level: BidPriceVolumeLevel<IS>,
	highest_ask_level: AskPriceVolumeLevel<IS>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct LiquidityEstimation<IS: InstrumentSpec>(Option<Levels<IS>>);

impl<IS: InstrumentSpec> LiquidityEstimation<IS> {
	pub fn reset(
		&mut self,
	) {
		self.0 = None;
	}

	pub fn collect_furthest_bid_ask(
		&mut self,
	) -> Option<(
		BidPriceVolumeLevel<IS>,
		AskPriceVolumeLevel<IS>,
	)> {
		let Levels {
			highest_ask_level,
			lowest_bid_level,
		} = self.0.take()?;

		Some((lowest_bid_level, highest_ask_level))
	}

	pub fn furthest_bid_ask(
		&self,
	) -> Option<(
		&BidPriceVolumeLevel<IS>,
		&AskPriceVolumeLevel<IS>,
	)> {
		let Some(Levels {
			lowest_bid_level,
			highest_ask_level,
		}) = &self.0 else {
			return None;
		};

		Some((lowest_bid_level, highest_ask_level))
	}

	#[instrument(skip_all)]
	pub fn walk_trades<'a>(
		&mut self,
		mut trades: impl Iterator<Item = &'a TradeTradeTimestamp<IS>>,
	)
	where
		IS: 'a,
	{
		if self.0.is_none() {
			let Some(
				first_trade_trade_timestamp,
			) = trades.next() else {
				return;
			};

			let highest_ask_level = AskPriceVolumeLevel {
				price: first_trade_trade_timestamp.trade().price,
				directionless_volume: *first_trade_trade_timestamp
					.trade()
					.aggressive_volume.as_directionless_volume(),
			};

			let lowest_bid_level = BidPriceVolumeLevel {
				price: highest_ask_level.price,
				directionless_volume: highest_ask_level.directionless_volume,
			};

			self.0 = Some(Levels {
				highest_ask_level,
				lowest_bid_level,
			});
		}

		let levels = self.0.as_mut().unwrap();

		for trade_trade_timestamp in trades {
			let trade_price = &trade_trade_timestamp.trade().price;

			let trade_aggressive_volume = &trade_trade_timestamp.trade().aggressive_volume;

			if trade_price == &levels.highest_ask_level.price {
				levels.highest_ask_level.directionless_volume = levels.highest_ask_level.directionless_volume.add_unchecked(
					*trade_aggressive_volume.as_directionless_volume(),
				);
			} else if trade_price > &levels.highest_ask_level.price {
				levels.highest_ask_level.price = *trade_price;

				levels.highest_ask_level.directionless_volume = *trade_aggressive_volume.as_directionless_volume();
			}

			if trade_price == &levels.lowest_bid_level.price {
				levels.lowest_bid_level.directionless_volume = levels.lowest_bid_level.directionless_volume.add_unchecked(
					*trade_aggressive_volume.as_directionless_volume(),
				);
			} else if trade_price < &levels.lowest_bid_level.price {
				levels.lowest_bid_level.price = *trade_price;

				levels.lowest_bid_level.directionless_volume = *trade_aggressive_volume.as_directionless_volume();
			}
		}
	}
}

impl<IS: InstrumentSpec> Default for LiquidityEstimation<IS> {
	fn default() -> Self {
		Self(None)
	}
}