import pandas as pd

class EMABasedStrategy(QCAlgorithm):
    
    def Initialize(self):       #Initialize Dates, Cash, Equities, Fees, Allocation, Parameters, Indicators, Charts
        
        # Set Start Date, End Date, and Cash
        #-------------------------------------------------------
        self.SetTimeZone(TimeZones.NewYork)     #EDIT: Added Timezon
        self.SetStartDate(2012, 1, 1)   # Set Start Date
        self.SetEndDate(2013, 1, 1)      # Set End Date
        self.SetCash(10000)            # Set Strategy Cash
        self.SetBenchmark("SPY")
        #-------------------------------------------------------
        
        # Set Custom Universe
        #-------------------------------------------------------
        self.AddUniverse(self.CoarseSelectionFilter, self.FineSelectionFilter)
        self.UniverseSettings.Resolution = Resolution.Hour    #Needs to change to Resolution.Minute once code works, leaving Daily for now to minimize data
        self.UniverseSettings.SetDataNormalizationMode = DataNormalizationMode.SplitAdjusted
        self.UniverseSettings.FeeModel = ConstantFeeModel(0.0)
        self.UniverseSettings.Leverage = 1
        self.SetBrokerageModel(BrokerageName.Alpaca, AccountType.Cash)      #EDIT: Added Brokerage, appears to have set fees to zero
        #-------------------------------------------------------
        
        # Set Contants 
        #-------------------------------------------------------
        self.EMA_Period_Fast = 5
        self.EMA_Period_Slow = 50
        self.EMA_Period_Medium = 10

        self.__numberOfSymbols     = 5
        self.__numberOfSymbolsFine = 2
        #-------------------------------------------------------
        
        # Define Percentage Allocation and variables
        #-------------------------------------------------------
        self.percentagebuy = 0.02
        self.percentagesell = 0.05
        self.stopLossPercent = 0.90
        self.indicators = {}
        self.buyOrderTicketBySymbol = {}
        self.sellOrderTicketBySymbol = {}
        self.trailingStopLossPrice = {}
        #-------------------------------------------------------
        
        self.Schedule.On(self.DateRules.On(2011, 12, 29), self.TimeRules.At(13, 0), self.SpecificTime);
        
        self.daily_equity = pd.Series([])
        
    def SpecificTime(self):
        #self.Log("SpecificTime: Fired at : {0}".format(self.Time)) 
        self.Liquidate()
        
    def CoarseSelectionFilter(self, coarse):
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)   # sort descending by daily dollar volume
        return [ x.Symbol for x in sortedByDollarVolume[:self.__numberOfSymbols] ]  # return the symbol objects of the top entries from our sorted collection
    
    def FineSelectionFilter(self, fine):  # sort the data by P/E ratio and take the top 'NumberOfSymbolsFine'
        sortedByPeRatio = sorted(fine, key=lambda x: x.OperationRatios.OperationMargin.Value, reverse=False)    # sort descending by P/E ratio
        self.universe = [ x.Symbol for x in sortedByPeRatio[:self.__numberOfSymbolsFine] ]  # take the top entries from our sorted collection
        return self.universe
    
    def OnSecuritiesChanged(self, changes):
        # Create indicator for each new security
        for security in changes.AddedSecurities:
            self.indicators[security.Symbol] = SymbolData(security.Symbol, self, self.EMA_Period_Fast, self.EMA_Period_Slow, self.EMA_Period_Medium)
        
        # for security in changes.RemovedSecurities:
        #     if security.Invested:
        #         self.Liquidate(security.Symbol, "Universe Removed Security")
            
        #     if security in self.indicators:
        #         self.indicators.pop(security.Symbol, None)
    
    
    def OnData(self, data):         #Entry Point for Data and algorithm - Check Data, Define Buy Quantity, Process Volume, Check Portfolio, Check RSI, Execute Buy/Sell orders, Chart Plots
        for symbol in self.universe:
            
            if not data.ContainsKey(symbol):    #Tested and Valid/Necessary
                continue
            
            if data[symbol] is None:            #Tested and Valid/Necessary
                continue
            
            if not symbol in self.indicators:    #Tested and Valid/Necessary
                continue
            
            # Ensure indicators are ready to update rolling windows
            if not self.indicators[symbol].slow_ema.IsReady:
                continue
            
            # Update EMA rolling windows
            self.indicators[symbol].fast_ema_window.Add(self.indicators[symbol].get_fast_EMA())
            self.indicators[symbol].slow_ema_window.Add(self.indicators[symbol].get_slow_EMA())
            self.indicators[symbol].medium_ema_window.Add(self.indicators[symbol].get_medium_EMA())
            
            # Check for Indicator Readiness within Rolling Window
            #-------------------------------------------------------
            if not (self.indicators[symbol].fast_ema_window.IsReady and self.indicators[symbol].slow_ema_window.IsReady and self.indicators[symbol].medium_ema_window.IsReady): 
                continue
            
            #EXECUTE TRADING LOGIC HERE -
            if self.Portfolio[symbol].Invested:
                if self.Portfolio[symbol].IsLong:
                    #self.Debug('symbol in long -->' + str(symbol) + '')
                    if self.indicators[symbol].is_buy_signal_liquidate():
                        self.Liquidate(symbol)
                    elif self.Portfolio[symbol].Price > self.buyOrderTicketBySymbol[symbol].AverageFillPrice + (self.buyOrderTicketBySymbol[symbol].AverageFillPrice*(1-self.stopLossPercent)) \
                        and self.Portfolio[symbol].Price > self.trailingStopLossPrice[symbol]: 
                        # self.Portfolio[symbol] returns EquityHoldings class
                        updateFields = UpdateOrderFields()
                        updateFields.StopPrice = self.Portfolio[symbol].Price * self.stopLossPercent
                        self.trailingStopLossPrice[symbol] = updateFields.StopPrice
                        self.buyOrderTicketBySymbol[symbol].Update(updateFields)
                elif self.indicators[symbol].is_sell_signal_liquidate() and self.Portfolio[symbol].IsShort:
                    #self.Debug('sell action  ->' + str(symbol))
                    self.Liquidate(symbol)
            if self.Portfolio.MarginRemaining > self.percentagebuy * self.Portfolio.TotalPortfolioValue:
                if not self.Portfolio[symbol].Invested:
                    if self.indicators[symbol].is_buy_signal():
                        self.buyquantity = round((self.percentagebuy*self.Portfolio.TotalPortfolioValue)/data[symbol].Close)
                        self.ticket = self.MarketOrder(symbol, self.buyquantity)
                        self.buyOrderTicketBySymbol[symbol] = self.ticket
                        self.trailingStopLossPrice[symbol] = self.ticket.AverageFillPrice
                        self.sellOrderTicketBySymbol.pop(symbol,None)
                        #self.Debug('trigger buy ordeer -->' + str(symbol))
                    elif self.indicators[symbol].is_sell_signal():
                        self.sellQuantity = round((self.percentagesell*self.Portfolio.TotalPortfolioValue)/data[symbol].Close)
                        self.ticket = self.MarketOrder(symbol, -1*self.sellQuantity)
                        self.sellOrderTicketBySymbol[symbol] = self.ticket
                        self.trailingStopLossPrice[symbol] = self.ticket.AverageFillPrice
                        self.buyOrderTicketBySymbol.pop(symbol,None)
                        #self.Debug('trigger sell ordeer -->' + str(symbol))
    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        #self.Log("{0}: {1}: {2}".format(self.Time, order.Type, orderEvent))
        
    def OnEndOfDay(self):
        self.daily_equity = self.daily_equity.append(pd.Series(self.Portfolio.TotalPortfolioValue, index=[self.Time]))
        
    def OnEndOfAlgorithm(self):
        self.Log(self.daily_equity.to_string())

class SymbolData(object):
    
    rolling_window_length = 2
    
    def __init__(self, symbol, context, fast_ema_period, slow_ema_period, medium_ema_period):
        
        self.symbol = symbol
        self.fast_ema_period = fast_ema_period
        self.slow_ema_period = slow_ema_period
        self.medium_ema_period = medium_ema_period
        self.fast_ema = context.SMA(symbol, self.fast_ema_period, Resolution.Daily)    #, fillDataForward = True, leverage = 1, extendedMarketHours = False)
        self.slow_ema = context.SMA(symbol, self.slow_ema_period, Resolution.Daily)    #, fillDataForward = True, leverage = 1, extendedMarketHours = False)
        self.medium_ema = context.SMA(symbol, self.medium_ema_period, Resolution.Daily)
        self.fast_ema_window = RollingWindow[float](self.rolling_window_length)
        self.slow_ema_window = RollingWindow[float](self.rolling_window_length)
        self.medium_ema_window = RollingWindow[float](self.rolling_window_length)
        
        # Warm up EMA indicators
        history = context.History([symbol], slow_ema_period + self.rolling_window_length, Resolution.Hour)
        for time, row in history.loc[symbol].iterrows():
            self.fast_ema.Update(time, row["close"])
            self.slow_ema.Update(time, row["close"])
            self.medium_ema.Update(time, row["close"])
            
            # Warm up rolling windows
            if self.fast_ema.IsReady:
                self.fast_ema_window.Add(self.fast_ema.Current.Value)
            if self.slow_ema.IsReady:
                self.slow_ema_window.Add(self.slow_ema.Current.Value)
            if self.medium_ema.IsReady:
                self.medium_ema_window.Add(self.medium_ema.Current.Value)
    
    def get_fast_EMA(self):
        return self.fast_ema.Current.Value
        
    def get_slow_EMA(self):
        return self.slow_ema.Current.Value
        
    def get_medium_EMA(self):
        return self.medium_ema.Current.Value
        
    def is_buy_signal(self):
        if self.medium_ema_window[0] > self.slow_ema_window[0] and \
            self.medium_ema_window[1] < self.slow_ema_window[1]:
                if self.fast_ema_window[0] > self.slow_ema_window[0]:
                    return True

    def is_sell_signal(self):
        if self.medium_ema_window[0] < self.slow_ema_window[0] and \
            self.medium_ema_window[1] > self.slow_ema_window[1]:
                if self.fast_ema_window[0] < self.slow_ema_window[0]:
                    return True
    
    def is_buy_signal_liquidate(self):
        return self.fast_ema_window[0] < self.slow_ema_window[0] and \
            self.fast_ema_window[1] > self.slow_ema_window[1]
            
    def is_sell_signal_liquidate(self):
        return self.fast_ema_window[0] > self.slow_ema_window[0] and \
            self.fast_ema_window[1] < self.slow_ema_window[1]