stock-trek 0.8.9

Stock Trek time-series analysis
Documentation
use crate::{
    allocations::allocation::{Allocation, AllocationTrait},
    cex::{asset_id::AssetId, cex_id::CexId},
};
use std::collections::HashMap;

#[derive(Debug, Clone, Default)]
pub struct InMemoryAllocation {
    cex_assets: HashMap<CexId, Allocations>,
}
impl InMemoryAllocation {
    pub fn new(cex_assets: HashMap<CexId, Allocations>) -> Self {
        Self { cex_assets }
    }
    pub fn builder() -> Builder {
        Builder::new()
    }
}

impl From<InMemoryAllocation> for Allocation {
    fn from(value: InMemoryAllocation) -> Self {
        Box::new(value)
    }
}

#[derive(Debug, Clone, Default)]
pub struct Allocations {
    asset_allocations: HashMap<AssetId, f64>,
}
impl Allocations {
    pub fn new(asset_allocations: HashMap<AssetId, f64>) -> Self {
        Self { asset_allocations }
    }
}

impl AllocationTrait for InMemoryAllocation {
    fn allocation_for_asset_total(&self, asset_id: &AssetId) -> f64 {
        self.cex_assets
            .values()
            .map(|assets| assets.asset_allocations.get(asset_id).unwrap_or(&0.0))
            .sum()
    }
    fn allocation_for_asset_in_cex(&self, asset_id: &AssetId, cex_id: &CexId) -> f64 {
        self.cex_assets
            .get(cex_id)
            .and_then(|assets| assets.asset_allocations.get(asset_id))
            .copied()
            .unwrap_or(0.0)
    }
}

#[derive(Clone, Default)]
pub struct Builder {
    cex_allocations: HashMap<CexId, Allocations>,
}

impl Builder {
    pub fn new() -> Self {
        Self {
            cex_allocations: HashMap::new(),
        }
    }
    pub fn allocation(&mut self, cex_id: CexId, asset_id: AssetId, quantity: f64) -> &mut Self {
        assert!(quantity > 0.0, "allocation must be greater than 0.0");
        assert!(
            quantity <= 100.0,
            "allocation must be less or equal to 100.0"
        );
        self.cex_allocations
            .entry(cex_id)
            .or_insert_with(|| Allocations::new(HashMap::new()))
            .asset_allocations
            .entry(asset_id)
            .and_modify(|prev| *prev += quantity)
            .or_insert(quantity);
        self
    }
    pub fn build(&self) -> InMemoryAllocation {
        InMemoryAllocation::new(self.cex_allocations.clone())
    }
}