# region imports
from AlgorithmImports import *
# endregion

class SimpleDynamicMomentumAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2020, 1, 1)  # Set start date
        self.SetEndDate(2024, 1, 1)    # Set end date
        self.SetCash(100000)           # Set initial capital

        # Define the universe of assets
        self.symbols = [self.AddEquity(ticker, Resolution.Daily).Symbol for ticker in ["AAPL", "MSFT", "GOOGL", "AMZN", "META"]]

        self.lookback = 90  # Lookback period for momentum calculation (e.g., 3 months)
        self.rebalance_period = 30  # Rebalance period (e.g., monthly)
        self.next_rebalance = self.Time + timedelta(days=self.rebalance_period)
        
        self.stop_loss_percentage = 0.1  # 10% stop-loss

        self.entry_prices = {}  # Store the entry prices for positions

        # Market index to gauge overall market conditions
        self.market = self.AddEquity("SPY", Resolution.Daily).Symbol

        # Moving averages for market condition
        self.short_sma = self.SMA(self.market, 50, Resolution.Daily)
        self.long_sma = self.SMA(self.market, 200, Resolution.Daily)

        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(10, 0), self.Rebalance)

    def OnData(self, data):
        self.CheckStopLoss(data)

    def CheckStopLoss(self, data):
        for symbol in list(self.entry_prices.keys()):
            if symbol in data and data[symbol] is not None:
                if data[symbol].Price < self.entry_prices[symbol] * (1 - self.stop_loss_percentage):
                    self.Liquidate(symbol)
                    self.Debug(f"Stop-loss triggered for {symbol.Value} at {data[symbol].Price}")
                    del self.entry_prices[symbol]

    def Rebalance(self):
        if self.Time < self.next_rebalance:
            return

        # Determine market condition
        if self.short_sma.Current.Value > self.long_sma.Current.Value:
            # Bullish market condition
            long_weight = 0.8
            short_weight = 0.2
        else:
            # Bearish market condition
            long_weight = 0.2
            short_weight = 0.8

        # Calculate momentum for each stock in the universe
        momentum = {}
        for symbol in self.symbols:
            history = self.History(symbol, self.lookback, Resolution.Daily)
            if not history.empty:
                momentum[symbol] = history['close'].pct_change(self.lookback).iloc[-1]

        # Sort symbols based on momentum
        sorted_symbols = sorted(momentum.items(), key=lambda x: x[1], reverse=True)
        
        # Determine the number of long and short positions
        num_long = int(len(sorted_symbols) * long_weight)
        num_short = int(len(sorted_symbols) * short_weight)

        # Go long on top-ranked symbols
        long_symbols = [symbol for symbol, mom in sorted_symbols[:num_long]]
        # Go short on bottom-ranked symbols
        short_symbols = [symbol for symbol, mom in sorted_symbols[-num_short:]]

        # Set target holdings
        long_weight_per_position = long_weight / num_long if num_long > 0 else 0
        short_weight_per_position = short_weight / num_short if num_short > 0 else 0

        for symbol in self.symbols:
            if symbol in long_symbols:
                self.SetHoldings(symbol, long_weight_per_position)
                self.entry_prices[symbol] = self.Securities[symbol].Price  # Set entry price
            elif symbol in short_symbols:
                self.SetHoldings(symbol, -short_weight_per_position)
                self.entry_prices[symbol] = self.Securities[symbol].Price  # Set entry price
            else:
                self.Liquidate(symbol)
                if symbol in self.entry_prices:
                    del self.entry_prices[symbol]  # Remove entry price

        # Update next rebalance time
        self.next_rebalance = self.Time + timedelta(days=self.rebalance_period)

    def OnEndOfAlgorithm(self):
        self.Debug("Algorithm finished running.")