import numpy as np

class QuantumNadionCompensator(QCAlgorithm):

    def Initialize(self):
        
        self.SetStartDate(2017, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash

        _ = [self.AddCrypto(x, Resolution.Hour, Market.Bitfinex) for x in ['BTCUSD', 'ETHUSD']]
        self.SetBrokerageModel(BrokerageName.Bitfinex, AccountType.Margin)
        
        # Kalman filter parameters 
        self.theta = np.ones(2)          # Hidden state - weights of the linear model
        self.R = None                    # Variance-covariance matrix
        self.delta = 1e-6                # System noise
        self.wt = self.delta / (1 - self.delta) * np.eye(2)
        self.vt = 1e-7                   # Measurement noise
        
        # Strategy parameters
        self.invest_type = None         # long/short BTCUSD
        self.invest_portf = 0.1         # Share of the portfolio to be invested into each asset

    def OnData(self, data):
        
        # Kalman filtering update
        obs = np.array([data['ETHUSD'].Close, 1])
        y = data['BTCUSD'].Close
        
        if self.R is None:
            self.R = np.eye(2)#np.zeros((2,2))
            #self.R[0,0] = 100
        else:
            self.R = self.C + self.wt
        
        yhat = obs.dot(self.theta)
        et = y - yhat
        
        # Variance of the prediction at t=i
        Qt = obs.dot(self.R).dot(obs.T) + self.vt
        pred_err = np.sqrt(Qt)
        
        # Posterior values of states theta
        At = self.R.dot(obs.T) / Qt
        self.theta = self.theta + At.flatten() * et
        self.C = self.R - At * obs.dot(self.R)
        
        # Implement pairs trading strategy
        if self.invest_type is None:
            if et < -pred_err:
                # Undervalued BTC => long BTC, short ETH
                self.SetHoldings('BTCUSD', self.invest_portf, False)
                self.SetHoldings('ETHUSD', -self.theta[0] * self.invest_portf * data['ETHUSD'].Close / data['BTCUSD'].Close, False)
                self.invest_type = 'long'
            elif et > pred_err:
                # Overvalued BTC => long BTC, short ETH
                self.SetHoldings('BTCUSD', -self.invest_portf, False)
                self.SetHoldings('ETHUSD', self.theta[0] * self.invest_portf * data['ETHUSD'].Close / data['BTCUSD'].Close, False)
                self.invest_type = 'short'
        elif self.invest_type == 'long' and et > -pred_err or self.invest_type == 'short' and et < pred_err:
                self.SetHoldings('BTCUSD', 0, False)
                self.SetHoldings('ETHUSD', 0, False)
                self.invest_type = None