# -*- coding: utf-8 -*-
"""
Created on Fri Jul 18 23:23:08 2014

@author: rlabbe
"""
from math import radians, sin, cos, sqrt, exp
import numpy.random as random
import matplotlib.pyplot as plt
import filterpy.kalman as kf
import numpy as np

def ball_kf(x, y, omega, v0, dt, r=0.5, q=0.02):

    g = 9.8 # gravitational constant
    f1 = kf.KalmanFilter(dim_x=5, dim_z=2)

    ay = .5*dt**2
    f1.F = np.array ([[1, dt,  0,  0,  0],   # x   = x0+dx*dt
                      [0,  1,  0,  0,  0],   # dx  = dx
                      [0,  0,  1, dt, ay],   # y   = y0 +dy*dt+1/2*g*dt^2
                      [0,  0,  0,  1, dt],   # dy  = dy0 + ddy*dt 
                      [0,  0,  0,  0, 1]])   # ddy = -g.

    f1.H = np.array([
                [1, 0, 0, 0, 0],
                [0, 0, 1, 0, 0]])
    
    f1.R *= r
    f1.Q *= q

    omega = radians(omega)
    vx = cos(omega) * v0
    vy = sin(omega) * v0

    f1.x = np.array([[x,vx,y,vy,-9.8]]).T
    
    return f1
    
    
class BaseballPath(object):
    def __init__(self, x0, y0, launch_angle_deg, velocity_ms, noise=(1.0,1.0)): 
        """ Create baseball path object in 2D (y=height above ground)
        
        x0,y0 initial position
        launch_angle_deg angle ball is travelling respective to ground plane
        velocity_ms speeed of ball in meters/second
        noise amount of noise to add to each reported position in (x,y)
        """
        
        omega = radians(launch_angle_deg)
        self.v_x = velocity_ms * cos(omega)
        self.v_y = velocity_ms * sin(omega)

        self.x = x0
        self.y = y0

        self.noise = noise


    def drag_force (self, velocity):
        """ Returns the force on a baseball due to air drag at
        the specified velocity. Units are SI
        """
        B_m = 0.0039 + 0.0058 / (1. + exp((velocity-35.)/5.))
        return B_m * velocity


    def update(self, dt, vel_wind=0.):
        """ compute the ball position based on the specified time step and
        wind velocity. Returns (x,y) position tuple.
        """

        # Euler equations for x and y
        self.x += self.v_x*dt
        self.y += self.v_y*dt

        # force due to air drag
        v_x_wind = self.v_x - vel_wind
       
        v = sqrt (v_x_wind**2 + self.v_y**2)
        F = self.drag_force(v)

        # Euler's equations for velocity
        self.v_x = self.v_x - F*v_x_wind*dt
        self.v_y = self.v_y - 9.81*dt - F*self.v_y*dt

        return (self.x + random.randn()*self.noise[0], 
                self.y + random.randn()*self.noise[1])



def plot_ball():
    y = 1.
    x = 0.
    theta = 35. # launch angle
    v0 = 50.
    dt = 1/10.   # time step
    
    ball = BaseballPath(x0=x, y0=y, launch_angle_deg=theta, velocity_ms=v0, noise=[.3,.3])
    f1 = ball_kf(x,y,theta,v0,dt,r=1.)
    f2 = ball_kf(x,y,theta,v0,dt,r=10.)
    t = 0
    xs = []
    ys = []
    xs2 = []
    ys2 = []
    
    while f1.x[2,0] > 0:
        t += dt
        x,y = ball.update(dt)
        z = np.mat([[x,y]]).T
    
        f1.update(z)
        f2.update(z)
        xs.append(f1.x[0,0])
        ys.append(f1.x[2,0])
        xs2.append(f2.x[0,0])
        ys2.append(f2.x[2,0])    
        f1.predict() 
        f2.predict()
        
        p1 = plt.scatter(x, y, color='green', marker='o', s=75, alpha=0.5)
    
    p2, = plt.plot (xs, ys,lw=2)
    p3, = plt.plot (xs2, ys2,lw=4, c='r')
    plt.legend([p1,p2, p3], ['Measurements', 'Kalman filter(R=0.5)', 'Kalman filter(R=10)'])
    plt.show()