# region imports
from AlgorithmImports import *
# endregion
class MomentumEffectAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2007,1,1)
        #self.set_end_date(2012,6,1)
        self.set_cash(100_000)
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)

        self.long_symbols = []
        self._selection_data_by_symbol = {}
        self._lookback = 201
        self._portfolio_size = 10

        self.settings.seed_initial_prices = True
        self.universe_settings.resolution = Resolution.DAILY
        self.settings.automatic_indicator_warm_up = True
        
        self.set_security_initializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.RAW))
        self._spy = self.add_equity("SPY", Resolution.DAILY)
        self._spy.sma = self.sma(self._spy, 200, Resolution.DAILY)
        self.universe_settings.schedule.on(self.date_rules.week_end(self._spy))
        self._universe = self.add_universe(self.universe.etf(self._spy), self._select_assets)
        
        self.set_warm_up(timedelta(201))
        self.schedule.on(self.date_rules.week_start('SPY'), self.time_rules.after_market_open('SPY', 1), self._rebalance)

    def _select_assets(self, fundamentals):
        
        universe_symbols = []
        
        filtered = [f for f in fundamentals if f.has_fundamental_data and f.price > 1]
        for f in filtered:
            universe_symbols.append(f.symbol)
            if f.symbol not in self._selection_data_by_symbol:
                self._selection_data_by_symbol[f.symbol] = SelectionData(self, f.symbol)
            self._selection_data_by_symbol[f.symbol].update(self.time, f.adjusted_price)

        symbols_to_remove = [s for s in self._selection_data_by_symbol.keys() if s not in universe_symbols]
        for symbol in symbols_to_remove:
            self._selection_data_by_symbol.pop(symbol)
        
        
        self.long_symbols = [
            symbol for symbol, selection_data in self._selection_data_by_symbol.items() if (selection_data.is_above_sma and selection_data.is_below_rsi)
        ]

        return self.long_symbols

    def on_securities_changed(self, changes):
        
        for security in changes.added_securities:
            security.momentum = self.momp(security, self._lookback, Resolution.DAILY)

        for security in changes.removed_securities:
            if security.Symbol in self._selection_data_by_symbol:
                self._selection_data_by_symbol.pop(security.Symbol)
            self.deregister_indicator(security.momentum)
            #self.liquidate(security, 'Removed from universe')

    def _rebalance(self):
        if self.is_warming_up:
            return
        if not self.long_symbols:
            return    
        if self._spy.price < self._spy.sma.current.value * 0.98:
            self.liquidate(tag = "Bear market")
            return

        targets = []
        securities_long = []
        for symbol in self.long_symbols:
            security = self.securities[symbol]
            # Check if indicator is ready
            if hasattr(security, 'momentum') and security.momentum.IsReady:
                securities_long.append(security)
        
        if securities_long:
            selected_long = sorted(securities_long, key=lambda x: x.momentum)[-self._portfolio_size:]
            long_weight = 0.99 / len(selected_long) 
            for sec in selected_long:
                targets.append(PortfolioTarget(sec.Symbol, long_weight))

        if targets:
            self.set_holdings(targets, True)

class SelectionData(object):

    def __init__(self, algorithm, symbol):
        self._sma = SimpleMovingAverage(200)
        self._rsi = RelativeStrengthIndex(3)
        self.is_above_sma = False
        self.is_below_rsi = False

        self.is_ready = False
        
        algorithm.warm_up_indicator(symbol, self._sma, Resolution.DAILY)
        algorithm.warm_up_indicator(symbol, self._rsi, Resolution.DAILY)
    
    def update(self, time, value):
        if self._sma.update(time, value) and self._rsi.update(time, value):
            self.is_above_sma = value > self._sma.current.value
            self.is_below_rsi = 50 > self._rsi.current.value