from AlgorithmImports import *

class EnhancedEMACrossoverRSIAlgorithm(QCAlgorithm):

    def Initialize(self):
        # Set backtest dates
        self.SetStartDate(2023, 1, 1)
        self.SetEndDate(2024, 1, 31)  # Adjusted end date to a trading day

        # Set starting cash
        self.SetCash(100000)

        # Universe settings
        self.num_coarse = 100  # Reduced from 1000 to 100
        self.symbol_data = {}

        # Risk management parameters
        self.stop_loss_pct = 0.025  # 2.5% stop-loss
        self.take_profit_pct = 0.10  # 10% take-profit
        self.risk_per_trade = 0.01  # Risk 1% of portfolio per trade

        # Add SPY for always invested rule
        self.spy = self.AddEquity("SPY", Resolution.Daily).Symbol
        self.spy_entry_price = None  # Initialize entry price for SPY

        # Schedule universe selection
        self.AddUniverse(self.CoarseSelectionFunction)

        # Rebalance only on trading days
        self.Schedule.On(self.DateRules.EveryDay(self.spy), self.TimeRules.AfterMarketOpen(self.spy, 30), self.Rebalance)

        # Set universe settings
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.Leverage = 2
        self.UniverseSettings.FillForward = False

        # ATR for volatility filter
        self.volatility_threshold = 0.05  # 5% ATR threshold
        self.atr_period = 14  # ATR period

    def CoarseSelectionFunction(self, coarse):
        # Filter to top 100 liquid stocks
        sorted_coarse = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        selected = [x.Symbol for x in sorted_coarse[:self.num_coarse]]
        return selected

    def OnSecuritiesChanged(self, changes):
        # Add indicators for new securities
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            security.SetLeverage(2)  # Set leverage to 2
            if symbol not in self.symbol_data:
                ema5 = self.EMA(symbol, 5, Resolution.Daily)
                ema8 = self.EMA(symbol, 8, Resolution.Daily)
                ema13 = self.EMA(symbol, 13, Resolution.Daily)
                rsi = self.RSI(symbol, 14, MovingAverageType.Exponential, Resolution.Daily)
                atr = self.ATR(symbol, self.atr_period, MovingAverageType.Simple, Resolution.Daily)
                self.symbol_data[symbol] = {
                    'ema5': ema5,
                    'ema8': ema8,
                    'ema13': ema13,
                    'rsi': rsi,
                    'atr': atr,
                    'entry_price': None
                }

        # Remove data for removed securities
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.symbol_data:
                if self.Portfolio[symbol].Invested:
                    self.Liquidate(symbol)
                self.symbol_data.pop(symbol)

    def Rebalance(self):
        long_symbols = []
        short_symbols = []

        # Calculate signals
        for symbol, indicators in self.symbol_data.items():
            if not all([indicators['ema5'].IsReady, indicators['ema8'].IsReady, indicators['ema13'].IsReady,
                        indicators['rsi'].IsReady, indicators['atr'].IsReady]):
                continue

            ema5 = indicators['ema5'].Current.Value
            ema8 = indicators['ema8'].Current.Value
            ema13 = indicators['ema13'].Current.Value
            rsi = indicators['rsi'].Current.Value
            atr = indicators['atr'].Current.Value
            price = self.Securities[symbol].Price

            if price <= 0:
                continue

            # Volatility filter
            if atr / price > self.volatility_threshold:
                continue  # Skip if volatility is too high

            # Long condition
            if ema5 > ema8 > ema13 and rsi > 60:
                long_symbols.append(symbol)
            # Short condition
            elif ema5 < ema8 < ema13 and rsi < 40:
                short_symbols.append(symbol)

        # Ensure we're always invested
        total_symbols = long_symbols + short_symbols
        if not total_symbols:
            # If no signals, invest in SPY
            if not self.Portfolio[self.spy].Invested:
                self.SetHoldings(self.spy, 1)
                # Set entry price for SPY
                self.spy_entry_price = self.Securities[self.spy].Price
            return
        else:
            if self.Portfolio[self.spy].Invested:
                self.Liquidate(self.spy)
                self.spy_entry_price = None

        # Calculate target allocation per position
        total_positions = len(long_symbols) + len(short_symbols)
        if total_positions == 0:
            return

        # Leverage limit of 2
        target_percent = min(1.0, 2.0) / total_positions

        for symbol in long_symbols:
            self.EnterPosition(symbol, target_percent)

        for symbol in short_symbols:
            self.EnterPosition(symbol, -target_percent)

        # Liquidate positions no longer in signals
        invested_symbols = [x.Symbol for x in self.Portfolio.Values if x.Invested and x.Symbol != self.spy]
        for symbol in invested_symbols:
            if symbol not in total_symbols:
                self.Liquidate(symbol)
                if symbol in self.symbol_data:
                    self.symbol_data[symbol]['entry_price'] = None

    def EnterPosition(self, symbol, target_percent):
        # Set holdings based on target percent
        self.SetHoldings(symbol, target_percent)
        if symbol in self.symbol_data:
            self.symbol_data[symbol]['entry_price'] = self.Securities[symbol].Price

    def OnData(self, data):
        # Check stop-loss and take-profit for SPY
        if self.Portfolio[self.spy].Invested and data.ContainsKey(self.spy) and data[self.spy] is not None:
            price = data[self.spy].Close
            entry_price = self.spy_entry_price
            if entry_price is not None:
                pnl = (price - entry_price) / entry_price
                if pnl <= -self.stop_loss_pct:
                    # Stop-loss hit
                    self.Liquidate(self.spy)
                    self.spy_entry_price = None
                elif pnl >= self.take_profit_pct:
                    # Take-profit hit
                    self.Liquidate(self.spy)
                    self.spy_entry_price = None

        # Check stop-loss and take-profit for other symbols
        for symbol in list(self.symbol_data.keys()):
            if symbol not in data or not data[symbol]:
                continue
            if not self.Portfolio[symbol].Invested:
                continue
            indicators = self.symbol_data[symbol]
            entry_price = indicators.get('entry_price')
            if entry_price is None:
                continue  # Skip if entry_price is not set
            price = data[symbol].Close
            holdings = self.Portfolio[symbol]
            direction = 1 if holdings.IsLong else -1
            pnl = (price - entry_price) * direction / entry_price

            if pnl <= -self.stop_loss_pct:
                # Stop-loss hit
                self.Liquidate(symbol)
                self.symbol_data[symbol]['entry_price'] = None
            elif pnl >= self.take_profit_pct:
                # Take-profit hit
                self.Liquidate(symbol)
                self.symbol_data[symbol]['entry_price'] = None