import numpy as np
import decimal as d 
class PairsDualListedArbitrage(QCAlgorithm):

    def Initialize(self):
        
        self.SetStartDate(2016,1,1)  
        self.SetEndDate(2017,12,1)    
        self.SetCash(1000)
        
        self.TakeProfit = d.Decimal(1.025)
        self.StopLoss = d.Decimal(1.01)
        
        self.StringSymbols = ["DISCA", "DISCK"]
        self.symbols = [
                self.AddEquity(self.StringSymbols[0], Resolution.Minute),
                self.AddEquity(self.StringSymbols[1], Resolution.Minute)
                ]
                
        self.fast = self.SMA(self.StringSymbols[1], 20, Resolution.Daily)
        self.slow = self.SMA(self.StringSymbols[1], 50, Resolution.Daily)
        
        self.AddEquity("SPY", Resolution.Minute)
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.AfterMarketOpen("SPY",0),
                Action(self.SetBenchmark))
                
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.AfterMarketOpen("SPY", 60),
                Action(self.Rebalance1))

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.AfterMarketOpen("SPY", 380),
                Action(self.Reset))
        self.SetWarmUp(int(6.5*60*50))

    def OnData(self, data):
        if self.IsWarmingUp: return
    
    def SetBenchmark(self):
        # Returns a spread for the premium of owning X stock over the 
        # past 100 days
        # The mean of those 100 spreads becomes the Benchmark Spread that we base our
        # predictions off of
        self.mainHistory = self.History([self.StringSymbols[1]], 50)  # Class B shares
        self.dualHistory = self.History([self.StringSymbols[0]], 50)   # Class A shares
        
        mainBars = self.mainHistory.loc[self.StringSymbols[1]]
        mainHistoricalClose = mainBars["close"]
        dualBars = self.dualHistory.loc[self.StringSymbols[0]]
        dualHistoricalClose = dualBars["close"]
        
        spread = dualHistoricalClose - mainHistoricalClose
        self.benchmarkSpread = spread.mean()
        self.benchmarkSpread = float(self.benchmarkSpread)
        
    def Rebalance1(self):
        main = self.Securities[self.StringSymbols[1]]
        dual = self.Securities[self.StringSymbols[0]]
        
        mainPrice = main.Price
        dualPrice = dual.Price
        
        currentSpread = dualPrice - mainPrice
        currentSpread = float(currentSpread)
        self.PercentDiff = (currentSpread - self.benchmarkSpread)/self.benchmarkSpread
        
        self.IsUpTrend = self.fast.Current.Value > self.slow.Current.Value
        if self.IsUpTrend:
            if self.PercentDiff > .070:
                if self.Portfolio[self.StringSymbols[1]].Quantity <= 0:
                    self.MarketOrder(self.StringSymbols[1], int(self.Portfolio.TotalPortfolioValue/mainPrice))
                    self.LimitOrder(self.StringSymbols[1], -1*abs((self.Portfolio[self.StringSymbols[1]].Quantity)), d.Decimal(mainPrice)*self.TakeProfit, tag="Short")
                    self.StopMarketOrder(self.StringSymbols[1], -1*abs((self.Portfolio[self.StringSymbols[1]].Quantity)), d.Decimal(mainPrice)/self.StopLoss, tag="Short")
                
    def Reset(self):
        self.Liquidate(self.StringSymbols[1])

    def OnOrderEvent(self, orderEvent):
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        
        if order.Status == OrderStatus.Filled:
            if order.Type == OrderType.Limit or order.Type == OrderType.StopMarket:
                self.Transactions.CancelOpenOrders(order.Symbol)