#!/usr/bin/env python
import numpy as np
from numpy import sin, cos, tan, exp, sqrt, sign, arctan as atan, arctan2 as atan2, abs as Abs
from builtins import min as Min, max as Max

# This file is auto-generated by ekf/codegen.py. DO NOT EDIT.


def Heaviside(x):
    return 1 * (x > 0)


def DiracDelta(x, v=1):
    return x == 0 and v or 0


def initial_state():
    x = np.float32(
        [0.0, 0.0, 1.00000, 0.0500000, 0.0300000, 1.00000, 0.0, 1.00000]
    )
    P = np.diag(
        [1.00000, 1.00000, 1.00000, 0.0100000, 0.0100000, 1.00000, 1.00000, 1.00000]
    )

    return x, P


def predict(x, P, Delta_t, u_M, u_delta):
    (v, delta, k1, k2, k3, srv_a, srv_b, srv_r) = x

    tmp0 = Abs(u_M)
    tmp1 = k2*tmp0
    tmp2 = tmp0*Heaviside(u_M)
    tmp3 = Delta_t*v
    tmp4 = Delta_t*srv_r
    tmp5 = -delta + srv_a*u_delta + srv_b

    F = np.eye(8)
    F[0, 0] += -Delta_t*(k3 + tmp1)
    F[0, 2] += Delta_t*tmp2
    F[0, 3] += -tmp0*tmp3
    F[0, 4] += -tmp3
    F[1, 1] += -tmp4
    F[1, 5] += tmp4*u_delta
    F[1, 6] += tmp4
    F[1, 7] += Delta_t*tmp5

    Q = np.zeros((8, 8))
    Q[0, 0] = 0.500000000000000
    Q[0, 1] = 1.00000000000000e-5
    Q[0, 2] = 0.00300000000000000
    Q[0, 3] = 1.00000000000000e-6
    Q[0, 4] = 1.00000000000000e-6
    x[0] += -Delta_t*(-k1*tmp2 + k3*v + tmp1*v)
    x[1] += tmp4*tmp5

    P = np.dot(F, np.dot(P, F.T)) + Delta_t * Q
    return x, P


def step(x, u, Delta_t):
    (v, delta, k1, k2, k3, srv_a, srv_b, srv_r) = x
    (u_M, u_delta) = u

    tmp0 = Abs(u_M)
    tmp1 = k2*tmp0
    tmp2 = Heaviside(u_M)
    tmp3 = tmp0*tmp2
    tmp4 = -delta + srv_a*u_delta + srv_b
    tmp5 = Delta_t*srv_r
    tmp6 = Delta_t*v
    tmp7 = sign(u_M)

    F = np.eye(8)
    F[0, 0] += -Delta_t*(k3 + tmp1)
    F[0, 2] += Delta_t*tmp3
    F[0, 3] += -tmp0*tmp6
    F[0, 4] += -tmp6
    F[1, 1] += -tmp5
    F[1, 5] += tmp5*u_delta
    F[1, 6] += tmp5
    F[1, 7] += Delta_t*tmp4

    J = np.zeros((8, 2))
    J[0, 0] = Delta_t*(k1*tmp0*DiracDelta(u_M) + k1*tmp2*tmp7 - k2*tmp7*v)
    J[1, 1] = srv_a*tmp5
    x[0] += -Delta_t*(-k1*tmp3 + k3*v + tmp1*v)
    x[1] += tmp4*tmp5
    return x, F, J


def update_IMU(x, P, g_z):
    v = x[0]
    delta = x[1]

    yk = np.float32(
        [-delta*v + g_z])

    Hk = np.float32([
        [delta, v, 0, 0, 0, 0, 0, 0]])

    Rk = np.diag([
        (v + 0.001)**2])

    S = np.dot(Hk, np.dot(P, Hk.T)) + Rk

    LL = -np.dot(yk, np.dot(np.linalg.inv(S), yk)) - 0.5 * np.log((2 * np.pi)**8 * np.linalg.det(S))
    K = np.linalg.lstsq(S, np.dot(Hk, P))[0].T
    x += np.dot(K, yk)
    KHk = np.dot(K, Hk)
    P = np.dot((np.eye(len(x)) - KHk), P)
    return x, P, LL


def update_encoders(x, P, dsdt, wperiod):
    v = x[0]
    tmp0 = -v

    yk = np.float32(
        [dsdt + tmp0, tmp0 + wperiod])

    Hk = np.float32([
        [1, 0, 0, 0, 0, 0, 0, 0],
        [1, 0, 0, 0, 0, 0, 0, 0]])

    Rk = np.diag([
        13.6900000000000, 12.2500000000000])

    S = np.dot(Hk, np.dot(P, Hk.T)) + Rk

    LL = -np.dot(yk, np.dot(np.linalg.inv(S), yk)) - 0.5 * np.log((2 * np.pi)**8 * np.linalg.det(S))
    K = np.linalg.lstsq(S, np.dot(Hk, P))[0].T
    x += np.dot(K, yk)
    KHk = np.dot(K, Hk)
    P = np.dot((np.eye(len(x)) - KHk), P)
    return x, P, LL


