Skip to main content

trader_bot/domain/
wallet.rs

1//! Wallet and portfolio types.
2
3use rust_decimal::Decimal;
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7use super::{Money, Position};
8
9/// Represents a desired portfolio allocation.
10///
11/// The keys are asset symbols and values are target percentages (0-100).
12#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
13pub struct DesiredAllocation {
14    allocations: HashMap<String, Decimal>,
15}
16
17impl DesiredAllocation {
18    /// Creates a new empty allocation.
19    #[must_use]
20    pub fn new() -> Self {
21        Self::default()
22    }
23
24    /// Creates an allocation from a HashMap.
25    #[must_use]
26    pub fn from_map(allocations: HashMap<String, Decimal>) -> Self {
27        Self { allocations }
28    }
29
30    /// Sets the target percentage for a symbol.
31    pub fn set(&mut self, symbol: impl Into<String>, percentage: Decimal) {
32        self.allocations.insert(symbol.into(), percentage);
33    }
34
35    /// Gets the target percentage for a symbol.
36    #[must_use]
37    pub fn get(&self, symbol: &str) -> Option<Decimal> {
38        self.allocations.get(symbol).copied()
39    }
40
41    /// Returns all symbols in the allocation.
42    pub fn symbols(&self) -> impl Iterator<Item = &String> {
43        self.allocations.keys()
44    }
45
46    /// Returns an iterator over (symbol, percentage) pairs.
47    pub fn iter(&self) -> impl Iterator<Item = (&String, &Decimal)> {
48        self.allocations.iter()
49    }
50
51    /// Returns the number of assets in the allocation.
52    #[must_use]
53    pub fn len(&self) -> usize {
54        self.allocations.len()
55    }
56
57    /// Returns true if the allocation is empty.
58    #[must_use]
59    pub fn is_empty(&self) -> bool {
60        self.allocations.is_empty()
61    }
62
63    /// Calculates the total percentage (should sum to 100 for a valid allocation).
64    #[must_use]
65    pub fn total_percentage(&self) -> Decimal {
66        self.allocations.values().sum()
67    }
68
69    /// Validates that the allocation sums to approximately 100%.
70    #[must_use]
71    pub fn is_valid(&self) -> bool {
72        let total = self.total_percentage();
73        let diff = (total - Decimal::from(100)).abs();
74        diff < Decimal::new(1, 2) // Allow 0.01% tolerance
75    }
76
77    /// Normalizes the allocation to sum to exactly 100%.
78    #[must_use]
79    pub fn normalized(&self) -> Self {
80        let total = self.total_percentage();
81        if total.is_zero() {
82            return self.clone();
83        }
84
85        let normalized: HashMap<String, Decimal> = self
86            .allocations
87            .iter()
88            .map(|(k, v)| (k.clone(), *v * Decimal::from(100) / total))
89            .collect();
90
91        Self {
92            allocations: normalized,
93        }
94    }
95}
96
97/// Represents a wallet (portfolio) containing multiple positions.
98#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
99pub struct Wallet {
100    /// The base currency for value calculations.
101    base_currency: String,
102    /// Cash balance.
103    cash: Money,
104    /// Positions in the wallet.
105    positions: Vec<Position>,
106}
107
108impl Wallet {
109    /// Creates a new empty wallet.
110    #[must_use]
111    pub fn new(base_currency: impl Into<String>) -> Self {
112        let currency: String = base_currency.into();
113        Self {
114            cash: Money::zero(&currency),
115            base_currency: currency,
116            positions: Vec::new(),
117        }
118    }
119
120    /// Creates a wallet with initial cash balance.
121    #[must_use]
122    pub fn with_cash(base_currency: impl Into<String>, cash: Money) -> Self {
123        let currency: String = base_currency.into();
124        Self {
125            cash,
126            base_currency: currency,
127            positions: Vec::new(),
128        }
129    }
130
131    /// Returns the base currency.
132    #[must_use]
133    pub fn base_currency(&self) -> &str {
134        &self.base_currency
135    }
136
137    /// Returns the cash balance.
138    #[must_use]
139    pub const fn cash(&self) -> &Money {
140        &self.cash
141    }
142
143    /// Returns all positions.
144    #[must_use]
145    pub fn positions(&self) -> &[Position] {
146        &self.positions
147    }
148
149    /// Returns a mutable reference to positions.
150    pub fn positions_mut(&mut self) -> &mut Vec<Position> {
151        &mut self.positions
152    }
153
154    /// Adds a position to the wallet.
155    pub fn add_position(&mut self, position: Position) {
156        // Check if position with same symbol already exists
157        if let Some(existing) = self
158            .positions
159            .iter_mut()
160            .find(|p| p.symbol() == position.symbol())
161        {
162            // Update existing position
163            let new_qty = existing.quantity() + position.quantity();
164            existing.update_quantity(new_qty);
165            existing.update_price(position.current_price().clone());
166        } else {
167            self.positions.push(position);
168        }
169    }
170
171    /// Removes a position by symbol.
172    pub fn remove_position(&mut self, symbol: &str) -> Option<Position> {
173        if let Some(idx) = self.positions.iter().position(|p| p.symbol() == symbol) {
174            Some(self.positions.remove(idx))
175        } else {
176            None
177        }
178    }
179
180    /// Gets a position by symbol.
181    #[must_use]
182    pub fn get_position(&self, symbol: &str) -> Option<&Position> {
183        self.positions.iter().find(|p| p.symbol() == symbol)
184    }
185
186    /// Gets a mutable position by symbol.
187    pub fn get_position_mut(&mut self, symbol: &str) -> Option<&mut Position> {
188        self.positions.iter_mut().find(|p| p.symbol() == symbol)
189    }
190
191    /// Updates the cash balance.
192    pub fn set_cash(&mut self, cash: Money) {
193        self.cash = cash;
194    }
195
196    /// Adds to the cash balance.
197    pub fn add_cash(&mut self, amount: Money) {
198        if let Some(new_cash) = self.cash.checked_add(&amount) {
199            self.cash = new_cash;
200        }
201    }
202
203    /// Subtracts from the cash balance.
204    pub fn subtract_cash(&mut self, amount: Money) {
205        if let Some(new_cash) = self.cash.checked_sub(&amount) {
206            self.cash = new_cash;
207        }
208    }
209
210    /// Calculates the total market value of all positions.
211    #[must_use]
212    pub fn total_positions_value(&self) -> Money {
213        let total: Decimal = self
214            .positions
215            .iter()
216            .map(|p| p.market_value().amount())
217            .sum();
218        Money::new(total, &self.base_currency)
219    }
220
221    /// Calculates the total portfolio value (cash + positions).
222    #[must_use]
223    pub fn total_value(&self) -> Money {
224        let positions_value = self.total_positions_value();
225        self.cash
226            .checked_add(&positions_value)
227            .unwrap_or_else(|| positions_value)
228    }
229
230    /// Calculates the current allocation percentages.
231    #[must_use]
232    pub fn current_allocation(&self) -> DesiredAllocation {
233        let total = self.total_value().amount();
234        if total.is_zero() {
235            return DesiredAllocation::new();
236        }
237
238        let allocations: HashMap<String, Decimal> = self
239            .positions
240            .iter()
241            .map(|p| {
242                let pct = p.market_value().amount() / total * Decimal::from(100);
243                (p.symbol().to_string(), pct)
244            })
245            .collect();
246
247        DesiredAllocation::from_map(allocations)
248    }
249
250    /// Calculates the cash percentage of total portfolio.
251    #[must_use]
252    pub fn cash_percentage(&self) -> Decimal {
253        let total = self.total_value().amount();
254        if total.is_zero() {
255            return Decimal::ZERO;
256        }
257        self.cash.amount() / total * Decimal::from(100)
258    }
259
260    /// Returns the number of positions.
261    #[must_use]
262    pub fn position_count(&self) -> usize {
263        self.positions.len()
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270    use rust_decimal_macros::dec;
271
272    fn make_usd(amount: Decimal) -> Money {
273        Money::new(amount, "USD")
274    }
275
276    #[test]
277    fn test_desired_allocation() {
278        let mut alloc = DesiredAllocation::new();
279        alloc.set("AAPL", dec!(30));
280        alloc.set("GOOGL", dec!(30));
281        alloc.set("MSFT", dec!(40));
282
283        assert_eq!(alloc.get("AAPL"), Some(dec!(30)));
284        assert_eq!(alloc.total_percentage(), dec!(100));
285        assert!(alloc.is_valid());
286    }
287
288    #[test]
289    fn test_allocation_normalization() {
290        let mut alloc = DesiredAllocation::new();
291        alloc.set("A", dec!(10));
292        alloc.set("B", dec!(20));
293        alloc.set("C", dec!(30)); // Total = 60
294
295        let normalized = alloc.normalized();
296        let total = normalized.total_percentage();
297        assert!((total - dec!(100)).abs() < dec!(0.01));
298    }
299
300    #[test]
301    fn test_wallet_creation() {
302        let wallet = Wallet::with_cash("USD", make_usd(dec!(10000)));
303        assert_eq!(wallet.base_currency(), "USD");
304        assert_eq!(wallet.cash().amount(), dec!(10000));
305        assert_eq!(wallet.positions().len(), 0);
306    }
307
308    #[test]
309    fn test_wallet_add_position() {
310        let mut wallet = Wallet::with_cash("USD", make_usd(dec!(5000)));
311        let pos = Position::new("AAPL", "NASDAQ", dec!(10), make_usd(dec!(150)));
312        wallet.add_position(pos);
313
314        assert_eq!(wallet.position_count(), 1);
315        assert_eq!(wallet.get_position("AAPL").unwrap().quantity(), dec!(10));
316    }
317
318    #[test]
319    fn test_wallet_total_value() {
320        let mut wallet = Wallet::with_cash("USD", make_usd(dec!(5000)));
321        wallet.add_position(Position::new(
322            "AAPL",
323            "NASDAQ",
324            dec!(10),
325            make_usd(dec!(150)),
326        ));
327        wallet.add_position(Position::new(
328            "GOOGL",
329            "NASDAQ",
330            dec!(5),
331            make_usd(dec!(100)),
332        ));
333
334        // Cash: 5000, AAPL: 10*150=1500, GOOGL: 5*100=500
335        // Total = 7000
336        assert_eq!(wallet.total_value().amount(), dec!(7000));
337    }
338
339    #[test]
340    fn test_wallet_current_allocation() {
341        let mut wallet = Wallet::with_cash("USD", make_usd(dec!(0)));
342        wallet.add_position(Position::new(
343            "AAPL",
344            "NASDAQ",
345            dec!(10),
346            make_usd(dec!(100)),
347        ));
348        wallet.add_position(Position::new(
349            "GOOGL",
350            "NASDAQ",
351            dec!(10),
352            make_usd(dec!(100)),
353        ));
354
355        let alloc = wallet.current_allocation();
356        assert_eq!(alloc.get("AAPL"), Some(dec!(50)));
357        assert_eq!(alloc.get("GOOGL"), Some(dec!(50)));
358    }
359
360    #[test]
361    fn test_wallet_merge_positions() {
362        let mut wallet = Wallet::new("USD");
363        wallet.add_position(Position::new(
364            "AAPL",
365            "NASDAQ",
366            dec!(10),
367            make_usd(dec!(150)),
368        ));
369        wallet.add_position(Position::new(
370            "AAPL",
371            "NASDAQ",
372            dec!(5),
373            make_usd(dec!(155)),
374        ));
375
376        assert_eq!(wallet.position_count(), 1);
377        assert_eq!(wallet.get_position("AAPL").unwrap().quantity(), dec!(15));
378    }
379}