import numpy as np
from datetime import timedelta
from QuantConnect.Securities.Option import OptionPriceModels
import pandas as pd

# A pairs trading strategy for dual listed stocks
class DualListedArb(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2017, 6,1)  
        self.SetEndDate(2017,12,1)    
        self.SetCash(10000)  
        
        self.TriggerSignal = .025
        # Adding the dual class shares of Viacom 
        self.AddEquity("VIA", Resolution.Minute)
        self.AddEquity("VIAB", Resolution.Minute)
        
        self.AddEquity("SPY", Resolution.Minute)
        # Set a benchmark at the beginning of every trading day
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                self.TimeRules.At(0,0),
                Action(self.SetBenchmark))
        self.SetWarmUp(int(6.5*60*100))

    def OnData(self, data):
        # Called every minute
        if self.IsWarmingUp: return
        if data["VIAB"] is None: return
        if data["VIA"] is None: return
        if not data.ContainsKey("VIAB"): return
        if not data.ContainsKey("VIA"): return      
        mainPrice = data["VIAB"].Price 
        dualPrice = data["VIA"].Price

        self.currentSpread = dualPrice - mainPrice
        self.currentSpread = float(self.currentSpread)
        self.percentDiff = (self.currentSpread-self.benchmarkSpread)/self.benchmarkSpread
        # If the current spread is >2.5% of the benchmark spread in the 
        # SetBenchmark method, buy the stock without the premium
        if self.percentDiff > .025:
            if self.Portfolio["VIAB"].Quantity <= 0:
                self.SetHoldings("VIAB", 1.0, tag="Long")
        if self.percentDiff < .025:
            if self.Portfolio["VIAB"].Quantity > 0:
                self.SetHoldings("VIAB", 0, tag="Short")
                    
    def SetBenchmark(self):
        # Returns a spread for the premium of owning VIA 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(["VIAB"], 100)
        self.dualHistory = self.History(["VIA"], 100)
        
        mainBars = self.mainHistory.loc["VIAB"]
        mainHistoricalClose = mainBars["close"]
        dualBars = self.dualHistory.loc["VIA"]
        dualHistoricalClose = dualBars["close"]
        
        spread = dualHistoricalClose - mainHistoricalClose
        self.benchmarkSpread = spread.mean()
        self.benchmarkSpread = float(self.benchmarkSpread)