from QuantConnect.Indicators import RelativeStrengthIndex

class TransdimensionalOptimizedCoil(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 12, 21)
        self.SetEndDate(2020, 1, 1)
        self.SetCash(100000)
        
        self.AddAlpha(MyAlphaModel())
        
        symbols = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        self.UniverseSettings.Resolution = Resolution.Minute
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        
        self.SetExecution(ImmediateExecutionModel())
        self.SetWarmUp(240)


class MyAlphaModel(AlphaModel):
    symbol_data_by_symbol = {}
    rsi_length = 14
    
    def Update(self, algorithm, data):
        if algorithm.IsWarmingUp:
            return []
        
        if data.Time.hour == 9 and data.Time.minute == 31:
            for symbol in self.symbol_data_by_symbol:
                algorithm.Log(f"RSI: {self.symbol_data_by_symbol[symbol].rsi}")
                
        return []
    
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            self.symbol_data_by_symbol[added.Symbol] = SymbolData(added.Symbol, algorithm, self.rsi_length)
            
        for removed in changes.RemovedSecurities:
            symbol_data = self.symbol_data_by_symbol.pop(removed.Symbol, None)
            if symbol_data:
                symbol_data.dispose()
        
        
class SymbolData:
    def __init__(self, symbol, algorithm, rsi_length):
        self.symbol = symbol
        self.algorithm = algorithm 
        
        self.rsi = RelativeStrengthIndex(rsi_length, MovingAverageType.Simple)
        
        # Warm up RSI
        history = algorithm.History(symbol, rsi_length, Resolution.Daily).loc[symbol]
        for idx, row in history.iterrows():
            self.rsi.Update(idx, row.close)
        
        # Setup daily indicator consolidator
        self.consolidator = TradeBarConsolidator(timedelta(1))
        self.consolidator.DataConsolidated += self.CustomDailyHandler
        algorithm.SubscriptionManager.AddConsolidator(self.symbol, self.consolidator)
    
    def CustomDailyHandler(self, sender, consolidated):
        self.rsi.Update(consolidated.Time, consolidated.Close)
        
    def dispose(self):
        self.algorithm.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)