#region imports
from AlgorithmImports import *

from scipy import stats
from math import floor
from datetime import timedelta
from collections import deque
import itertools as it
from decimal import Decimal
#endregion
# https://quantpedia.com/Screener/Details/12


class PairsTradingAlgorithm(QCAlgorithm):
    
    def initialize(self):
        self.set_start_date(2014,1,1)
        self.set_end_date(2018,1,1)
        self.set_cash(100000)
       
        tickers = [ 'XLK', 'QQQ', 'BANC', 'BBVA', 'BBD', 'BCH', 'BLX', 'BSBR', 'BSAC', 'SAN',
                    'CIB', 'BXS', 'BAC', 'BOH', 'BMO', 'BK', 'BNS', 'BKU', 'BBT','NBHC', 'OFG',
                    'BFR', 'CM', 'COF', 'C', 'VLY', 'WFC', 'WAL', 'WBK','RBS', 'SHG', 'STT', 'STL', 'SCNB', 'SMFG', 'STI']
                    # 'DKT', 'DB', 'EVER', 'KB', 'KEY', , 'MTB', 'BMA', 'MFCB', 'MSL', 'MTU', 'MFG', 
                    # 'PVTD', 'PB', 'PFS', 'RF', 'RY', 'RBS', 'SHG', 'STT', 'STL', 'SCNB', 'SMFG', 'STI',
                    # 'SNV', 'TCB', 'TD', 'USB', 'UBS', 'VLY', 'WFC', 'WAL', 'WBK', 'WF', 'YDKN', 'ZBK']
        self._threshold = 2
        self._symbols = []
        for i in tickers:
            self._symbols.append(self.add_equity(i, Resolution.DAILY).symbol)
        
        formation_period = 252

        self._history_price = {}
        for symbol in self._symbols:
            hist = self.history([symbol], formation_period+1, Resolution.DAILY)
            if hist.empty: 
                self._symbols.remove(symbol)
            else:
                self._history_price[str(symbol)] = deque(maxlen=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)]) < 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.add_equity("SPY", Resolution.DAILY) 
        self.schedule.on(self.date_rules.month_start("SPY"), self.time_rules.after_market_open("SPY"), self._rebalance)
        self._count = 0
        self._sorted_pairs = None
        
    def on_data(self, data):
        # Update the price series everyday
        for symbol in self._symbols:
            if data.bars.contains_key(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
        
        for i in self._sorted_pairs:
            # 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
            # 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.calculate_order_quantity(i[0], 0.1))
                    self.sell(i[0], quantity) 
                    self.buy(i[1], floor(ratio*quantity))                
            
            elif spread[-1] < mean - self._threshold * std: 
                quantity = int(self.calculate_order_quantity(i[0], 0.1))
                if not self.portfolio[i[0]].invested and not self.portfolio[i[1]].invested:
                    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.liquidate(i[0]) 
                    self.liquidate(i[1])                

    def _rebalance(self):
        # schedule the event to fire every half year to select pairs with the smallest historical distance
        if self._count % 6 == 0:
            distances = {}
            for i in self._symbol_pairs:
                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]
        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)