from QuantConnect.Data.UniverseSelection import *
import pandas as pd
from decimal import Decimal


class RSISignalAlgorithm(QCAlgorithm):
    
    def Initialize(context):
        # // Set parameters for RSI Indicator and Signals, thereof.
        context.rsiPeriod = 21
        context.lowerThreshold = 40
        context.upperThreshold = 80
        
        # // Other parameters required for backtesting/project
        context.MinDollarVolume = 10000000          # 10M ADV
        context.PerTradeEquity = Decimal(0.01)      # Percent of Equity per Trade
        context.MinCashHoldings = Decimal(0.02)
        context.StopLossPercent = Decimal(0.1)      # 10% Stop Loss
        context.TrailingLeadPercent = Decimal(0.1)  # Trailing Stop after 10%
        context.TrailingStopPercent = Decimal(0.05) # at 5%
        
        # // Set the initial cash, start, and end dates for backtesting
        context.SetCash(1000000)         # 1M starting capital
        context.SetStartDate(2016, 7, 1) # Backtest start date
        context.SetEndDate(2018, 7, 1)   # Backtest end date
        
        # // Subscribe to data for securities returned in our selected universe
        # // We are deliberately setting Leverage to 1, as leverage was not
        # // specified for this project.
        context.UniverseSettings.Resolution = Resolution.Daily
        context.UniverseSettings.MinimumTimeInUniverse = 0
        context.UniverseSettings.Leverage = 1
        context.UniverseSettings.FillForward = False
        context.AddUniverse(context.universe_filter_coarse)
        # // Use raw prices for all securities.
        # context.SetSecurityInitializer(
        #     lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw))        
        
        # // Placeholder to collect indicator and other data for securities
        context.Data = {}
        
        # // Placeholder to collect various stats for securities for later use
        context.TrailingLeads = {}
        context.StopLossOrders = {}
        context.TrailingLeadReached = {}
        
        # // Finally, set a warm up period for the entire universe
        context.SetWarmUp(context.rsiPeriod)
        
    def universe_filter_coarse(context, coarse_data):
        # // Select all securities with Dollar Volume above specified threshold
        # // We make sure to arrange them according to decreasing ADV.
        selected = filter(lambda x: x.DollarVolume > context.MinDollarVolume, coarse_data)
        selected = sorted(list(selected), key=lambda x: x.DollarVolume, reverse=True)
        
        # // Initialize indicator and other data for all selected securities
        for cf in selected:
            if cf.Symbol not in context.Data:
                context.Data[cf.Symbol] = SymbolData(cf.Symbol, context.rsiPeriod,
                    context.upperThreshold, context.lowerThreshold)
            context.Data[cf.Symbol].update(cf)
            
        # // Since, we are going to buy securities with 1% of our equity on each
        # // trade, and we will be only doing this when we have enough cash
        # // available, we can return all available securities for our universe.
        return [ x.Symbol for x in selected ]
        
    # // Whenever new securities are added or removed in our universe,
    # // this function will be called with corresponding changes.
    def OnSecuritiesChanged(context, changes):
        for security in changes.RemovedSecurities:
            data = context.Data[security.Symbol]
            # // Sell a security, if we are invested in it and indicator
            # // value crossed the specified threshold for selling.
            if security.Invested and data.aboveThreshold:
                # context.Debug("Selling Ticker = %s, with Indicator Value = %s" % (data.symbol, data.rsi))
                context.Liquidate(security.Symbol, "SELL Threshold")
                context.mark_security_sold(security.Symbol)
                
        for security in changes.AddedSecurities:
            data = context.Data[security.Symbol]
            # // We will buy a security with 1% percent of our equity, when:
            # // - we have enough cash available (say, 2% of our equity)
            # // - indicator value crossed below the specified threshold for buy
            # // - we are not already holding the security (not specified in project)
            cash = context.Portfolio.Cash
            equity = context.Portfolio.TotalPortfolioValue
            cash_threshold = context.MinCashHoldings * equity
            if not security.Invested and data.belowThreshold and cash > cash_threshold:
                quantity = int(context.PerTradeEquity * equity/data.price)
                # // buy with a LimitOrder? or MarketOnOpen order? - Not sure.
                # // Going with MarketOnOpen order at the moment.
                context.MarketOnOpenOrder(security.Symbol, quantity, "BUY Threshold")
                # context.Debug("Buying Ticker = %s, with Indicator Value = %s" % (data.symbol, data.rsi))
                
                # // Also, since we need stop loss - place a StopMarketOrder right away.
                # // We, also, need to store these stop market orders in a dict,
                # // so that we can cancel them later, if required.
                stop_price = data.price * (1 - context.StopLossPercent)
                trailing_lead = data.price * (1 + context.TrailingLeadPercent)
                context.StopLossOrders[data.symbol] = context.StopMarketOrder(
                    data.symbol, -quantity, stop_price, "STOP LOSS")
                context.TrailingLeadReached[data.symbol] = False
                context.TrailingLeads[data.symbol] = trailing_lead

    # // This function is called on each tick, and is a good place to add logic for
    # // trailing stop loss functionality. At each tick, for each security, we
    # // check if we have reached a new trailing lead price, otherwise, we check
    # // if we have hit trailing stop price and sell the security as well as cancel
    # // any stop loss we may have (this is very important).
    def OnData(context, _data):
        for symbol, data in context.Data.items():
            # // check if we are holding the security?
            # // At every buy, we issue a stop loss order and cache this ticket.
            if context.Portfolio[symbol].HoldStock:
                if symbol in context.StopLossOrders:
                    sl_ticket = context.StopLossOrders[symbol]
                    tp_reached = context.TrailingLeadReached[symbol]
                    trailing_lead = context.TrailingLeads[symbol]
                    trailing_stop = trailing_lead * (1 - context.TrailingStopPercent)
                    # // stop loss was filled, we should remove any tracking
                    # // information for this symbol.
                    if sl_ticket.Status == "Filled":
                        context.mark_security_sold(symbol)
                    # // Trailing Lead was hit earlier and now price is below the
                    # // trailing stop price - we need to sell and cancel our
                    # // stop limit order.
                    elif tp_reached and data.price < trailing_stop:
                        context.Liquidate(symbol, "TRAIL STOP")
                        context.mark_security_sold(symbol)
                    # // Trailing Stop has not yet been hit. If price exceeds
                    # // trailing lead, we update it.
                    elif data.price >= trailing_lead:
                        trailing_lead = data.price
                        context.TrailingLeadReached[symbol] = True
                # else:
                    # context.Debug("Security in Portfolio without a StopLoss: %s" % symbol)
                
    # // when a security is marked as sold, we should ensure that we delete any
    # // remnants of it from the cached data/dicts we preserved for later use.
    # // Also, cancel StopLoss order created for this security when buying.
    def mark_security_sold(context, symbol):
        sl_ticket = context.StopLossOrders[symbol]
        if sl_ticket.Status != "Filled":
            sl_ticket.Cancel()
        context.StopLossOrders.pop(symbol)
        context.TrailingLeads.pop(symbol)
        context.TrailingLeadReached.pop(symbol)


class SymbolData(object):
    def __init__(self, symbol, rsiPeriod, upperThreshold, lowerThreshold):
        self.symbol = symbol
        self.price = 0
        self.rsi = RelativeStrengthIndex(rsiPeriod)
        self.upperThreshold = upperThreshold
        self.lowerThreshold = lowerThreshold
        self.aboveThreshold = False
        self.belowThreshold = False
        
    def update(self, data):
        self.price = data.Price
        if self.rsi.Update(data.EndTime, data.Price):
            rsi = self.rsi.Current.Value
            self.aboveThreshold = rsi > self.upperThreshold
            self.belowThreshold = rsi < self.lowerThreshold