from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
from Selection.QC500UniverseSelectionModel import QC500UniverseSelectionModel

class SimpleRSITestQC500Universe(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2020, 1, 1) # Set Start Date
        self.SetEndDate(2020, 6, 5) # Set End Date
        self.SetCash(100000) # Set Strategy Cash
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.05))
        symbols = [ Symbol.Create("SPY", SecurityType.Equity, Market.USA), Symbol.Create("GE", SecurityType.Equity, Market.USA), Symbol.Create("BA", SecurityType.Equity, Market.USA) ]
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddAlpha(RsiAlphaModelTest())
        

class RsiAlphaModelTest(AlphaModel):

    def __init__(self, period = 14, resolution = Resolution.Daily):
        self.period = period
        self.resolution = resolution
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.symbolDataBySymbol = {}
        self.closeWindows = {}
        self.rsiWindows = {}
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)
        
    def Update(self, algorithm, data):
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if data.ContainsKey(symbol) and data[symbol] is not None:
                self.closeWindows[symbol].Add(data[symbol].Close)
            if self.closeWindows[symbol].Count>2:
                algorithm.Debug(self.closeWindows[symbol][2])
            rsi = symbolData.RSI
            
            self.rsiWindows[symbol].Add(rsi.Current.Value)
            
            if self.rsiWindows[symbol].IsReady:
                # plot oldest RSI value
                algorithm.Plot('RSI', symbol.Value, self.rsiWindows[symbol][19]) 
                
            previous_state = symbolData.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.Up))
                if state == State.TrippedHigh:
                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Down))
            symbolData.State = state
        return insights
        

    def OnSecuritiesChanged(self, algorithm, changes):
        
        # clean up data for removed securities
        symbols = [ x.Symbol for x in changes.RemovedSecurities ]
        if len(symbols) > 0:
            for subscription in algorithm.SubscriptionManager.Subscriptions:
                if subscription.Symbol in symbols:
                    self.symbolDataBySymbol.pop(subscription.Symbol, None)
                    subscription.Consolidators.Clear()
                
        # initialize data for added securities
        
        addedSymbols = [ x.Symbol for x in changes.AddedSecurities if x.Symbol not in self.symbolDataBySymbol]
        if len(addedSymbols) == 0: return
        
        history = algorithm.History(addedSymbols, self.period + 20, self.resolution)
        
        for symbol in addedSymbols:
            rsi = algorithm.RSI(symbol, self.period, MovingAverageType.Wilders, self.resolution)
            #rsi.Updated += self.RsiUpdated(symbol=symbol, sender=sender, updated=updated)
            self.rsiWindows[symbol] = RollingWindow[float](20)
            self.closeWindows[symbol] = RollingWindow[float](self.period)
            # symbolTradeBarsHistory = history.loc[symbol]
            # symbolClose = symbolTradeBarsHistory["close"]
            # symbolTime = symbolTradeBarsHistory["time"]
            # for historyIndex in range(self.period):
            #     self.closeWindows[symbol].Add(symbolClose[historyIndex])
            #     self.rsiWindows[symbol].Add(symbolTime[historyIndex],symbolClose[historyIndex])
            # 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[symbol].itertuples():
                self.closeWindows[symbol].Add(tuple.close)
                rsi.Update(tuple.Index, tuple.close)
                if rsi.IsReady:
                    self.rsiWindows[symbol].Add(rsi.Current.Value)
            self.symbolDataBySymbol[symbol] = SymbolData(symbol, rsi)
            # symbolTradeBarsHistory = None
            # symbolClose = None
        
        for k in self.closeWindows.keys():
            algorithm.Debug(str(k) + ' ' + str(self.closeWindows[k][0]) + ' ' + str(self.closeWindows[k][1]) + ' ' + str(self.closeWindows[k][2]) + ' ' + str(self.closeWindows[k][3]) + ' ' + str(self.closeWindows[k][4]) + ' ' + str(self.closeWindows[k][5]))
    
    def GetState(self, rsi, previous):
        if rsi.Current.Value > 70:
            return State.TrippedHigh
        if rsi.Current.Value < 30:
            return State.TrippedLow
        if previous == State.TrippedLow:
            if rsi.Current.Value > 35:
                return State.Middle
        if previous == State.TrippedHigh:
            if rsi.Current.Value < 65:
                return State.Middle
            
        return previous
        
        #def RsiUpdated(self, symbol, sender, updated):
        # self.rsiWindows[symbol].Add(updated)


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


class State(Enum):
    '''Defines the state. This is used to prevent signal spamming and aid in bounce detection.'''
    TrippedLow = 0
    Middle = 1
    TrippedHigh = 2