%matplotlib inline
from __future__ import division, print_function
import matplotlib.pyplot as plt
import numpy as np
from numpy.random import randn
from numpy import array, asarray

import sys
sys.path.insert(0, '../code') # allow us to import book_format
import book_format
book_format.load_style()

pos = 1
vel = 1
p_pos = 1
p_move = .5

data = [(0,pos+randn()*p_pos) for i in range(1000)]
data = np.asarray(data)

pos += vel


plt.scatter(data[:,0], data[:,1], marker=',')
plt.show()

from filterpy.kalman import KalmanFilter
from numpy.random import multivariate_normal
import matplotlib.cm as cm

vel = 20.
kf = KalmanFilter(dim_x=2, dim_z=1)
kf.F = array([[1, 1], [0, 1]], dtype=float)
kf.H = array([[1, 0]])
kf.x = array([[0., vel]]).T
kf.P *= 100
kf.R *= 10
kf.Q *= .0001

steps = 6
colors = cm.winter(np.linspace(0, .5, steps))
pcolors = cm.Reds(np.linspace(.5, 1, steps))
for i in range(steps):
    kf.predict()
    x,y = multivariate_normal(mean=kf.x.T[0], cov=kf.P, size=1000).T
    plt.scatter(x,y, marker=',', c=pcolors[i])
    
    z = (i+1)*vel
    kf.update(z)
    x,y = multivariate_normal(mean=kf.x.T[0], cov=kf.P, size=1000).T
    plt.scatter(x,y, marker='.', c=colors[i], alpha=0.7)

def update(mu1, var1, mu2, var2):
    if var1 == 0.0:
        var1=1.e-80
        
    if var2 == 0:
        var2 = 1e-80
        
    mean = (var1*mu2 + var2*mu1) / (var1+var2)
    variance = 1. / (1/var1 + 1/var2)
    return (mean, variance)



def predict(mu1, var1, mu2, var2):
    return (mu1+mu2, var1+var2)

pos = 0
pos_p = 100.
move = 1
move_p = .1
sense_p = 8

N = 10
colors = cm.winter(np.linspace(0, .5, N))
pcolors = cm.Reds(np.linspace(.5, 1, N))
for i in range(N):
    pos, pos_p = predict(pos, pos_p, move, move_p)
    x,y = multivariate_normal(mean=(pos,i), cov=[[0.001, 0],[0, p_pos]], size=1000).T
    plt.scatter(x,y, marker=',', c=colors[i])    
    print(pos, pos_p)
    
    pos, pos_p = update(pos, pos_p, i, sense_p)
    x,y = multivariate_normal(mean=(i, pos), cov=asarray([[0.001, 0],[0, pos_p]]).T, size=1000).T
    plt.scatter(x,y, marker='.', c=pcolors[i], alpha=0.5)