from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Common")

from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Data.UniverseSelection import *
from QuantConnect.Data import *
from QuantConnect.Indicators import *

class upTrendUniverseExample(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019,1,4)
        self.SetEndDate(2019,1,10)
        self.SetCash(100000)
        self.obj = {}
        
        # Choose number of securities in universe
        self.numberOfSecurities = 10
        
        # Set lookback period for EMA and ATR
        self.period = 8

        # Add a universe using CoarseSelectionFunction
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction)
    
    # Retrieve a list of stocks that have the highest dollar volume
    def CoarseSelectionFunction(self, coarse):
        CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData]
        sortedByDollarVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=False) 
        top = sortedByDollarVolume[:self.numberOfSecurities]
        return [i.Symbol for i in top]

    def OnData(self, slice):
        for x in self.ActiveSecurities.Values:
            # Update values on SymbolData object
            avg = self.obj[x.Symbol]
            avg.updateATR(slice[x.Symbol]) 
            avg.updateEMA(self.Time, x.Price)
            
            # Return if already invested in stock
            if self.Portfolio[x.Symbol].Invested: return
        
            # Invest with equal weights in each stock that is trending upwards (price > EMA)
            if avg.IsAboveEMA and avg.IsAtrReady:
                
                quantity = self.CalculateOrderQuantity(x.Symbol,1/self.numberOfSecurities)
                self.Buy(x.Symbol,quantity)
                # Set a limit order for Take-Profit, using the Average True Range (ATR) indicator
                self.LimitOrder(x.Symbol, -quantity, x.Price + avg.atr.Current.Value*2)
                # Set a Stop Market Order for Stop-Loss, using the Average True Range (ATR) indicator
                self.StopMarketOrder(x.Symbol, -quantity, x.Price - avg.atr.Current.Value*1)
                
        # Reset changes
        self.changes = None

    def OnSecuritiesChanged(self, changes):
        self.changes = changes

        # List of added symbols
        self.addedSymbols = [x.Symbol for x in changes.AddedSecurities]  
        
        # Get history for added symbols
        history = self.History(self.addedSymbols, self.period, Resolution.Daily)

        for symbol in self.addedSymbols:
            if symbol not in self.obj.keys():
                # Store the symbolData object in obj
                self.obj[symbol] = symbolData(symbol, self.period)
                # Retrieve TradeBars using History()
                bars = self.History(str(symbol),self.period,Resolution.Daily) 
                # WarmUp the ATR Indicator
                self.obj[symbol].WarmUpIndicatorBar(bars)
                # WarmUp the EMA Indicator
                if str(symbol) in history.index:
                    self.obj[symbol].WarmUpIndicator(history.loc[str(symbol)])

    # Use OnOrderEvent() to cancel remaining order if limit order or stop loss order is executed
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        
        if order.Status == OrderStatus.Filled:
            if order.Type == OrderType.Limit or OrderType.StopLimit:
                self.Transactions.CancelOpenOrders(order.Symbol)
            
            if order.Status == OrderStatus.Canceled:
                self.Log(str(orderEvent))

        
class symbolData(object):
    def __init__(self, symbol, period):
        self.symbol = symbol
        self.atr = AverageTrueRange(period)
        self.ema = ExponentialMovingAverage(period)
        self.IsAboveEMA = False
        self.IsAtrReady = False

    # Warm up the EMA indicator using pandas DataFrame
    def WarmUpIndicator(self, history):
        for tuple in history.itertuples():
            item = IndicatorDataPoint(self.symbol, tuple.Index, float(tuple.close))
            self.ema.Update(item)
    
    # Warm up the ATR indicator using TradeBars
    def WarmUpIndicatorBar(self, bars):
            for bar in bars:
                self.atr.Update(bar)
     
    # Update the EMA indicator               
    def updateEMA(self, time, price):
        if self.ema.Update(time, price):
            self.IsAboveEMA = price > self.ema.Current.Value
    
    # Update the ATR indicator
    def updateATR(self, bar):
        if self.atr.Update(bar):
             self.IsAtrReady = self.atr.Current.Value > 0