import numpy as np
def kf(Jp, hp, y, A, b, S_d, H, S_o):
    m, n = H.shape; Q = np.zeros([2*n+m, 2*n+m])
    Q[:n, :n] = Jp; Q[n:, :n][:n, :] = Q.T[n:, :n][:n, :] = A
    Q[n:-m, n:-m] = S_d; Q[-m:, -m:] = S_o
    Qinv = np.linalg.inv(Q); G = np.block([-np.eye(n), H.T])
    return (- G @ Qinv[n:, n:] @ G.T, - G @ Qinv[n:, :] @ np.block([hp, -b, y]))

# J  A'  0    0 | h
# A -s_d 0   -I | 0
# 0  0  -s_o  H | y
#   -I   H'   _ | _

if __name__ == '__main__':
    import scipy.linalg as la
    n, m = 4, 2
    A = np.array([[1, 0, 1, 0],
                  [0, 1, 0, 1],
                  [0, 0, 1, 0],
                  [0, 0, 0, 1]])
    H = np.array([[1, 0, 0, 0],
                  [0, 1, 0, 0]])
    Q = np.eye(n)
    R = 10 * np.eye(m)


    Jn = 1/10 * np.eye(n); hn = Jn @ np.array([10, 10, 1, 0])
    xn = la.cholesky(la.inv(Jn), lower=True) @ np.random.randn(n)

    trace = []
    for ii in range(100):
        xn = A @ xn + la.cholesky(Q, lower=True) @ np.random.randn(n)
        yn = H @ xn + la.cholesky(R, lower=True) @ np.random.randn(m)
        #
        Jn, hn = kf(Jn, hn, yn, A, np.zeros(n), Q, H, R)
        #
        trace.append((xn, la.solve(Jn, hn)))

    trace = np.array(trace)

    import matplotlib.pyplot as plt

    plt.plot(*trace[:, 0, [0, 1]].T, label='actual')
    plt.plot(*trace[:, 1, [0, 1]].T, label='inferred')
    plt.legend()
    plt.show()