'''
Based on Pairs Trading with Stocks strategy by Jin Wu 2018
https://www.quantconnect.com/learning/articles/investment-strategy-library/pairs-trading-with-stocks
'''
#region imports
from AlgorithmImports import *
#endregion
# https://quantpedia.com/Screener/Details/12

import numpy as np
import pandas as pd
from scipy import stats
from math import floor
from datetime import timedelta
from collections import deque
import itertools as it
from decimal import Decimal

class PairsTradingAlgorithm(QCAlgorithm):
    
    def Initialize(self):
        
        self.SetStartDate(2017,1,1)
        self.SetEndDate(2017,6,1)
        self.SetCash(10000)
       
        tickers = ['COF','BRK.B', 'JPM', 'V', 'MA', 'BAC', 'WFC', 'SPGI', 'GS', 'MS', 
                    'BLK', 'AXP', 'MMC', 'C', 'PYPL','CB', 'FISV', 'PGR', 'SCHW', 'CME', 
                    'AON','ICE', 'MCO', 'PNC', 'AJG','TRV', 'USB', 'AFL','AIG', 'MSCI', 
                    'MET', 'TFC']
                    #  
                    # 
                    # ]
        self.threshold = 2
        self.symbols = []
        for i in tickers:
            self.symbols.append(self.AddEquity(i, Resolution.Daily).Symbol)
        
        self.pairs = {}
        self.formation_period = int(self.GetParameter("days")) #252

        self.history_price = {}
        for symbol in self.symbols:
            hist = self.History([symbol], self.formation_period+1, Resolution.Daily)
            if hist.empty: 
                self.symbols.remove(symbol)
            else:
                self.history_price[str(symbol)] = deque(maxlen=self.formation_period)
                for tuple in hist.loc[str(symbol)].itertuples():
                    self.history_price[str(symbol)].append(float(tuple.close))
                if len(self.history_price[str(symbol)]) < self.formation_period:
                    self.symbols.remove(symbol)
                    self.history_price.pop(str(symbol))

        self.symbol_pairs = list(it.combinations(self.symbols, 2))
        
        # Add the benchmark
        self.AddEquity("SPY", Resolution.Daily) 
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), self.Rebalance)
        self.count = 0
        self.sorted_pairs = None
        
        
    def OnData(self, data):
        # Update the price series everyday
        self.Log("Hist: "+str(self.history_price))
        for symbol in self.symbols:
            if data.Bars.ContainsKey(symbol) and str(symbol) in self.history_price:
                self.history_price[str(symbol)].append(float(data[symbol].Close)) 
        if self.sorted_pairs is None: return
        
        x=0
        #self.Log("Len Sorted pairs: "+str(len(self.sorted_pairs)))
        for i in self.sorted_pairs:
            #self.Log("OnData Self.sorted x{}: {} {}".format(x,str(i[0]),str(i[1])))
            # calculate the spread of two price series
            spread = np.array(self.history_price[str(i[0])]) - np.array(self.history_price[str(i[1])])
            mean = np.mean(spread)
            std = np.std(spread)
            ratio = self.Portfolio[i[0]].Price / self.Portfolio[i[1]].Price    
            #self.Log("pairs {}: {} {} {} {}".format(x,self.sorted_pairs[x][0], self.sorted_pairs[x][1],spread[-1] > mean + self.threshold*std,
                                                             #spread[-1] < mean - self.threshold*std))
            
            # long-short position is opened when pair prices have diverged by two standard deviations
            if spread[-1] > mean + self.threshold * std:
                if not self.Portfolio[i[0]].Invested and not self.Portfolio[i[1]].Invested:
                    quantity = int(self.CalculateOrderQuantity(i[0], 0.2))
                    self.Log("Will buy {} and sell {}".format(self.sorted_pairs[x][1], self.sorted_pairs[x][0]))
                    self.Sell(i[0], quantity) 
                    self.Buy(i[1],  floor(ratio*quantity))                
            
            elif spread[-1] < mean - self.threshold * std: 
                self.Log("Entered 2nd if")
                quantity = int(self.CalculateOrderQuantity(i[0], 0.2))
                if not self.Portfolio[i[0]].Invested and not self.Portfolio[i[1]].Invested:
                    self.Log("Will buy {} and sell {}".format(self.sorted_pairs[x][0], self.sorted_pairs[x][1]))
                    self.Sell(i[1], quantity) 
                    self.Buy(i[0], floor(ratio*quantity))  
                    
            # the position is closed when prices revert back
            elif self.Portfolio[i[0]].Invested and self.Portfolio[i[1]].Invested:
                    self.Log("Liquidating: {} {}".format(self.sorted_pairs[x][0], self.sorted_pairs[x][1]))
                    self.Liquidate(i[0]) 
                    self.Liquidate(i[1]) 
            
            x=x+1
            
                    
                    

    def Rebalance(self):
        # schedule the event to fire every half year to select pairs with the smallest historical distance
        if self.count % 3 == 0:
            self.Log("Symbols: ")
           

            distances = {}
            y = 0
            for i in self.symbol_pairs:
                #self.Debug("Pair {}: {} {}".format(y,str(i[0]), str(i[1])))
                #self.Debug("History {}: {} {}".format(y,self.history_price[str(i[0])], self.history_price[str(i[1])]))
                if  self.history_price[str(i[0])] and self.history_price[str(i[1])]:
                    distances[i] = Pair(i[0], i[1], self.history_price[str(i[0])],  self.history_price[str(i[1])]).distance()
                    self.sorted_pairs = sorted(distances, key = lambda x: distances[x])[:4]
                else:
                    self.Debug("Empty history")
                y = y+1
            for x in self.sorted_pairs:
                self.Log("Self.sorted: {} {}".format(str(x[0]),str(x[1])))
           
  


        self.count += 1
            
class Pair:
    def __init__(self, symbol_a, symbol_b, price_a, price_b):
        self.symbol_a = symbol_a
        self.symbol_b = symbol_b
        self.price_a = price_a
        self.price_b = price_b
    
    def distance(self):
        # calculate the sum of squared deviations between two normalized price series
        norm_a = np.array(self.price_a)/self.price_a[0]
        norm_b = np.array(self.price_b)/self.price_b[0]
        return sum((norm_a - norm_b)**2)