from datetime import datetime

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

from System import *
from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Indicators import *
from sklearn import linear_model
import numpy as np
from scipy import stats


class PairsTradingAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        
        '''Initialise the data and resolution required, as well as the cash and start-end dates for your algorithm. All algorithms must initialized.'''
        
        self.SetStartDate(2014,1,5)
        self.SetEndDate(2014,11,20)
        self.SetCash(100000)
        self.numdays = 120  # set the length of training period
        self.ticker = ["XLK","QQQ"]
        self.syl = []
        self.threshold = 2
        self.coef = None
        self.mean, self.std = None, None
        for i in self.ticker:
            equity = self.AddSecurity(SecurityType.Equity, i, Resolution.Daily)
            self.syl.append(equity.Symbol)
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday, DayOfWeek.Monday),self.TimeRules.At(10, 0), Action(self._set_signal))
        
        #self.Schedule.On(self.DateRules.MonthStart(self.syl[0]), self.TimeRules.AfterMarketOpen(self.syl[0],30), Action(self._set_signal))
        #self.Schedule.On(self.DateRules.MonthStart(self.syl[1]), self.TimeRules.AfterMarketOpen(self.syl[1],30), Action(self._set_signal))
      
        #self.Schedule.On(self.DateRules.EveryDay(self.syl[0]),self.TimeRules.AfterMarketOpen(self.syl[0],30),Action(self._set_signal))
        #self.Schedule.On(self.DateRules.EveryDay(self.syl[1]),self.TimeRules.AfterMarketOpen(self.syl[1],30),Action(self._set_signal))
    
    def _set_signal(self):
        history = self.History(self.numdays,Nullable[Resolution](Resolution.Daily))
        logprice = {}
        for j in range(len(self.syl)):
            close = []
            for slice in history:
                bar = slice[self.syl[j]]
                close.append(bar.Close)
            logprice[self.ticker[j]] = np.log([float(z) for z in close])    # generate the log return series of stock price
        # run linear regression over the two history log price series
        reg = linear_model.LinearRegression()
        x,y = logprice[self.ticker[0]],logprice[self.ticker[1]]
        reg.fit([[n] for n in x],y)
        self.coef = reg.coef_    # reg.intercept_
        self.spread = y - self.coef * x  #compute the spread series based on regression result
        self.mean, self.std = np.mean(self.spread),np.std(self.spread)

    def OnData(self,data):
            
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.'''
        if self.mean == None or self.std == None: return
        else:
            logprice_x = np.log(float(self.Portfolio[self.syl[0]].Price))
            logprice_y = np.log(float(self.Portfolio[self.syl[1]].Price))
        
            syl_x, syl_y = self.syl[0],self.syl[1]
            quantity = float(self.CalculateOrderQuantity(syl_y,0.4))
        
            current_spread = float(logprice_y) - self.coef * float(logprice_x) 
            if self.mean == None or self.std == None: return
            else:
                if current_spread > self.mean + self.threshold * self.std:
                    self.Sell(syl_y, 1 * quantity) 
                    self.Buy(syl_x, self.coef * quantity)
                    self.Log('Trade I : ' + str(syl_y)+' ,  '+str(quantity)+' ,  '+str(syl_x)+' ,  '+str(self.coef*quantity)) 
                elif current_spread < self.mean - self.threshold * self.std:
                    self.Buy(syl_y, 1 * quantity )
                    self.Sell(syl_x, self.coef * quantity)
                    self.Log('Trade II : ' + str(syl_y)+' ,  '+str(quantity)+' ,  '+str(syl_x)+' ,  '+str(self.coef*quantity)) 
                if self.Portfolio[self.syl[0]].Quantity != 0 and self.Portfolio[self.syl[1]].Quantity != 0:
                    if  current_spread < 0.5 * self.std and current_spread > -0.5 * self.std:
                        self.Liquidate(self.syl[0])
                        self.Liquidate(self.syl[1])