# region imports
from AlgorithmImports import *
# endregion

class JumpingMagentaKitten(QCAlgorithm):

    def Initialize(self) -> None:
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        self.SetEndDate(2021, 2, 1)    # Set End Date
        self.SetCash(100000)  # Set Strategy Cash

        [self.AddEquity(x) for x in ["AAPL", "MSFT"]]
        
        self.AddAlpha(CustomAlpha())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

class CustomAlpha(AlphaModel):

    def __init__(self):
        self.symbol_data = {}
    
    def Update(self, algorithm, data):
        insight = []

        for symbol, symbol_data in self.symbol_data.items():
            if not data.ContainsKey(symbol): continue

            sma = symbol_data.SMA
            if sma.IsReady and sma.Current.Value > data[symbol].Close:
                insight.append(Insight.Price(symbol, timedelta(minutes=1), InsightDirection.Up))

        return insight

    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            symbol = added.Symbol
            self.symbol_data[symbol] = SymbolData(algorithm, symbol)

        for removed in changes.RemovedSecurities:
            symbol_data = self.symbol_data.pop(removed.Symbol, None)
            if symbol_data:
                symbol_data.Dispose()

class SymbolData:

    def __init__(self, algo, symbol):
        self.algo = algo
        self.symbol = symbol

        # Let's say we have a 10-SMA strategy of 30-minute bar
        self.SMA = SimpleMovingAverage(10)
        
        self.consolidator = TradeBarConsolidator(timedelta(minutes=30))
        self.consolidator.DataConsolidated += self.ConsolidationHandler

        # Register the consolidator for automatic updates
        algo.SubscriptionManager.AddConsolidator(symbol, self.consolidator)

        # Warmup SMA indicator
        history = algo.History(symbol, 300, Resolution.Minute)
        for row in history.itertuples():
            trade_bar = TradeBar(row.Index[1], symbol, row.open, row.high, row.low, row.close, row.volume)
            self.consolidator.Update(trade_bar)

    def ConsolidationHandler(self, sender, bar):
        self.SMA.Update(bar.EndTime, bar.Close)

    # Dispose method to release memory from indicator and remove consolidator from consuming resources
    def Dispose(self):
        self.SMA.reset()
        self.algo.SubscriptionManager.RemoveConsolidator(self.symbol, self.consolidator)