# region imports
from AlgorithmImports import *
# endregion

class MuscularYellowGreenCow(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2023, 1, 1)
        self.set_end_date(2024, 12, 1)
        self.set_cash(100000)
        
        # Add ES futures with hourly resolution
        self._es = self.add_future(Futures.Indices.SP_500_E_MINI, Resolution.HOUR)
        self._es.set_filter(0, 90)
        
        # Momentum indicators
        self._roc = None  # Rate of Change
        self._fast_sma = None
        self._slow_sma = None
        
        # Current contract
        self._contract = None
        
        # Position tracking
        self._entry_price = None
        self._position_side = None  # 1 for long, -1 for short
        
        # Momentum parameters
        self._roc_period = 20
        self._fast_period = 10
        self._slow_period = 30
        self._momentum_threshold = 0
    
    def on_securities_changed(self, changes):
        for security in changes.added_securities:
            # Initialize indicators for the new contract
            if security.symbol.security_type == SecurityType.FUTURE:
                self._contract = security.symbol
                self._roc = self.roc(self._contract, self._roc_period, Resolution.HOUR)
                self._fast_sma = self.sma(self._contract, self._fast_period, Resolution.HOUR)
                self._slow_sma = self.sma(self._contract, self._slow_period, Resolution.HOUR)
    
    def on_data(self, data: Slice):
        # Wait for contract to be set and indicators to be ready
        if self._contract is None or self._roc is None:
            return
        
        if not self._roc.is_ready or not self._fast_sma.is_ready or not self._slow_sma.is_ready:
            return
        
        # Skip if no data for current contract
        if not data.bars.contains_key(self._contract):
            return
        
        # Get current values
        roc_value = self._roc.current.value
        fast_sma = self._fast_sma.current.value
        slow_sma = self._slow_sma.current.value
        price = self.securities[self._contract].price
        
        # Momentum signals
        bullish_momentum = roc_value > self._momentum_threshold and fast_sma > slow_sma
        bearish_momentum = roc_value < -self._momentum_threshold and fast_sma < slow_sma
        
        # Current position
        holdings = self.portfolio[self._contract]
        
        # Entry logic
        if not holdings.invested:
            if bullish_momentum:
                # Go long
                quantity = self._calculate_position_size()
                self.market_order(self._contract, quantity)
                self._entry_price = price
                self._position_side = 1
                self.debug(f"Long entry at {price}, ROC: {roc_value:.2f}")
            
            elif bearish_momentum:
                # Go short
                quantity = self._calculate_position_size()
                self.market_order(self._contract, -quantity)
                self._entry_price = price
                self._position_side = -1
                self.debug(f"Short entry at {price}, ROC: {roc_value:.2f}")
        
        # Exit logic
        else:
            if self._entry_price is None:
                return
            
            # Calculate P&L in points
            if self._position_side == 1:  # Long position
                pnl_points = price - self._entry_price
                
                # Exit long on bearish momentum or stop loss
                if bearish_momentum or pnl_points < -20:
                    self.liquidate(self._contract)
                    self.debug(f"Long exit at {price}, P&L: {pnl_points:.2f} points")
                    self._entry_price = None
                    self._position_side = None
            
            elif self._position_side == -1:  # Short position
                pnl_points = self._entry_price - price
                
                # Exit short on bullish momentum or stop loss
                if bullish_momentum or pnl_points < -20:
                    self.liquidate(self._contract)
                    self.debug(f"Short exit at {price}, P&L: {pnl_points:.2f} points")
                    self._entry_price = None
                    self._position_side = None
    
    def _calculate_position_size(self):
        # Risk 2% of portfolio per trade
        # ES contract multiplier is 50
        risk_amount = self.portfolio.total_portfolio_value * 0.02
        stop_loss_points = 20
        quantity = int(risk_amount / (stop_loss_points * 50))
        return max(1, quantity)