from __future__ import print_function
import numpy as np
from copy import copy
from matplotlib import pyplot as plt
from noiseestimation.playback_sensor import PlaybackSensor
from bicycle_ukf import BicycleUKFBias

# parameters
skip_samples = 500
used_taps = 100
measurement_var = 1e-5
R_proto = np.array([[1, 0],
                    [0, 3]])
sim_var = 1e-3
# num_samples = skip_samples + 1600
num_samples = 11670 / 3
dt = 0.01


def setup():
    sim = PlaybackSensor("data/vehicle_state.json",
                         fields=["fYawrate", "fVx"],
                         control_fields=["fStwAng", "fAx"])
    # set up kalman filter
    tracker = BicycleUKFBias(dt)
    Q_factor = np.array(1e-6)
    tracker.Q = np.diag([1, 1, 1, 1e-6] * Q_factor)
    tracker.R = R_proto * (sim_var + measurement_var) * 1
    tracker.x = np.array([0, 0, 1, 1*np.pi/180]).T
    tracker.P = np.eye(4) * 1e0
    tracker.P[3, 3] = 1e-3

    return sim, tracker


def filtering(sim, tracker):
    # perform sensor simulation and filtering
    Rs = [R_proto * sim_var] * num_samples
    readings, filtered, residuals, Ps, Fs, Ks = [], [], [], [], [], []
    for idx, R in enumerate(Rs):
        time, reading = sim.read(R)
        measurement = reading[0:2]
        controls = reading[2:]
        # skip low velocities
        if measurement[1, 0] < 0.5:
            continue
        tracker.predict(fx_args=controls)
        if idx % 1 == 0:
            tracker.update(measurement[:, 0])

        readings.append(reading)
        filtered.append(copy(tracker.x[:, np.newaxis]))
        Ps.append(copy(tracker.P))
        # if not np.all(np.linalg.eigvals(tracker.P) >= 0):
        #     print(tracker.Pz)
        #     print(tracker.K)
        #     print(tracker.P)
        # residuals.append(tracker.y)
        # Ks.append(tracker.K)

    readings = np.asarray(readings)
    filtered = np.asarray(filtered)
    residuals = np.asarray(residuals)
    Ps = np.asarray(Ps)
    Fs = np.asarray(Fs)
    Ks = np.asarray(Ks)
    return readings, filtered, residuals, Ps, Fs, Ks


def plot_results(readings, filtered, Ps):
    plot_filtered_values(readings, filtered, Ps)
    # plot_position(readings, filtered)


def plot_filtered_values(readings, filtered, Ps):
    f, axarr = plt.subplots(4, 2)
    axarr[0, 0].set_title("Schwimmwinkel (deg)")
    axarr[0, 0].plot(
        filtered[:, 0] * 180.0 / np.pi,
        'go')
    axarr[0, 0].set_ylim((-20, 20))
    axarr[0, 1].set_title("Geschaetze Varianz des Schwimmwinkels")
    axarr[0, 1].plot(
        Ps[:, 0, 0]
    )
    axarr[0, 1].set_ylim((0, 0.005))

    axarr[1, 0].set_title("Gierrate (deg/s)")
    axarr[1, 0].plot(
        readings[:, 0] * 180.0 / np.pi,
        'kx'
    )
    axarr[1, 0].plot(
        # filtered[:, 1] * 180.0 / np.pi,
        (filtered[:, 1] + filtered[:, 3]) * 180.0 / np.pi,
        'r-', label="with bias")
    axarr[1, 0].plot(
        filtered[:, 1] * 180.0 / np.pi,
        '-', label="no bias")
    axarr[1, 0].legend(loc="upper right")
    axarr[1, 1].set_title("Geschaetze Varianz der Gierrate")
    axarr[1, 1].plot(
        Ps[:, 1, 1]
    )

    axarr[2, 0].set_title("Gierratenbias (deg/s)")
    axarr[2, 0].plot(
        filtered[:, 3] * 180.0 / np.pi,
        'r-')
    axarr[2, 1].set_title("Geschaetze Varianz des Gierratenbias")
    axarr[2, 1].plot(
        Ps[:, 3, 3]
    )

    axarr[3, 0].set_title("Geschwindigkeit (m/s)")
    axarr[3, 0].plot(readings[:, 1], 'kx')
    axarr[3, 0].plot(filtered[:, 2], 'b-')
    axarr[3, 1].set_title("Geschaetze Varianz der Geschwindigkeit")
    axarr[3, 1].plot(
        Ps[:, 2, 2]
    )
    plt.show()


def plot_position(readings, filtered):
    # skip last value in loops
    yaw_angles = [0] * len(filtered)
    for idx, yawrate in enumerate(filtered[:-1, 1, 0]):
        yaw_angles[idx + 1] = yaw_angles[idx] + dt * yawrate

    positions = np.zeros((len(filtered), 2))
    for idx in range(len(filtered) - 1):
        sideslip = filtered[idx, 0, 0]
        angle = yaw_angles[idx] + sideslip
        velocity = filtered[idx, 2, 0]
        delta = dt * velocity * np.array([np.cos(angle),
                                          np.sin(angle)])
        positions[idx + 1] = positions[idx] + delta

    plt.plot(positions[:, 0], positions[:, 1], 'b-')
    plt.xlabel("x (m)")
    plt.ylabel("y (m)")
    plt.show()


def run_tracker():
    sim, tracker = setup()
    readings, filtered, residuals, Ps, Fs, Ks = filtering(sim, tracker)
    plot_results(readings, filtered, Ps)


if __name__ == "__main__":
    run_tracker()
