from datetime import timedelta, datetime

class SMAPairsTrading(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 7, 1) 
        self.SetCash(10000)
        
        # Using the ManualUniverseSelectionModel(), add the symbols "PEP" and "KO" 
        symbols = [Symbol.Create("PEP", SecurityType.Equity, Market.USA), Symbol.Create("KO", SecurityType.Equity, Market.USA)]
        self.AddUniverseSelection(ManualUniverseSelectionModel(symbols))

        # In Universe Settings, set the resolution to hour
        self.UniverseSettings.Resolution = Resolution.Minute
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.Raw
        
        # Create an instance of the PairsTradingAlphamodel()
        self.AddAlpha(PairsTradingAlphaModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())
        self.SetRiskManagement(TrailingStopRiskManagementModel(0.03))
    
    # Use OnEndOfDay() to Log() your positions at the close of each trading day.
    def OnEndOfDay(self, symbol):
        self.Log("Taking a position of " + str(self.Portfolio[symbol].Quantity) + " units of symbol " + str(symbol))

class PairsTradingAlphaModel(AlphaModel):

    def __init__(self):
        # Initialize an empty list self.pair = [ ]
        self.pair = [ ]
        # Create a 500-day Simple Moving Average Indicator monitoring the spread SMA 
        self.spreadMean = SimpleMovingAverage(500)
        # Create a 500-day Standard Deviation Indicator monitoring the spread Std 
        self.spreadStd = StandardDeviation(500)
        # Set self.period to a 2 hour timedelta 
        self.period = timedelta(hours=2)
        
    def Update(self, algorithm, data):
        # Set the price difference calculation to self.spread.
        spread = self.pair[1].Price - self.pair[0].Price
        # Update the spreadMean indicator with the spread
        self.spreadMean.Update(algorithm.Time, spread)
        # Update the spreadStd indicator with the spread
        self.spreadStd.Update(algorithm.Time, spread) 
        
        # Save our upper threshold and lower threshold
        upperthreshold = self.spreadMean.Current.Value + 2*self.spreadStd.Current.Value
        lowerthreshold = self.spreadMean.Current.Value - 2*self.spreadStd.Current.Value

        # Emit an Insight.Group() if the spread is greater than the upperthreshold 
        if spread > upperthreshold:
            return Insight.Group(
                [
                    Insight.Price(self.pair[0].Symbol, self.period, InsightDirection.Up),
                    Insight.Price(self.pair[1].Symbol, self.period, InsightDirection.Down)
                ])
        
        # Emit an Insight.Group() if the spread is less than the lowerthreshold
        if spread < lowerthreshold:
            return Insight.Group(
                [
                    Insight.Price(self.pair[0].Symbol, self.period, InsightDirection.Down),
                    Insight.Price(self.pair[1].Symbol, self.period, InsightDirection.Up)
                ])

        # If the spread is not greater than the upper or lower threshold, do not return Insights
        return []
    
    def OnSecuritiesChanged(self, algorithm, changes):
        # Set self.pair to the changes.AddedSecurities changes
        self.pair = [x for x in changes.AddedSecurities]
        
        # Call for 500 days of history data for each symbol in the pair and save to the variable history
        history = algorithm.History([x.Symbol for x in self.pair], 500)
        # Unstack the Pandas data frame to reduce it to the history close price
        history = history.close.unstack(level=0)
        # Iterate through the history tuple and update the mean and standard deviation with historical data
        for tuple in history.itertuples():
            self.spreadMean.Update(tuple[0], tuple[2]-tuple[1])
            self.spreadStd.Update(tuple[0], tuple[2]-tuple[1])