from AlgorithmImports import *

class MarketNeutralPennyStocksAlgorithm(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2023, 9, 1)
        self.set_end_date(2024, 5, 1)
        self.SetCash(10_000)
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.Margin)
        self.SetSecurityInitializer(BrokerageModelSecurityInitializer(
            self.BrokerageModel, FuncSecuritySeeder(self.GetLastKnownPrices)
        ))

        self.UniverseSettings.Resolution = Resolution.Daily

        self.momentum_indicators = {}
        self.lookback_period = 252
        self.num_coarse = 300
        self.num_fine = 300
        self.num_long = 10
        self.num_short = 10

        self.current_month = 6
        self.rebalance_flag = True

        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

    def CoarseSelectionFunction(self, coarse):
        if self.current_month == self.Time.month:
            return Universe.Unchanged

        self.rebalance_flag = True
        self.current_month = self.Time.month

        selected = sorted(
            [x for x in coarse if x.HasFundamentalData and x.Price < 1 and x.MarketCap < 1e9],
            key=lambda x: x.DollarVolume, 
            reverse=True
        )
        self.Debug(f"Coarse selected: {len(selected)} securities")
        return [x.Symbol for x in selected[:self.num_coarse]]

    def FineSelectionFunction(self, fine):
        selected = sorted(fine, key=lambda f: f.MarketCap, reverse=True)
        self.Debug(f"Fine selected: {len(selected)} securities")
        return [x.Symbol for x in selected[:self.num_fine]]

    def OnData(self, data):
        for symbol, momentum in self.momentum_indicators.items():
            if symbol in data and data[symbol] is not None:
                momentum.Update(self.Time, data[symbol].Close)

        if not self.rebalance_flag:
            return

        sorted_momentum = sorted(
            [k for k, v in self.momentum_indicators.items() if v.IsReady],
            key=lambda x: self.momentum_indicators[x].Current.Value, 
            reverse=True
        )
        self.Debug(f"Sorted momentum: {len(sorted_momentum)} securities")
        long_selected = sorted_momentum[:self.num_long]
        self.Debug(f"LONG {long_selected}")
        short_selected = sorted_momentum[-self.num_short:]
        self.Debug(f"SHORT {short_selected}")

        for symbol in list(self.Portfolio.Keys):
            if symbol not in long_selected and symbol not in short_selected:
                self.Liquidate(symbol, 'Not selected')

        initial_investment = 500
        additional_investment = 250

        for symbol in long_selected:
            current_investment = self.Portfolio[symbol].Invested
            if current_investment:
                self.MarketOrder(symbol, additional_investment / data[symbol].Close)
                stop_price = data[symbol].Close * 1.18
                self.stop_limit_order(symbol, self.Portfolio[symbol].Quantity , stop_price, stop_price)
                self.Debug(f"Long {symbol} Price {data[symbol].Close} Qty {int(additional_investment / data[symbol].Close)} Stop price {stop_price}")
            else:
                self.MarketOrder(symbol, initial_investment / data[symbol].Close)
                stop_price = data[symbol].Close * 1.18
                self.stop_limit_order(symbol, self.Portfolio[symbol].Quantity , stop_price, stop_price)
                self.Debug(f"Long {symbol} Price {data[symbol].Close} Qty {int(initial_investment / data[symbol].Close)} Stop price {stop_price}")

        for symbol in short_selected:
            current_investment = self.Portfolio[symbol].Invested
            if current_investment:
                if self.securities[symbol].is_tradable:
                    self.MarketOrder(symbol, - int(additional_investment / data[symbol].Close))
                    stop_price = data[symbol].Close * 1.02
                    self.stop_limit_order(symbol, self.Portfolio[symbol].Quantity , stop_price, stop_price)
                    self.Debug(f"Short {symbol} Price {data[symbol].Close} Qty {- int(additional_investment / data[symbol].Close)} Stop price {stop_price}")
            else:
                if self.securities[symbol].is_tradable:
                    self.MarketOrder(symbol, - int(initial_investment / data[symbol].Close))
                    stop_price = data[symbol].Close * 1.02
                    self.stop_limit_order(symbol, self.Portfolio[symbol].Quantity , stop_price, stop_price)
                    self.Debug(f"Short {symbol} Price {data[symbol].Close} Qty {- int(initial_investment / data[symbol].Close)} Stop price {stop_price}")
        
        self.rebalance_flag = False

    def OnSecuritiesChanged(self, changes):
        for security in changes.RemovedSecurities:
            symbol = security.Symbol
            if symbol in self.momentum_indicators:
                self.momentum_indicators.pop(symbol)
                self.Liquidate(symbol, 'Removed from universe')

        for security in changes.AddedSecurities:
            symbol = security.Symbol
            if symbol not in self.momentum_indicators:
                self.momentum_indicators[symbol] = MomentumPercent(self.lookback_period)
                history = self.History(symbol, self.lookback_period, Resolution.Daily)
                
                # Check if history data is available for the symbol
                if not history.empty and symbol in history.index.levels[0]:
                    for time, row in history.loc[symbol].iterrows():
                        self.momentum_indicators[symbol].Update(time, row['close'])
                else:
                    self.Debug(f"No historical data available for symbol {symbol}")