import numpy as np
from math import floor

class KalmanFilter:
    def __init__(self):
        self.delta = 1e-4
        self.wt = self.delta / (1 - self.delta) * np.eye(2)
        self.vt = 1e-3
        self.theta = np.zeros(2)
        self.P = np.zeros((2, 2))
        self.R = None
        self.qty = 2000

    def update(self, price_one, price_two):
        # Create the observation matrix of the latest prices
        # of TLT and the intercept value (1.0)
        F = np.asarray([price_one, 1.0]).reshape((1, 2))
        y = price_two

        # The prior value of the states \theta_t is
        # distributed as a multivariate Gaussian with
        # mean a_t and variance-covariance R_t
        if self.R is not None:
            self.R = self.C + self.wt
        else:
            self.R = np.zeros((2, 2))

        # Calculate the Kalman Filter update
        # ----------------------------------
        # Calculate prediction of new observation
        # as well as forecast error of that prediction
        yhat = F.dot(self.theta)
        et = y - yhat

        # Q_t is the variance of the prediction of
        # observations and hence \sqrt{Q_t} is the
        # standard deviation of the predictions
        Qt = F.dot(self.R).dot(F.T) + self.vt
        sqrt_Qt = np.sqrt(Qt)

        # The posterior value of the states \theta_t is
        # distributed as a multivariate Gaussian with mean
        # m_t and variance-covariance C_t
        At = self.R.dot(F.T) / Qt
        self.theta = self.theta + At.flatten() * et
        self.C = self.R - At * F.dot(self.R)
        hedge_quantity = int(floor(self.qty*self.theta[0]))
        
        return et, sqrt_Qt, hedge_quantity

import numpy as np
from math import floor
from KalmanFilter import KalmanFilter
class VerticalParticleInterceptor(QCAlgorithm):

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

        self.symbols = [self.AddEquity(x, Resolution.Minute).Symbol for x in ['VIA', 'VIAB']]
        self.kf = KalmanFilter()
        self.invested = None
    
        self.Schedule.On(self.DateRules.EveryDay('VIA'), self.TimeRules.BeforeMarketClose('VIA', 5), self.UpdateAndTrade)
    
    def UpdateAndTrade(self):
        
        # Get recent price and holdings information
        via = self.CurrentSlice[self.symbols[0]].Close
        viab = self.CurrentSlice[self.symbols[1]].Close
        holdings = self.Portfolio[self.symbols[0]]
        
        forecast_error, prediction_std_dev, hedge_quantity = self.kf.update(via, viab)
        
        if not holdings.Invested:
            # Long the spread
            if forecast_error < -prediction_std_dev:
                insights = Insight.Group([Insight(self.symbols[0], timedelta(1), InsightType.Price, InsightDirection.Down),
                                           Insight(self.symbols[1], timedelta(1), InsightType.Price, InsightDirection.Up)])
                self.EmitInsights(insights)
                self.MarketOrder(self.symbols[1], self.kf.qty)
                self.MarketOrder(self.symbols[0], -hedge_quantity)

            # Short the spread
            elif forecast_error > prediction_std_dev:
                insights = Insight.Group([Insight(self.symbols[0], timedelta(1), InsightType.Price, InsightDirection.Up),
                                           Insight(self.symbols[1], timedelta(1), InsightType.Price, InsightDirection.Down)])
                self.EmitInsights(insights)
                self.MarketOrder(self.symbols[1], -self.kf.qty)
                self.MarketOrder(self.symbols[0], hedge_quantity)

        if holdings.Invested:
            # Close long position
            if holdings.IsShort and (forecast_error >= -prediction_std_dev):
                insights = Insight.Group([Insight(self.symbols[0], timedelta(1), InsightType.Price, InsightDirection.Flat),
                                           Insight(self.symbols[1], timedelta(1), InsightType.Price, InsightDirection.Flat)])
                self.EmitInsights(insights)
                self.Liquidate()
                self.invested = None
            
            # Close short position
            elif holdings.IsLong and (forecast_error <= prediction_std_dev):
                insights = Insight.Group([Insight(self.symbols[0], timedelta(1), InsightType.Price, InsightDirection.Flat),
                                           Insight(self.symbols[1], timedelta(1), InsightType.Price, InsightDirection.Flat)])
                self.EmitInsights(insights)
                self.Liquidate()
                self.invested = None