#region imports
from AlgorithmImports import *

from math import floor
from collections import deque
import itertools as it
#endregion
# https://quantpedia.com/Screener/Details/55


class PairsTradingwithCountryETFsAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2011, 1, 1)  
        self.set_end_date(2018, 8, 1)  
        self.set_cash(100000) 
        # choose ten sector ETFs
        tickers =  ["GAF",  # SPDR S&P Emerging Middle East & Africa ETF 2007.4
                    "ENZL", # iShares MSCI New Zealand Investable Market Index Fund 2010.9
                    "NORW",  # Global X FTSE Norway 30 ETF 2011
                    "EWY",  # iShares MSCI South Korea Index ETF 2000.6
                    "EWP",  # iShares MSCI Spain Index ETF 1996
                    "EWD",  # iShares MSCI Sweden Index ETF 1996
                    "EWL",  # iShares MSCI Switzerland Index ETF 1996
                    "GXC",  # SPDR S&P China ETF 2007.4
                    "EWC",  # iShares MSCI Canada Index ETF 1996
                    "EWZ",  # iShares MSCI Brazil Index ETF 2000.8
                    # "AND",  # Global X FTSE Andean 40 ETF 2011.3
                    "AIA",  # iShares S&P Asia 50 Index ETF 1996
                    "EWO",  # iShares MSCI Austria Investable Mkt Index ETF 1996
                    "EWK",  # iShares MSCI Belgium Investable Market Index ETF 1996
                    "ECH",  # iShares MSCI Chile Investable Market Index ETF 2018  2008
                    # "EGPT", # Market Vectors Egypt Index ETF 2011
                    "EWJ",  # iShares MSCI Japan Index ETF 1999
                    "EZU",  # iShares MSCI Eurozone ETF 2000
                    "EWW",  # iShares MSCI Mexico Inv. Mt. Idx 2000
                    # "ERUS", # iShares MSCI Russia ETF 2011
                    "IVV",  # iShares S&P 500 Index 2001
                    "AAXJ", # iShares MSCI All Country Asia ex Japan Index ETF 2008.8
                    "EWQ",  # iShares MSCI France Index ETF 2000
                    "EWH",  # iShares MSCI Hong Kong Index ETF 1999
                    # "EPI",  # WisdomTree India Earnings ETF 2008.3
                    "EIDO",  # iShares MSCI Indonesia Investable Market Index ETF 2008.3
                    "EWI",  # iShares MSCI Italy Index ETF 1996
                    "ADRU"] # BLDRS Europe 100 ADR Index ETF 2003

        self._threshold = 0.5
        self._symbols = []
        for i in tickers:
            self._symbols.append(self.add_equity(i, Resolution.DAILY).symbol)
        
        formation_period = 121

        self._history_price = {}
        for symbol in self._symbols:
            hist = self.history([symbol.value], formation_period+1, Resolution.DAILY)
            if hist.empty: 
                self._symbols.remove(symbol)
            else:
                self._history_price[symbol.value] = deque(maxlen=formation_period)
                for tuple in hist.loc[str(symbol)].itertuples():
                    self._history_price[symbol.value].append(float(tuple.close))
                if len(self._history_price[symbol.value]) < formation_period:
                    self._symbols.remove(symbol)
                    self._history_price.pop(symbol.value)

        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._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 symbol.value in self._history_price:
                self._history_price[symbol.value].append(float(data[symbol].close)) 
        if self._sorted_pairs is None: return
        
        for i in self._sorted_pairs:
            pair = Pair(i[0], i[1], self._history_price[i[0].value],  self._history_price[i[1].value])
            index_a = pair.index_a[-1]
            index_b = pair.index_b[-1]
            delta = pair.distance()
            if index_a - index_b > self._threshold*delta:
                if not self.portfolio[pair.symbol_a].invested and not self.portfolio[pair.symbol_b].invested:
                    ratio = self.portfolio[pair.symbol_a].price / self.portfolio[pair.symbol_b].price
                    quantity = int(self.calculate_order_quantity(pair.symbol_a, 0.05))
                    self.sell(pair.symbol_a, quantity) 
                    self.buy(pair.symbol_b,  floor(ratio*quantity))                  
            
            elif index_a - index_b < -self._threshold*delta:
                if not self.portfolio[pair.symbol_a].invested and not self.portfolio[pair.symbol_b].invested:
                    ratio = self.portfolio[pair.symbol_b].price / self.portfolio[pair.symbol_a].price
                    quantity = int(self.calculate_order_quantity(pair.symbol_b, 0.05))
                    self.sell(pair.symbol_b, quantity) 
                    self.buy(pair.symbol_a, 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(pair.symbol_a) 
                    self.liquidate(pair.symbol_b)  
      
    def _rebalance(self):
        # schedule the event to fire every half year to select pairs with the smallest historical distance
        distances = {}
        for i in self._symbol_pairs:
            if i[0].value in self._history_price and i[1].value in self._history_price:
                distances[i] = Pair(i[0], i[1], self._history_price[i[0].value],  self._history_price[i[1].value]).distance()
        self._sorted_pairs = sorted(distances, key=lambda x: distances[x])[:5]


class Pair:

    def __init__(self, symbol_a, symbol_b, price_a, price_b):
        self.symbol_a = symbol_a
        self.symbol_b = symbol_b
        price_a = np.array(price_a)
        price_b = np.array(price_b)
        # compute normalized cumulative price indices 
        self.index_a = np.cumprod(price_a[1:]/price_a[:-1])
        self.index_b = np.cumprod(price_b[1:]/price_b[:-1])
        
    def distance(self):
        return 1/120*sum(abs(self.index_a - self.index_b))