# This source code is public domain
# Autor: Christian Schirm

import numpy
import matplotlib.pyplot as plt

# Generate polynomial
nSteps = 301
coeff = [-50, 70, -16, 1]
sigmaNoise = 50
sigmaPrior = 100
xMax = 10

ts = numpy.linspace(0,xMax,nSteps)
deltaT = ts[1] - ts[0]
nPoly = len(coeff)
A = numpy.array([ts**i for i in range(nPoly)])
y_polynomial = coeff @ A

# Noise
numpy.random.seed(1)
noise = sigmaNoise*numpy.random.randn(nSteps)

# Add noise to the signal
y = y_polynomial + noise

# Prepare Kalman estimation
D = numpy.zeros((nPoly,nPoly))
D[(numpy.arange(nPoly-1), numpy.arange(nPoly-1)+1)] = 1
Dt = D*deltaT
F = numpy.identity(nPoly) + Dt + Dt @ Dt/2 + Dt @ Dt @ Dt/6
H = numpy.zeros((1,nPoly))
H[0,0] = 1

# Initialize Kalman estimation
x = numpy.zeros(nPoly)
components = A / nSteps
# P = sigmaPrior**2 * numpy.identity(nPoly)
P = sigmaPrior**2 * numpy.linalg.inv(components @ components.T)  # Constant variance prior model

# Start Kalman iteration
yEst = []
ySigma = []
for i in range(len(y)):
    # Propagate
    if i > 0:
        x = F @ x
        P = F @ P @ F.T

    # Estimate
    K = P @ H.T @ numpy.linalg.inv(H @ P @ H.T + sigmaNoise**2)
    x = x + K @ (y[i] - H @ x)
    P = (numpy.identity(nPoly) - K @ H) @ P
    ySigma.append(P[0,0])
    yEst.append(x[0])

ySigma = numpy.sqrt(ySigma)

# Plot
plt.figure(figsize=(5,3.5))
plt.plot(ts,y_polynomial,'C3-', label='Polynom 3. Grades', zorder=1)
plt.plot(ts,y,'.-', color='C1', markersize=4, linewidth=0.4, alpha=0.6, label='Polynom + Rauschen', zorder=2)
plt.plot(ts,yEst,'C0-',  label='Kalman-Schätzung', zorder=3)
plt.fill_between(ts,y_polynomial-ySigma, y_polynomial+ySigma, color='0.2', alpha=0.17,
        label='Fehlerschätzung\n(relativ zu wahrer Kurve)', lw=0, zorder=0)
plt.xlabel('Zeit')
plt.legend(loc=4)
plt.tight_layout()
plt.savefig('Kalman_Polynom_Test.svg')
plt.savefig('Kalman_Polynom_Test.png')
# plt.show()