from System.Drawing import Color

class JumpingFluorescentOrangeGaur(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 4, 23)  # Set Start Date
        self.SetEndDate(2021, 4, 23)
        self.SetCash(100000)  # Set Strategy Cash
        self.UniverseSettings.ExtendedMarketHours = True
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddEquity('SPY', Resolution.Minute, Market.USA, True, 1, True).Symbol
        self.tradingSymbols = []
        
        self.SetAlpha(TradingModel(self.tradingSymbols))
        #self.tradingModel = TradingModel(self.tradingSymbols)
        
        self.Schedule.On(
            self.DateRules.EveryDay(),
            self.TimeRules.At(6, 50),
            self.addSymbols
            )

    def addSymbols(self):
        self.aapl = self.AddEquity('AAPL', Resolution.Minute, Market.USA, True, 1, True).Symbol
        self.tradingSymbols.append(self.aapl)
        #return TradingModel(self.tradingSymbols)
        
    def OnData(self, data):
        pass
    
class TradingModel(AlphaModel):
    
    def __init__(self, universe):
        self.tradeSymbols = universe
        self.symbolDataBySymbol = {}
        self.startIndicators = {}
        self.allIndicators = {}
        self.count = 0
    
    def Update(self, algorithm, data):
        insights = []
        for symbol in self.tradeSymbols:
            if symbol in self.symbolDataBySymbol.keys():
                symbolData = self.symbolDataBySymbol[symbol]
                if data.Bars.ContainsKey(symbol):
                    algorithm.Plot('Trade Plot', 'Price', data.Bars[symbol].Close)
                    algorithm.PlotIndicator('Trade Plot', symbolData.slowFiveEMA, Identity(symbol))
                
        '''
        self.count += 1
        if self.count == 120:
            for symbol, indicators in self.allIndicators.items():
                    self.indicatorValues = [indicator.Current.Value for indicator in indicators]
                    algorithm.Log(f"Our indicator values for {symbol} are: {self.indicatorValues}")
                    '''
        return insights

        
    def OnSecuritiesChanged(self, algorithm, changes):
        self.changes = changes
        self.resolution = Resolution.Minute
        self.fastPeriod = 20
        self.slowPeriod = 200
        self.barPeriod = TimeSpan.FromMinutes(1)
        self.rollingWindowSize = 6
        
        symbols = [added.Symbol for added in changes.AddedSecurities if added.Symbol not in self.symbolDataBySymbol]
        #symbolsValues = [symbol.Value for symbol in symbols]
        #algorithm.Log(f"Our symbols are: {symbolsValues}")
        historyFastfive = algorithm.History(symbols, self.fastPeriod * 10, self.resolution)
        algorithm.Debug(f"historyFastfive first value for {symbols}: {historyFastfive['close'].head(-20)}")
        if historyFastfive.empty:
            algorithm.Log(f"empty")
            return
        historySlowfive = algorithm.History(symbols, self.slowPeriod * 5, self.resolution)
        algorithm.Debug(f"historySlowfive first value for {symbols}: {historySlowfive['close'].head(-20)}")
        if historySlowfive.empty:
            return
        historySlowtwo = algorithm.History(symbols, self.slowPeriod * 5, self.resolution)
        algorithm.Debug(f"historySlowtwo first value for {symbols}: {historySlowtwo['close'][0]}")
        if historySlowtwo.empty:
            return
        historySlowone = algorithm.History(symbols, self.slowPeriod * 5, self.resolution)
        algorithm.Debug(f"historySlowone first value for {symbols}: {historySlowone['close'][0]}")
        if historySlowone.empty:
            return
    
        tickers = historyFastfive.index.levels[0]
        for ticker in tickers:
            symbol = SymbolCache.GetSymbol(ticker)
            symbolData = SymbolData(symbol, algorithm)
            symbolData.fastFiveEMA = ExponentialMovingAverage(self.fastPeriod)
            symbolData.slowFiveEMA = ExponentialMovingAverage(self.slowPeriod)
            symbolData.slowFiveSMA = SimpleMovingAverage(self.slowPeriod)
            symbolData.slowTwoEMA = ExponentialMovingAverage(self.slowPeriod)
            symbolData.slowTwoSMA = SimpleMovingAverage(self.slowPeriod)
            symbolData.slowOneEMA = ExponentialMovingAverage(self.slowPeriod)
            symbolData.slowOneSMA = SimpleMovingAverage(self.slowPeriod)
            symbolData.VWAP = algorithm.VWAP(symbol)
            symbolData.RegisterIndicator(algorithm)
            symbolData.WarmUpIndicators(algorithm, historyFastfive.loc[ticker], \
                                        historySlowfive.loc[ticker], historySlowtwo.loc[ticker], \
                                        historySlowone.loc[ticker])
            self.symbolDataBySymbol[symbol] = symbolData
            for symbol, symbolData in self.symbolDataBySymbol.items():
                self.allIndicators[symbol] = [
                    symbolData.fastFiveEMA, symbolData.slowFiveEMA, 
                    symbolData.slowFiveSMA, symbolData.slowTwoEMA, 
                    symbolData.slowOneEMA, symbolData.slowTwoSMA,
                    symbolData.slowOneSMA, symbolData.VWAP
                    ]

class SymbolData:
    
    def __init__(self, symbol, algorithm):
        
        self.Symbol = symbol
        self.algorithm = algorithm
        self.resolution = Resolution.Minute
        self.slowPeriod = 200
        self.fastPeriod = 20
        self.fastFiveEMA = None
        self.slowFiveEMA = None
        self.slowFiveSMA = None
        self.slowTwoEMA = None
        self.slowTwoSMA = None
        self.slowOneEMA = None
        self.slowOneSMA = None
        self.VWAP = None
        
        self.fiveMinuteConsolidator = TradeBarConsolidator(timedelta(minutes = 5))
        self.twoMinuteConsolidator = TradeBarConsolidator(timedelta(minutes = 2))
        self.oneMinuteConsolidator = TradeBarConsolidator(timedelta(minutes = 1))
        
        #self.fiveMinuteConsolidator.DataConsolidated += self.OnDataConsolidated
        #self.twoMinuteConsolidator.DataConsolidated += self.OnDataConsolidated
        #self.oneMinuteConsolidator.DataConsolidated += self.OnDataConsolidated
        
        algorithm.SubscriptionManager.AddConsolidator(self.Symbol, self.fiveMinuteConsolidator)
        algorithm.SubscriptionManager.AddConsolidator(self.Symbol, self.twoMinuteConsolidator)
        algorithm.SubscriptionManager.AddConsolidator(self.Symbol, self.oneMinuteConsolidator)
        
    def WarmUpIndicators(self, algorithm, historyFastFive, historySlowFive, historySlowTwo, historySlowOne):
        for tuple in historyFastFive.itertuples():
            self.fastFiveEMA.Update(tuple.Index, tuple.close)
        for tuple in historySlowFive.itertuples():
            self.slowFiveEMA.Update(tuple.Index, tuple.close)
            self.slowFiveSMA.Update(tuple.Index, tuple.close)
        for tuple in historySlowTwo.itertuples():
            self.slowTwoEMA.Update(tuple.Index, tuple.close)
            self.slowTwoSMA.Update(tuple.Index, tuple.close)
        for tuple in historySlowOne.itertuples():
            self.slowOneEMA.Update(tuple.Index, tuple.close)
            self.slowOneSMA.Update(tuple.Index, tuple.close)
        
    def RegisterIndicator(self, algorithm):
        algorithm.RegisterIndicator(self.Symbol, self.fastFiveEMA, self.fiveMinuteConsolidator, Field.Close)
        algorithm.RegisterIndicator(self.Symbol, self.slowFiveEMA, self.fiveMinuteConsolidator, Field.Close)
        algorithm.RegisterIndicator(self.Symbol, self.slowFiveSMA, self.fiveMinuteConsolidator, Field.Close)
        algorithm.RegisterIndicator(self.Symbol, self.slowTwoEMA, self.twoMinuteConsolidator, Field.Close)
        algorithm.RegisterIndicator(self.Symbol, self.slowTwoSMA, self.twoMinuteConsolidator, Field.Close)
        algorithm.RegisterIndicator(self.Symbol, self.slowOneEMA, self.oneMinuteConsolidator, Field.Close)
        algorithm.RegisterIndicator(self.Symbol, self.slowOneSMA, self.oneMinuteConsolidator, Field.Close)