Overall Statistics
from Alphas.PearsonCorrelationPairsTradingAlphaModel import PearsonCorrelationPairsTradingAlphaModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

from sklearn import linear_model
import numpy as np
import pandas as pd
from scipy import stats
from math import floor
from datetime import timedelta

import itertools
import math
from random import shuffle



class PairsTradingAlgorithm(QCAlgorithm):
    

    def Initialize(self):


        self.stock_grouping = {0: ['NKE', 'AAP', 'QSR', 'GM', 'TJX', 'DHI', 'MGM', 'CBS']}
        
        self.pairs_list = []    


        for stocks in self.stock_grouping.values():
            result = itertools.combinations(stocks, 2)
            list_of_lists = [list(elem) for elem in list(result)]
            self.pairs_list +=  list_of_lists            
                    

        self.stock_to_grouping = {}
        for k,v in self.stock_grouping.items():
            for x in v:
                self.stock_to_grouping.setdefault(x,[]).append(k)
     
        for sector in self.stock_grouping.keys():
            self.Debug('waaaat')
       
        
        self.SetCash(100000)  # Set Strategy Cash
        self.AddAlpha(PearsonCorrelationPairsTradingAlphaModel(252, Resolution.Daily))
        self.SetExecution(ImmediateExecutionModel())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())
        symbols = [ Symbol.Create("NASDAQ", SecurityType.Equity, Market.USA) ]
        self.AddUniverseSelection(ManualUniverseSelectionModel(symbols))

        
      
        self.SetStartDate(2015,2,20)
        self.SetEndDate(2016,2,19)
        self.numdays = 60# set the length of training period

        stocks = self.stock_grouping.values()
        

        
        tickers = []
        
        for list_element in stocks:
            for tick in list_element:
                tickers += [tick]
        

        tickers = list(set(tickers))
  
        self.symbols = []
        
        

        self.tick_2_sym_dict = dict()
        for i in tickers:
            self.symbols.append(self.AddEquity(i,Resolution.Minute).Symbol)

        
        self.lastSlice = None
 

  



            

        

    def OnData(self, data):
        '''OnData event is the primary entry point for your algorithm. Each new data point will be pumped in here.
            Arguments:
                data: Slice object keyed by symbol containing the stock data'''
                
        self.lastSlice = data