import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat
from numpy.linalg import inv
%matplotlib inline

"""
Explanation: Assignment 7
Lorenzo Biasi,  Julius Vernie
End of explanation
"""


data = loadmat('data_files/Tut7_file1.mat')
locals().update(data)
data.keys()

p, T = z.shape

mu = np.zeros(z.shape)
K = np.zeros((4, 4, T))
V = np.zeros((4, 4, T))
L = np.zeros((4, 4, T))

K[...,0] = L0.dot(B.T.dot(inv(B.dot(L0.dot(B.T)) + Gamma)))
mu[..., [0]] = A.dot(mu0) + K[..., 0].dot(x[:, [0]] - B.dot(A.dot(mu0))) + C.dot(u[..., [0]])
V[..., 0] = (np.eye(4) - K[..., 0].dot(B)).dot(L0)
L[..., 0] = A.dot(V[..., 0].dot(A.T)) + Sigma

"""
Explanation: We load the variables and initilize the parameters we need
End of explanation
"""


for t in range(1, T):
    K[...,t] = L[..., t - 1].dot(B.T.dot(inv(B.dot(L[..., t - 1].dot(B.T)) + Gamma)))
    mu[..., [t]] = A.dot(mu[..., [t-1]]) + K[..., t].dot(x[:, [t]] - B.dot(A.dot(mu[..., [t-1]]))) + C.dot(u[..., [t]])
    V[..., t] = (np.eye(4) - K[..., t].dot(B)).dot(L[..., t-1])
    L[..., t] = A.dot(V[..., t].dot(A.T)) + Sigma

"""
Explanation: We run the filter
End of explanation
"""


plt.plot(mu.T)
plt.plot(z.T, color='red')

V_tilde = np.zeros(V.shape)
mu_tilde = np.zeros(mu.shape)
V_tilde[..., -1] = V[..., -1]
mu_tilde[..., [-1]] = mu[..., [-1]]

for t in range(T - 2, -1, -1):
    #print(t)
    W = V[..., t].dot(A.T.dot(inv(L[..., t])))
    V_tilde[..., t] = V[..., t] + W.dot(V_tilde[..., t+1] - L[..., t]).dot(W.T)
    mu_tilde[..., [t]] = mu[..., [t]] + W.dot(mu_tilde[..., [t+1]] - A.dot(mu[..., [t]]))

"""
Explanation: We can see a slight offset, we would expect that to be solved with the smoother step
End of explanation
"""


plt.plot(mu_tilde.T)
plt.plot(z.T, color='red')

"""
Explanation: we can see that the offset is still present and slightly worse
End of explanation
"""


print ('Non smoothed result:', np.sum((mu - z).T ** 2))
print('Smoothed result:', np.sum((mu_tilde - z).T ** 2))

print('Ratio, \n', np.sum((mu_tilde - z).T ** 2) / np.sum((mu - z).T ** 2))


"""
Explanation: The predition is clearly following the data more or less correctly, but there is a probelm with the offset, that makes $\tilde{\mu}$ worse than our $\mu$. This should not happen, we would rather expect the opposite.
End of explanation
"""


plt.plot(x.T)

"""
Explanation: After checking the algorithm many times i decided to look at our x to see if there was anything strange. And if you look closely at the first time steps there is some oddity.
End of explanation
"""


#plt.plot(x.T[:4, :])
plt.plot(np.diff(x[..., :10]).T)
np.diff(x[..., :4])

"""
Explanation: If someone looks at how x varies at the first time step you will see that it is almost constant and than it starts changing. this could explain the offset in our predictions.
End of explanation
"""


T = 99
z = z[:, :-1]
mu = np.zeros(z.shape)
K = np.zeros((4, 4, T))
V = np.zeros((4, 4, T))
L = np.zeros((4, 4, T))

K[...,0] = L0.dot(B.T.dot(inv(B.dot(L0.dot(B.T)) + Gamma)))
mu[..., [0]] = mu0
V[..., 0] = 0
L[..., 0] = L0

for t in range(1, T):
    #print(t)
    K[...,t] = L[..., t - 1].dot(B.T.dot(inv(B.dot(L[..., t - 1].dot(B.T)) + Gamma)))
    mu[..., [t]] = A.dot(mu[..., [t-1]]) + K[..., t].dot(x[:, [t + 1]] - B.dot(A.dot(mu[..., [t-1]]))) + C.dot(u[..., [t]])
    V[..., t] = (np.eye(4) - K[..., t].dot(B)).dot(L[..., t-1])
    L[..., t] = A.dot(V[..., t].dot(A.T)) + Sigma

plt.plot(mu.T)
plt.plot(z.T, color='red')
np.sum((mu - z)**2)
A.dot(mu[..., [t-1]]) + K[..., t].dot(x[:, [t + 1]] - B.dot(A.dot(mu[..., [t-1]]))) + C.dot(u[..., [t]])

V_tilde = np.zeros(V.shape)
mu_tilde = np.zeros(mu.shape)
V_tilde[..., -1] = V[..., -1]
mu_tilde[..., [-1]] = mu[..., [-1]]

for t in range(T - 2, -1, -1):
    W = V[..., t].dot(A.T.dot(inv(L[..., t])))
    V_tilde[..., t] = V[..., t] + W.dot(V_tilde[..., t+1] - L[..., t]).dot(W.T)
    mu_tilde[..., [t]] = mu[..., [t]] + W.dot(mu_tilde[..., [t+1]] - A.dot(mu[..., [t]]))

plt.plot(mu_tilde.T)
plt.plot(z.T)

print ('Non smoothed result:', np.sum((mu - z).T ** 2))
print('Smoothed result:', np.sum((mu_tilde - z).T ** 2))

print('Ratio, \n', np.sum((mu_tilde - z).T ** 2) / np.sum((mu - z).T ** 2))

"""
Explanation: To test my hunch I decided to remove one time step from the data, to make sure that $x_1$ was not used in the prediction.
End of explanation
"""


