# Stock-Bond Trading using Kalman Filter 

from math import floor
import numpy as np

# -------------------------------------------------
STOCK = 'QQQ'; BOND = 'TLT'; LEV = 1.00; T = 0.020;
# -------------------------------------------------

class KalmanFilterTrading(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2008, 1, 1)  
        self.SetEndDate(2021, 7, 3)
        self.SetCash(100000) 
        self.SetBrokerageModel(AlphaStreamsBrokerageModel())
        self.stock = self.AddEquity(STOCK, Resolution.Minute).Symbol
        self.bond = self.AddEquity(BOND, Resolution.Minute).Symbol
        self.kf = KalmanFilter()
        self.SetWarmUp(10, Resolution.Daily)

        self.Schedule.On(self.DateRules.EveryDay(self.stock), self.TimeRules.BeforeMarketClose(self.stock, 1), 
            self.Trade)

    
    def Trade(self):
        if self.IsWarmingUp: return 
        stock_price = self.Securities[self.stock].Price
        bond_price = self.Securities[self.bond].Price
        
        forecast_error, prediction_std_dev, hedge = self.kf.update(stock_price, bond_price)

        if forecast_error < -prediction_std_dev:
            self.SetHoldings(self.stock, LEV*hedge)
            self.SetHoldings(self.bond, LEV*(1 - hedge))

        elif forecast_error > prediction_std_dev:
            self.SetHoldings(self.stock, LEV*(1 - hedge))
            self.SetHoldings(self.bond, LEV*hedge)
            
        self.Plot("HEDGE", "hedge", float(hedge))
        self.Plot("ERROR", "forecast_error", float(forecast_error))
        self.Plot("STDEV", "prediction_std_dev", float(prediction_std_dev))


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.t = T

    def update(self, price_one, price_two):
        
        F = np.asarray([price_one, 1.0]).reshape((1, 2))
        y = price_two

        if self.R is not None:
            self.R = self.C + self.wt
        else:
            self.R = np.zeros((2, 2))

        yhat = F.dot(self.theta)
        et = y - yhat

        Qt = F.dot(self.R).dot(F.T) + self.vt
        sqrt_Qt = np.sqrt(Qt)
        
        At = self.R.dot(F.T) / Qt
        self.theta = self.theta + At.flatten() * et
        self.C = self.R - At * F.dot(self.R)
        hedge = self.t*self.theta[0]
        
        return et, sqrt_Qt, hedge     
        
        
'''
START
1/1/2008
END 
7/2/2021

STOCK = 'QQQ'; BOND = 'TLT'; LEV = 1.00; T = 0.025;

EveryDay

Sharpe Ratio
0.971
Total Trades
2401
Average Win
0.67%
Average Loss
-1.12%
Compounding Annual Return
13.774%
Drawdown
20.900%
Expectancy
0.154
Net Profit
471.759%
PSR
35.955%
Loss Rate
28%
Win Rate
72%
Profit-Loss Ratio
0.60
Alpha
0.128
Beta
0.181
Annual Standard Deviation
0.156
Annual Variance
0.024
Information Ratio
0.102
Tracking Error
0.226
Treynor Ratio
0.837
Total Fees
$18019.24
Estimated Strategy Capacity
$6700000000.00
Lowest Capacity Asset
TLT SGNKIKYGE9NP


'''