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
        symbol = self.AddEquity('SPY', Resolution.Minute, Market.USA, True, 1, True).Symbol
        self.tradingSymbols = []
        
        self.SetAlpha(TradingModel(self.tradingSymbols, exclude_symbol=symbol))
        #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, exclude_symbol):
        
        self.tradeSymbols = universe
        self.exclude_symbol = exclude_symbol
        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)
                    if symbolData.slowFiveEMA.IsReady:
                        algorithm.Plot('Trade Plot', 'EMA', symbolData.slowFiveEMA.Current.Value)
                
        '''
        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 = []
        for added in changes.AddedSecurities:
            if added.Symbol not in self.symbolDataBySymbol and added.Symbol != self.exclude_symbol:
                symbols.append(added.Symbol)
        if len(symbols) == 0:
            return 
        
        history = algorithm.History(symbols, self.slowPeriod * 5, self.resolution)

        tickers = history.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)
            if not history.empty:
                symbolData.WarmUpIndicators(algorithm, history.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))
        
        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, history):
        for tuple in history.itertuples():
            tradebar = TradeBar(tuple.Index, self.Symbol, tuple.open, tuple.high, tuple.low, tuple.close, tuple.volume)
            self.fiveMinuteConsolidator.Update(tradebar)
            self.twoMinuteConsolidator.Update(tradebar)
            self.oneMinuteConsolidator.Update(tradebar)
        
    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)