# RetrospectiveBlackBuffalo


class RetrospectiveBlackBuffalo(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 2) 
        self.SetEndDate(2017, 12, 29) 
        self.SetCash(100000)  
        self.AddEquity("MSFT", Resolution.Daily)
        self.AddAlpha(MacdAlphaModel())
        self.AddAlpha(RsiAlphaModel())   
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

     
class MacdAlphaModel(AlphaModel):

    def __init__(self,
                 fastPeriod = 12,
                 slowPeriod = 25,
                 signalPeriod = 9,
                 movingAverageType = MovingAverageType.Exponential,
                 resolution = Resolution.Daily):

        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.signalPeriod = signalPeriod
        self.movingAverageType = movingAverageType
        self.resolution = resolution
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        self.bounceThresholdPercent = 0.0001
        self.symbolData = {}

        

    def Update(self, algorithm, data):

        insights = []

        for key, sd in self.symbolData.items():
            if sd.Security.Price == 0:
                continue

            direction = InsightDirection.Flat
            normalized_signal = sd.MACD.Signal.Current.Value / sd.Security.Price

            if normalized_signal > self.bounceThresholdPercent:
                direction = InsightDirection.Up
            elif normalized_signal < -self.bounceThresholdPercent:
                direction = InsightDirection.Flat
                
            if direction == sd.PreviousDirection:
                continue

            insight = Insight.Price(sd.Security.Symbol, self.insightPeriod, direction)
            sd.PreviousDirection = insight.Direction
            insights.append(insight)

        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            self.symbolData[added.Symbol] = SymbolData(algorithm, added, self.fastPeriod, self.slowPeriod, self.signalPeriod, self.movingAverageType, self.resolution)

        for removed in changes.RemovedSecurities:
            data = self.symbolData.pop(removed.Symbol, None)
            if data is not None:
                algorithm.SubscriptionManager.RemoveConsolidator(removed.Symbol, data.Consolidator)
                

class SymbolData:
    def __init__(self, algorithm, security, fastPeriod, slowPeriod, signalPeriod, movingAverageType, resolution):
        self.Security = security
        self.MACD = MovingAverageConvergenceDivergence(fastPeriod, slowPeriod, signalPeriod, movingAverageType)

        self.Consolidator = algorithm.ResolveConsolidator(security.Symbol, resolution)
        algorithm.RegisterIndicator(security.Symbol, self.MACD, self.Consolidator)

        self.PreviousDirection = None
     
        
class RsiAlphaModel(AlphaModel):

    def __init__(self,
                 period = 14,
                 resolution = Resolution.Daily):
        self.period = period
        self.resolution = resolution
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.SymbolData2BySymbol ={}

        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)

    def Update(self, algorithm, data):

        insights = []
        for symbol, SymbolData2 in self.SymbolData2BySymbol.items():
            rsi = SymbolData2.RSI
            previous_state = SymbolData2.State
            state = self.GetState(rsi, previous_state)

            if state != previous_state and rsi.IsReady:
                if state == State.TrippedLow:
                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Flat))
                if state == State.TrippedHigh:
                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Up))

            SymbolData2.State = state

        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        symbols = [ x.Symbol for x in changes.RemovedSecurities ]
        if len(symbols) > 0:
            for subscription in algorithm.SubscriptionManager.Subscriptions:
                if subscription.Symbol in symbols:
                    self.SymbolData2BySymbol.pop(subscription.Symbol, None)
                    subscription.Consolidators.Clear()

        addedSymbols = [ x.Symbol for x in changes.AddedSecurities if x.Symbol not in self.SymbolData2BySymbol]
        if len(addedSymbols) == 0: return

        history = algorithm.History(addedSymbols, self.period, self.resolution)

        for symbol in addedSymbols:
            rsi = algorithm.RSI(symbol, self.period, MovingAverageType.Wilders, self.resolution)

            if not history.empty:
                ticker = SymbolCache.GetTicker(symbol)

                if ticker not in history.index.levels[0]:
                    Log.Trace(f'RsiAlphaModel.OnSecuritiesChanged: {ticker} not found in history data frame.')
                    continue

                for tuple in history.loc[ticker].itertuples():
                    rsi.Update(tuple.Index, tuple.close)

            self.SymbolData2BySymbol[symbol] = SymbolData2(symbol, rsi)


    def GetState(self, rsi, previous):
        if rsi.Current.Value > 50:
            return State.TrippedHigh
        if rsi.Current.Value < 50:
            return State.TrippedLow

        return previous


class SymbolData2:
    def __init__(self, symbol, rsi):
        self.Symbol = symbol
        self.RSI = rsi
        self.State = State.Middle


class State(Enum):
    TrippedLow = 0
    Middle = 1
    TrippedHigh = 1