# region imports
from math import floor

from AlgorithmImports import *
# endregion

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

    def update(self, price_one, price_two, quantity):
        # 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(quantity*self._theta[0]))
        
        return et, sqrt_Qt, hedge_quantity

# region imports
from AlgorithmImports import *
from KalmanFilter import KalmanFilter
# endregion


class VerticalParticleInterceptor(QCAlgorithm):

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5*365))
        self.set_cash(100_000)
        # Add the asests to trade.
        self._security_a = self.add_equity('UA')
        self._security_b = self.add_equity('UAA')
        # Add the kalman filter.
        self.kf = KalmanFilter()
        # Add a Scheduled Event to update the kalman filter and place orders.
        self.schedule.on(
            self.date_rules.every_day(self._security_a), 
            self.time_rules.before_market_close(self._security_a, 5), 
            self._update_and_trade
        )
    
    def _update_and_trade(self):
        # Get recent price and holdings information.
        quantity = int(self.portfolio.total_portfolio_value / 2 / self._security_b.price)
        forecast_error, prediction_std_dev, hedge_quantity = self.kf.update(self._security_a.price, self._security_b.price, quantity)
        # Check for entries.
        if not self.portfolio.invested and hedge_quantity:
            # Long the spread.
            if forecast_error < -prediction_std_dev:
                insights = Insight.group([
                    Insight.price(self._security_a, timedelta(1), InsightDirection.DOWN),
                    Insight.price(self._security_b, timedelta(1), InsightDirection.UP)
                ])
                self.emit_insights(insights)
                self.market_order(self._security_b, quantity)
                self.market_order(self._security_a, -hedge_quantity)
            # Short the spread
            elif forecast_error > prediction_std_dev:
                insights = Insight.group([
                    Insight.price(self._security_a, timedelta(1), InsightDirection.UP),
                    Insight.price(self._security_b, timedelta(1), InsightDirection.DOWN)
                ])
                self.emit_insights(insights)
                self.market_order(self._security_b, -quantity)
                self.market_order(self._security_a, hedge_quantity)
        # Check for exits.
        if self.portfolio.invested:
            if (self._security_a.holdings.is_short and (forecast_error >= -prediction_std_dev) or
                self._security_a.holdings.is_long and (forecast_error <= prediction_std_dev)):
                insights = Insight.group([
                    Insight.price(self._security_a, timedelta(1), InsightDirection.FLAT),
                    Insight.price(self._security_b, timedelta(1), InsightDirection.FLAT)
                ])
                self.emit_insights(insights)
                self.liquidate()