# -*- coding: utf-8 -*-

# Copyright (C) 2016 Ivan Gorinov

import numpy as np
import matplotlib.pyplot as plt
import random

from filterpy.kalman import KalmanFilter
from filterpy.common import Q_discrete_white_noise

class SequentialKF(KalmanFilter):
    def update_sequential(self, start, z_i, R_i=None, H_i=None):
        """
        Add a single input measurement (z_i) to the Kalman filter.
        In sequential processing, inputs are processed one at a time.

        Parameters
        ----------
        start : integer
            Index of the first measurement input updated by this call.

        z_i : np.array or scalar
            Measurement of inputs for this partial update.

        R_i : np.array, scalar, or None
            Optionally provide R_i to override the measurement noise of
            inputs for this one call, otherwise a slice of self.R will
            be used.

        H_i : np.array, or None
            Optionally provide H[i] to override the partial measurement
            function for this one call, otherwise a slice of self.H will
            be used.
        """

        if isscalar(z_i):
            length = 1
        else:
            length = len(z_i)
        z_i = np.reshape(z_i, [length, 1])
        stop = start + length

        if R_i is None:
            R_i = self.R[start:stop, start:stop]
        elif isscalar(R_i):
            R_i = eye(length) * R_i

        if H_i is None:
            H_i = self.H[start:stop]

        H_i = np.reshape(H_i, [length, self.dim_x])

        # error (residual) between measurement and prediction
        y_i = z_i - H_i @ self.x
        self.y[start:stop] = y_i

        # common subexpression for speed
        PHT = self.P @ H_i.T

        # project system uncertainty into the measurement subspace
        S_i = H_i @ PHT + R_i

        if length == 1:
            K_i = PHT * (1.0 / S_i)
        else:
            K_i = dot(PHT, linalg.inv(S_i))

        self.K[:,start:stop] = K_i
        I_KH = self._I - K_i @ H_i

        # x = x + K_i @ y_i
        # update state estimation with residual scaled by the kalman gain
        self.x += dot(K_i, y_i)

        # compute the posterior covariance
        self.P = I_KH @ self.P @ I_KH.T + K_i @ R_i @ K_i.T

        # save measurement component #i and the posterior state
        self.z[start:stop] = z_i
        self.x_post = self.x.copy()
        self.P_post = self.P.copy()

random.seed(65537)

#                                  __
# Accelerometer noise density, µg/√Hz
# (BMI160)
nd_a = 300

# Accelerometer bias, m/s²
bias_a = 0.01

# Earth gravity (at ±45° latitude)
g_n = 9.80665

# Accelerometer sample rate, Hz
fs = 100.0

# Barometer sample rate, Hz
f_b = 25.0

# Accelerometer sample period, s
dt = 1.0 / fs

n = 900
r_n = 1.0 / n

# Standard deviation of sensor readings (RMS noise)
std_a = nd_a * 1e-6 * g_n * np.sqrt(fs / 2)
std_h = 0.1

print(std_a)

kf1 = KalmanFilter(dim_x=2, dim_z=1)
kf1.H = np.array([[1, 0]])
kf1.F = np.array([[1, dt], [0, 1]])

kf2 = KalmanFilter(dim_x=3, dim_z=2)
kf2.H = np.array([[1, 0, 0], [0, 0, 1]])
kf2.F = np.array([[1, dt, dt * dt * 0.5], [0, 1, dt], [0, 0, 1]])

# initial process covariance
kf1.P = np.array([[1, 0], [0, 0.0004]])
kf2.P = np.array([[1, 0, 0], [0, 0.0001, 0], [0, 0, std_a * std_a]])

# Process noise matrix
std = 0.0039
var = std * std
kf1.Q = Q_discrete_white_noise(dim=2, dt=dt, var=0.03125)
kf2.Q = Q_discrete_white_noise(dim=3, dt=dt, var=var)

# Measurement covariance
kf1.R *= np.array([[std_h * std_h]])
kf2.R *= np.array([[std_h * std_h, 0], [0, std_a * std_a]])

n_b = int(n / 4)
tt = np.zeros(n)
t_a = np.zeros(n)
t_b = np.zeros(n_b)
h_sim = np.zeros(n)
v_sim = np.zeros(n)
a_sim = np.zeros(n)
measured_h = np.zeros(n_b)
measured_a = np.zeros(n)
v_est_kf1 = np.zeros(n)
v_est_kf2 = np.zeros(n)

v = 0
h = 0

j = 0
for i in range(n):
    t = (i * 0.01)

    tt[i] = t
    t_a[i] = t

    # One barometer sample per 4 accelerometer samples
    if (i % 4 == 0):
        t_b[j] = t
        measured_h[j] = h + random.gauss(0, std_h)
        j += 1

    a = 0
    if t >= 1 and t < 2:
        a = (1 - np.cos((t - 1) * 2 * np.pi)) / 2
    if t >= 3 and t < 4:
        a = (np.cos((t - 3) * 2 * np.pi) - 1) / 2
    if t >= 5 and t < 6:
        a = (np.cos((t - 5) * 2 * np.pi) - 1) / 2
    if t >= 7 and t < 8:
        a = (1 - np.cos((t - 7) * 2 * np.pi)) / 2

    a_sim[i] = (a)
    v += a * dt
    v_sim[i] = v
    h += v * dt + (a * dt * dt) / 2
    h_sim[i] = h
    measured_a[i] = a + random.gauss(0, std_a) + bias_a

# Compute the speed estimations

t = 0
j = 0
for i in range(n):
    kf1.predict()
    v_est_kf1[i] = kf1.x[1]

    kf2.predict()

    if (i % 4 == 0):
        kf1.update(np.array([measured_h[j]]))
        kf2.update_sequential(0, [measured_h[j]])
        j += 1

    kf2.update_sequential(1, [measured_a[i]])

    v_est_kf2[i] = kf2.x[1]


# Plot the results

plt.figure(1, figsize=(10, 15), dpi=80)

plt.subplot(311)
plt.axis([0, n * dt, -1.25, +1.25])
plt.plot(t_a, measured_a, 'y+')
plt.plot(tt, a_sim, 'r')
plt.title('Vertical acceleration, m/s²')
plt.legend(('Measurement', 'True'), loc='best')

plt.subplot(312)
plt.axis([0, n * dt, -0.5, +1.5])
plt.plot(t_b, measured_h, 'c+')
plt.plot(tt, h_sim, 'b')
plt.title('Altitude, m')
plt.legend(('Measurement', 'True'), loc='best')

plt.subplot(313)
plt.axis([0, n * dt, -0.75, +0.75])
plt.plot(tt, v_sim, 'k', linewidth=0.5)
plt.plot(tt, v_est_kf1, 'g')
plt.plot(tt, v_est_kf2, 'y')
plt.title('Vertical speed, m/s')
plt.legend(('True', 'Kalman filter (1 sensor)', 'Kalman filter (2 sensors)'), loc='best')

plt.savefig('kalman_vsi.svg')
plt.show()