vibe-invest/src/strategy.rs

//! Common trading strategy logic used by both the live bot and the backtester.
use std::collections::HashMap;
use crate::config::{
    get_sector, IndicatorParams, Timeframe, ATR_STOP_MULTIPLIER,
    ATR_TRAIL_ACTIVATION_MULTIPLIER, ATR_TRAIL_MULTIPLIER, MAX_LOSS_PCT, MAX_POSITION_SIZE,
    MIN_ATR_PCT, RISK_PER_TRADE, STOP_LOSS_PCT, TIME_EXIT_BARS,
    TRAILING_STOP_ACTIVATION, TRAILING_STOP_DISTANCE,
};
use crate::types::{Signal, TradeSignal};

/// Contains the core trading strategy logic.
pub struct Strategy {
    pub params: IndicatorParams,
    pub high_water_marks: HashMap<String, f64>,
    pub entry_atrs: HashMap<String, f64>,
    pub entry_prices: HashMap<String, f64>,
}

impl Strategy {
    pub fn new(timeframe: Timeframe) -> Self {
        Self {
            params: timeframe.params(),
            high_water_marks: HashMap::new(),
            entry_atrs: HashMap::new(),
            entry_prices: HashMap::new(),
        }
    }

    /// Volatility-adjusted position sizing using ATR (Kelly-inspired).
    ///
    /// Position size = (Risk per trade / ATR stop distance) * confidence.
    /// The confidence scaling now has a much wider range (0.4 to 1.0) so that
    /// weak Buy signals (confidence ~0.4) get 40% size while StrongBuy signals
    /// (confidence ~1.0) get full size. This is a fractional Kelly approach:
    /// bet more when conviction is higher, less when marginal.
    pub fn calculate_position_size(
        &self,
        price: f64,
        portfolio_value: f64,
        available_cash: f64,
        signal: &TradeSignal,
    ) -> f64 {
        if available_cash <= 0.0 {
            return 0.0;
        }

        let position_value = if signal.atr_pct > MIN_ATR_PCT {
            let atr_stop_pct = signal.atr_pct * ATR_STOP_MULTIPLIER;
            let risk_amount = portfolio_value * RISK_PER_TRADE;
            let vol_adjusted = risk_amount / atr_stop_pct;
            // Wide confidence scaling: 0.4x for weak signals, 1.0x for strongest.
            // Old code used 0.7 + 0.3*conf which barely differentiated.
            let confidence_scale = 0.4 + 0.6 * signal.confidence;
            let sized = vol_adjusted * confidence_scale;
            sized.min(portfolio_value * MAX_POSITION_SIZE)
        } else {
            portfolio_value * MAX_POSITION_SIZE
        };

        let position_value = position_value.min(available_cash);
        // Use fractional shares -- Alpaca supports them for paper trading.
        // Truncate to 4 decimal places to avoid floating point dust.
        ((position_value / price) * 10000.0).floor() / 10000.0
    }

    /// Check if stop-loss, trailing stop, or time exit should trigger.
    ///
    /// Exit priority (checked in order):
    /// 1. Hard max-loss cap (MAX_LOSS_PCT) -- absolute worst-case, gap protection
    /// 2. ATR-based stop-loss (ATR_STOP_MULTIPLIER * ATR) -- primary risk control
    /// 3. Fixed % stop-loss (STOP_LOSS_PCT) -- fallback when ATR unavailable
    /// 4. ATR trailing stop (ATR_TRAIL_MULTIPLIER * ATR from HWM) -- profit protection
    /// 5. Time-based exit (TIME_EXIT_BARS) -- only if position is LOSING
    ///
    /// Key design decisions:
    /// - Trailing stop activates early (1.5x ATR) but has wide distance (2.5x ATR)
    ///   so winners have room to breathe but profits are protected.
    /// - Time exit ONLY sells losers. Winners at the time limit are doing fine;
    ///   the trailing stop handles profit-taking on them.
    /// - Max loss is wide enough to avoid being hit by normal ATR-level moves.
    pub fn check_stop_loss_take_profit(
        &mut self,
        symbol: &str,
        current_price: f64,
        bars_held: usize,
    ) -> Option<Signal> {
        let entry_price = match self.entry_prices.get(symbol) {
            Some(&p) => p,
            None => return None,
        };

        let pnl_pct = (current_price - entry_price) / entry_price;
        let entry_atr = self.entry_atrs.get(symbol).copied().unwrap_or(0.0);

        // Update high water mark
        if let Some(hwm) = self.high_water_marks.get_mut(symbol) {
            if current_price > *hwm {
                *hwm = current_price;
            }
        }

        // 1. Hard max-loss cap (catastrophic gap protection)
        if pnl_pct <= -MAX_LOSS_PCT {
            return Some(Signal::StrongSell);
        }

        // 2. ATR-based initial stop-loss (primary risk control)
        if entry_atr > 0.0 {
            let atr_stop_price = entry_price - ATR_STOP_MULTIPLIER * entry_atr;
            if current_price <= atr_stop_price {
                return Some(Signal::StrongSell);
            }
        } else if pnl_pct <= -STOP_LOSS_PCT {
            // 3. Fixed percentage fallback
            return Some(Signal::StrongSell);
        }

        // 4. ATR-based trailing stop (profit protection)
        // Activates earlier than before (1.5x ATR gain) so profits are locked in.
        // Distance is wider (2.5x ATR from HWM) so normal retracements don't trigger it.
        let activation_gain = if entry_atr > 0.0 {
            (ATR_TRAIL_ACTIVATION_MULTIPLIER * entry_atr) / entry_price
        } else {
            TRAILING_STOP_ACTIVATION
        };

        if pnl_pct >= activation_gain {
            if let Some(&high_water) = self.high_water_marks.get(symbol) {
                let trail_distance = if entry_atr > 0.0 {
                    ATR_TRAIL_MULTIPLIER * entry_atr
                } else {
                    high_water * TRAILING_STOP_DISTANCE
                };
                let trailing_stop_price = high_water - trail_distance;
                if current_price <= trailing_stop_price {
                    return Some(Signal::Sell);
                }
            }
        }

        // 5. Time-based exit: only for LOSING positions (capital efficiency)
        // Winners at the time limit are managed by the trailing stop.
        // This prevents the old behavior of dumping winners just because they
        // haven't hit an arbitrary activation threshold in N bars.
        if bars_held >= TIME_EXIT_BARS && pnl_pct < 0.0 {
            return Some(Signal::Sell);
        }

        None
    }
    
    /// Count positions in a given sector.
    pub fn sector_position_count<'a, I>(&self, sector: &str, positions: I) -> usize
    where
        I: IntoIterator<Item = &'a String>,
    {
        positions
            .into_iter()
            .filter(|sym| get_sector(sym) == sector)
            .count()
    }
}