from enum import Enum
from datetime import timedelta
import numpy as np


class BasicTemplateFrameworkAlgorithm(QCAlgorithm):
    '''Basic template framework algorithm uses framework components to define the algorithm.'''

    def Initialize(self):
        ''' Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''

        # Set requested data resolution
        self.UniverseSettings.Resolution = Resolution.Hour
        self.SetStartDate(2019,6,7)   #Set Start Date
        self.SetEndDate(2019,10,11)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash

        symbols = [ Symbol.Create("EURUSD", SecurityType.Forex, Market.FXCM) ]

        # set algorithm framework models
        self.SetUniverseSelection(ManualUniverseSelectionModel(symbols))
        self.SetAlpha(RsiAlphaModel(period = 7,lower = 25,upper = 75,resolution = Resolution.Hour))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(Resolution.Daily))
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(0.01))
        
        PricePlot = Chart('Trade Plot')
        PricePlot.AddSeries(Series('Price', SeriesType.Line, 0))
        self.AddChart(PricePlot)
        
        RSIPlot = Chart('RSI')
        RSIPlot.AddSeries(Series('rsi', SeriesType.Line, 0))
        self.AddChart(RSIPlot)
        
        StatePlot = Chart('State')
        StatePlot.AddSeries(Series('state', SeriesType.Line, 0))
        self.AddChart(StatePlot)


    def OnOrderEvent(self, orderEvent):
        if orderEvent.Status == OrderStatus.Filled:
            self.Debug("Purchased Stock: {0}".format(orderEvent.Symbol))


class RsiAlphaModel(AlphaModel):
    '''Uses Wilder's RSI to create insights.
    Using default settings, a cross over below 30 or above 70 will trigger a new insight.'''

    def __init__(self,period = 14,lower = 25,upper = 75,resolution = Resolution.Daily):

        '''Initializes a new instance of the RsiAlphaModel class
        Args:
            period: The RSI indicator period'''
        self.period = period
        self.lower = lower
        self.upper = upper
        self.resolution = resolution
        self.insightPeriod = Time.Multiply(Extensions.ToTimeSpan(resolution), period)
        self.symbolDataBySymbol ={}
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        self.Name = '{}({},{})'.format(self.__class__.__name__, period, resolutionString)

    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The new data available
        Returns:
            The new insights generated'''
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            rsi = symbolData.RSI
            previous_state = symbolData.State
            state = self.GetState(rsi, previous_state)

            if state != previous_state and rsi.IsReady:
                if state == State.CrossOverAbove:
                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Up))
                if state == State.CrossOverBelow:
                    insights.append(Insight.Price(symbol, self.insightPeriod, InsightDirection.Down))
            symbolData.State = state
            
        algorithm.Plot('Trade Plot', 'Price' , algorithm.Securities[symbol].Price)    
        algorithm.Plot('RSI'       , 'rsi'   , rsi.Current.Value)
        algorithm.Plot('State'     , 'state' , self.GetState(rsi, previous_state).value)
        
        algorithm.Log(' Price:'+str(algorithm.Securities[symbol].Price)+' rsi:'+str(rsi.Current.Value)+' state:'+str(self.GetState(rsi, previous_state).value))

        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        '''Cleans out old security data and initializes the RSI for any newly added securities.
        Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''

        # 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, 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.symbolDataBySymbol[symbol] = SymbolData(symbol, rsi)


    def GetState(self, rsi, previous):
        ''' Determines the new state. This is basically cross-over detection logic that
        includes considerations for bouncing using the configured bounce tolerance.'''
        if rsi.Current.Value > self.upper:
            return State.TrippedHigh
        if rsi.Current.Value < self.lower:
            return State.TrippedLow
        if previous == State.TrippedLow:
            if rsi.Current.Value > self.lower:# and rsi.Current.Value < self.lower+5
                return State.CrossOverAbove
        if previous == State.TrippedHigh:
            if rsi.Current.Value < self.upper:#and rsi.Current.Value > self.upper-5
                return State.CrossOverBelow
        if previous == State.CrossOverBelow or previous == State.CrossOverAbove:
            return State.Middle

        return previous


class SymbolData:
    '''Contains data specific to a symbol required by this model'''
    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 = -2
    CrossOverAbove = -1
    CrossOverBelow = 1
    TrippedHigh = 2
    Middle = 0