#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(self.end_date - timedelta(5*365))
        self.set_cash(100_000) 
        # Define some parameters.
        self._threshold = 0.5
        formation_period = 121
        
        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
        ] 
        # Add the Country ETFs.
        for ticker in tickers:
            security = self.add_equity(ticker, Resolution.DAILY)
            # Add a RollingWindow to track historical prices over the formation period.
            security.history = RollingWindow[float](formation_period+1)
            # Warm up the RollingWindow
            for bar in self.history[TradeBar](security, formation_period+1):
                security.history.add(bar.close)
            # Add a consolidator to update the RollingWindow each day going forward.
            self.consolidate(
                security, Resolution.DAILY, 
                lambda bar: self.securities[bar.symbol].history.add(bar.close)
            )
        # Define all the possible pairs of ETFs.
        self._pairs = [Pair(self, a, b) for a, b in list(it.combinations(self.securities.values(), 2))]
        self._selected_pairs = []
        # Add a Scheduled Event to select pairs at the start of each month.
        self.schedule.on(self.date_rules.month_start("SPY"), self.time_rules.midnight, self._select_pairs)

    def on_data(self, data):
        # Trade the selected pairs.
        for pair in self._selected_pairs:
            index_a = pair.a_norm[-1]
            index_b = pair.b_norm[-1]
            delta = pair.distance()
            if index_a - index_b > self._threshold*delta:
                pair.enter(-1, 1)
            elif index_a - index_b < -self._threshold*delta:
                pair.enter(1, -1)
            else:  # Close the position when prices revert back.
                pair.exit()
      
    def _select_pairs(self):
        # At the start of each month, select pairs with the 
        # smallest historical distance. 
        ready_pairs = [p for p in self._pairs if p.is_ready]       
        new_selected_pairs = sorted(ready_pairs, key=lambda p: p.distance())[:5]
        for pair in self._selected_pairs:
            if pair not in new_selected_pairs:
                pair.exit()
        self._selected_pairs = new_selected_pairs


class Pair:

    def __init__(self, algorithm, security_a, security_b):
        self._algorithm = algorithm
        self._a = security_a
        self._b = security_b
        self._ticket_by_security = {}

    @property
    def is_ready(self):
        return self._a.history.is_ready and self._b.history.is_ready
        
    def distance(self):
        a = np.array(list(self._a.history)[::-1])
        b = np.array(list(self._b.history)[::-1])
        # compute normalized cumulative price indices 
        self.a_norm = a[1:] / a[0]
        self.b_norm = b[1:] / b[0]
        return 1/120*sum(abs(self.a_norm - self.b_norm))

    def enter(self, direction_a, direction_b):
        if self._ticket_by_security:
            # If the direction of each security is the same, do nothing.
            if np.sign(self._ticket_by_security[self._a].quantity) == direction_a:
                return
            # Otherwise exit the current positions.
            self.exit()
        # Enter the new position.
        order_size = 0.1 * self._algorithm.portfolio.total_portfolio_value
        quantity_a = direction_a * int(order_size / self._a.price)
        quantity_b = direction_b * int(order_size / self._b.price)
        if quantity_a and quantity_b:
            self._ticket_by_security = {
                self._a: self._algorithm.market_order(self._a, quantity_a),
                self._b: self._algorithm.market_order(self._b, quantity_b)
            }
        
    def exit(self):
        for security, ticket in self._ticket_by_security.items():
            self._algorithm.market_order(security, -ticket.quantity)
        self._ticket_by_security.clear()