#region imports
from AlgorithmImports import *
from arch.unitroot.cointegration import engle_granger
from pykalman import KalmanFilter
#endregion

class PairsTrade(QCAlgorithm):
    
    def Initialize(self):
        # Required: Five years of backtest history
        self.SetStartDate(2019, 1, 1)
        
        # Required: Alpha Streams Models:
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage)
        
        # Required: Significant AUM Capacity
        self.SetCash(1000000)
        
        # Required: Benchmark to SPY
        self.SetBenchmark("SPY")
        
        self.assets = ["XLK", "XLU"]
        
        # Add Equity ------------------------------------------------ 
        for asset in self.assets:
            self.AddEquity(asset, Resolution.Minute)
            
        # Add Options ------------------------------------------------ 
        self.option_symbols = []
        for asset in self.assets:
            option = self.AddOption(asset)
            option.SetFilter(-2, +2, 90, 180)  # ATM options with expiration between 90 and 180 days
            self.option_symbols.append(option.Symbol)
            
        # Instantiate our model
        self.recalibrate()
        
        # Set a variable to indicate the trading bias of the portfolio
        self.state = 0
        
        # Set Scheduled Event Method For Kalman Filter updating.
        self.Schedule.On(self.DateRules.WeekStart(), self.TimeRules.At(0, 0), self.recalibrate)
        
        # Set Scheduled Event Method For trading every 2 days.
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), self.TimeRules.BeforeMarketClose("XLK"), self.every_other_day_before_market_close)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Wednesday), self.TimeRules.BeforeMarketClose("XLK"), self.every_other_day_before_market_close)
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday), self.TimeRules.BeforeMarketClose("XLK"), self.every_other_day_before_market_close)
            
        self.option_contracts = {}
    
    def recalibrate(self):
        qb = self
        history = qb.History(self.assets, 252*2, Resolution.Daily)
        if history.empty: return
        
        # Select the close column and then call the unstack method
        data = history['close'].unstack(level=0)
        
        # Convert into log-price series to eliminate compounding effect
        log_price = np.log(data)
        
        ### Get Cointegration Vectors
        # Get the cointegration vector
        coint_result = engle_granger(log_price.iloc[:, 0], log_price.iloc[:, 1], trend="c", lags=0)
        coint_vector = coint_result.cointegrating_vector[:2]
        
        # Get the spread
        spread = log_price @ coint_vector
        
        ### Kalman Filter
        # Initialize a Kalman Filter. Using the first 20 data points to optimize its initial state. We assume the market has no regime change so that the transitional matrix and observation matrix is [1].
        self.kalman_filter = KalmanFilter(transition_matrices = [1],
                                          observation_matrices = [1],
                                          initial_state_mean = spread.iloc[:20].mean(),
                                          observation_covariance = spread.iloc[:20].var(),
                                          em_vars=['transition_covariance', 'initial_state_covariance'])
        self.kalman_filter = self.kalman_filter.em(spread.iloc[:20], n_iter=5)
        (filtered_state_means, filtered_state_covariances) = self.kalman_filter.filter(spread.iloc[:20])
        
        # Obtain the current Mean and Covariance Matrix expectations.
        self.current_mean = filtered_state_means[-1, :]
        self.current_cov = filtered_state_covariances[-1, :]
        
        # Initialize a mean series for spread normalization using the Kalman Filter's results.
        mean_series = np.array([None]*(spread.shape[0]-20))
        
        # Roll over the Kalman Filter to obtain the mean series.
        for i in range(20, spread.shape[0]):
            (self.current_mean, self.current_cov) = self.kalman_filter.filter_update(filtered_state_mean = self.current_mean,
                                                                                    filtered_state_covariance = self.current_cov,
                                                                                    observation = spread.iloc[i])
            mean_series[i-20] = float(self.current_mean)
        
        # Obtain the normalized spread series.
        normalized_spread = (spread.iloc[20:] - mean_series)
        
        ### Determine Trading Threshold
        # Initialize 50 set levels for testing.
        set_levels = self.GetParameter("set_levels", 50)
        
        s0 = np.linspace(0, max(normalized_spread), set_levels)
        
        # Calculate the profit levels using the 50 set levels.
        f_bar = np.array([None]*set_levels)
        for i in range(set_levels):
            f_bar[i] = len(normalized_spread.values[normalized_spread.values > s0[i]]) / normalized_spread.shape[0]
            
        # Set trading frequency matrix.
        D = np.zeros((set_levels-1, set_levels))
        for i in range(D.shape[0]):
            D[i, i] = 1
            D[i, i+1] = -1
            
        # Set level of lambda.
        l = 1.0
        
        # Obtain the normalized profit level.
        f_star = np.linalg.inv(np.eye(set_levels) + l * D.T@D) @ f_bar.reshape(-1, 1)
        s_star = [f_star[i]*s0[i] for i in range(set_levels)]
        self.threshold = s0[s_star.index(max(s_star))]
        
        # Set the trading weight. We would like the portfolio absolute total weight is 1 when trading.
        self.trading_weight = coint_vector / np.sum(abs(coint_vector))
    
    def every_other_day_before_market_close(self):
        qb = self
        
        # Get the real-time log close price for all assets and store in a Series
        series = pd.Series()
        for symbol in qb.Securities.Keys:
            series[symbol] = np.log(qb.Securities[symbol].Close)
            
        # Get the spread
        spread = np.sum(series * self.trading_weight)
        
        # Update the Kalman Filter with the Series
        (self.current_mean, self.current_cov) = self.kalman_filter.filter_update(filtered_state_mean = self.current_mean,
                                                                                 filtered_state_covariance = self.current_cov,
                                                                                 observation = spread)
            
        # Obtain the normalized spread.
        normalized_spread = spread - self.current_mean
    
        # ==============================
        
        # Mean-reversion
        if normalized_spread < -self.threshold:
            self.trade_options("call", "put")
            self.state = 1
                
        elif normalized_spread > self.threshold:
            self.trade_options("put", "call")
            self.state = -1
                
        # Out of position if spread recovered
        elif self.state == 1 and normalized_spread > -self.threshold or self.state == -1 and normalized_spread < self.threshold:
            self.Liquidate()
            self.state = 0
    
    def OnData(self, slice):
        # Store option chains
        for symbol in self.option_symbols:
            if symbol in slice.OptionChains:
                self.option_contracts[symbol] = slice.OptionChains[symbol]
    
    def trade_options(self, option_type_asset1, option_type_asset2):
        contracts = []
        for i, asset in enumerate(self.assets):
            symbol = self.option_symbols[i]
            if symbol in self.option_contracts:
                chain = self.option_contracts[symbol]
                contract = self.get_atm_option(chain, option_type_asset1 if i == 0 else option_type_asset2)
                if contract:
                    contracts.append(contract)
        
        if len(contracts) == 2:
            # Calculate the initial margin requirement for each contract
            initial_margin_0 = self.get_initial_margin_requirement(contracts[0], 1).Value
            initial_margin_1 = self.get_initial_margin_requirement(contracts[1], 1).Value

            # Calculate the weights for each asset
            weight_0 = abs(self.trading_weight[0])
            weight_1 = abs(self.trading_weight[1])

            # Calculate the maximum number of contracts based on 10% of the remaining buying power
            buying_power = self.Portfolio.Cash * 0.1
            total_margin_needed = weight_0 * initial_margin_0 + weight_1 * initial_margin_1
            max_num_contracts = buying_power / total_margin_needed

            num_contracts_0 = int(max_num_contracts * weight_0)
            num_contracts_1 = int(max_num_contracts * weight_1)

            if num_contracts_0 > 0 and num_contracts_1 > 0:
                # Calculate limit prices based on current market prices
                limit_price_0 = self.Securities[contracts[0].Symbol].Price
                limit_price_1 = self.Securities[contracts[1].Symbol].Price

                if option_type_asset1 == "call" and option_type_asset2 == "put":
                    self.LimitOrder(contracts[0].Symbol, num_contracts_0, limit_price_0)  # Buy ATM call for the undervalued asset
                    self.LimitOrder(contracts[1].Symbol, num_contracts_1, limit_price_1)  # Buy ATM put for the overvalued asset
                elif option_type_asset1 == "put" and option_type_asset2 == "call":
                    self.LimitOrder(contracts[0].Symbol, num_contracts_0, limit_price_0)  # Buy ATM put for the undervalued asset
                    self.LimitOrder(contracts[1].Symbol, num_contracts_1, limit_price_1)  # Buy ATM call for the overvalued asset
    
    def get_initial_margin_requirement(self, contract, quantity):
        parameters = InitialMarginParameters(self.Securities[contract.Symbol], quantity)
        initial_margin = self.Securities[contract.Symbol].BuyingPowerModel.GetInitialMarginRequirement(parameters)
        return initial_margin

    def get_atm_option(self, chain, option_type):
        if option_type == "call":
            calls = sorted([x for x in chain if x.Right == OptionRight.Call and x.Strike <= chain.Underlying.Price], key=lambda x: x.Strike)
            if calls:
                return calls[-1]
        elif option_type == "put":
            puts = sorted([x for x in chain if x.Right == OptionRight.Put and x.Strike >= chain.Underlying.Price], key=lambda x: x.Strike)
            if puts:
                return puts[0]
        return None