# region imports
from AlgorithmImports import *
# endregion

class MALongShortStrategy(QCAlgorithm):
    """Long-Short strategy on AAPL and MSFT using 20-day MA crossover"""

    def initialize(self):
        # Set start date, end date, and cash
        self.set_start_date(2018, 1, 1)
        self.set_end_date(2025, 10, 13)
        self.set_cash(100000)
        
        # Add equity data
        self.aapl = self.add_equity("AAPL", Resolution.DAILY).symbol
        self.msft = self.add_equity("MSFT", Resolution.DAILY).symbol
        
        # Create 20-day Simple Moving Average indicators
        self.aapl_sma = self.sma(self.aapl, 20, Resolution.DAILY)
        self.msft_sma = self.sma(self.msft, 20, Resolution.DAILY)
        
        # Warm up the indicators
        self.set_warm_up(20)
        
        # Track positions
        self.aapl_position = 0
        self.msft_position = 0
        
    def on_data(self, data: Slice):
        # Skip if we're still warming up
        if self.is_warming_up:
            return
            
        # Check if indicators are ready and we have data
        if not (self.aapl_sma.is_ready and self.msft_sma.is_ready):
            return
            
        if not (data.contains_key(self.aapl) and data.contains_key(self.msft)):
            return
        
        # Get current prices
        aapl_price = data[self.aapl].price
        msft_price = data[self.msft].price
        
        # Get MA values
        aapl_ma = self.aapl_sma.current.value
        msft_ma = self.msft_sma.current.value
        
        # AAPL Trading Logic
        # Long when price > MA, Short when price < MA
        if aapl_price > aapl_ma and self.aapl_position <= 0:
            # Go long AAPL (25% of portfolio)
            self.set_holdings(self.aapl, 0.25)
            self.aapl_position = 1
            self.debug(f"AAPL Long: Price {aapl_price:.2f} > MA {aapl_ma:.2f}")
            
        elif aapl_price < aapl_ma and self.aapl_position >= 0:
            # Go short AAPL (25% of portfolio)
            self.set_holdings(self.aapl, -0.25)
            self.aapl_position = -1
            self.debug(f"AAPL Short: Price {aapl_price:.2f} < MA {aapl_ma:.2f}")
        
        # MSFT Trading Logic
        # Long when price > MA, Short when price < MA
        if msft_price > msft_ma and self.msft_position <= 0:
            # Go long MSFT (25% of portfolio)
            self.set_holdings(self.msft, 0.25)
            self.msft_position = 1
            self.debug(f"MSFT Long: Price {msft_price:.2f} > MA {msft_ma:.2f}")
            
        elif msft_price < msft_ma and self.msft_position >= 0:
            # Go short MSFT (25% of portfolio)
            self.set_holdings(self.msft, -0.25)
            self.msft_position = -1
            self.debug(f"MSFT Short: Price {msft_price:.2f} < MA {msft_ma:.2f}")