use std::collections::HashMap;

use bt_kernel::SymbolId;

#[derive(Clone, Debug)]
pub struct PortfolioState {
    pub capital: f64,
    pub positions: HashMap<SymbolId, f64>,
}

impl PortfolioState {
    pub fn new(initial_capital: f64) -> Self {
        Self {
            capital: initial_capital,
            positions: HashMap::new(),
        }
    }

    pub fn position(&self, symbol: SymbolId) -> f64 {
        self.positions.get(&symbol).copied().unwrap_or(0.0)
    }

    pub fn set_position(&mut self, symbol: SymbolId, qty: f64) {
        if qty.abs() < 1e-15 {
            self.positions.remove(&symbol);
        } else {
            self.positions.insert(symbol, qty);
        }
    }

    pub fn equity(&self, close_prices: &HashMap<SymbolId, f64>) -> f64 {
        self.capital
            + self
                .positions
                .iter()
                .map(|(s, pos)| pos * close_prices.get(s).copied().unwrap_or(0.0))
                .sum::<f64>()
    }
}

#[derive(Clone, Debug)]
pub struct PortfolioTrace {
    pub symbol_positions: HashMap<SymbolId, Vec<f64>>,
    pub capitals: Vec<f64>,
    pub equities: Vec<f64>,
}

impl PortfolioTrace {
    pub fn new(symbols: &[SymbolId]) -> Self {
        let mut symbol_positions = HashMap::new();
        for &s in symbols {
            symbol_positions.insert(s, Vec::new());
        }
        Self {
            symbol_positions,
            capitals: Vec::new(),
            equities: Vec::new(),
        }
    }

    pub fn push(&mut self, state: &PortfolioState, equity: f64) {
        for (symbol, positions) in &mut self.symbol_positions {
            positions.push(state.position(*symbol));
        }
        self.capitals.push(state.capital);
        self.equities.push(equity);
    }
}
