import math
import sys

import numpy as np
from filterpy.kalman import KalmanFilter
from filterpy.stats import plot_covariance

import json

from matplotlib import pyplot as plt

from models.EggLandscape import EggLandscape
from models.RobotEKF import RobotEKF
from models.BearingsOnlyTracking import BearingsOnlyTracking

kp = json.load(open(sys.argv[1]))
model = kp['model']

matrices = {k: np.array(kp[k]) for k in kp.keys() if k != 'model'}
Xf = matrices['Xf']
GT = matrices['X']
Ps = matrices['Ps']


def create_filter(model):
    if model['name'] == "BikeLandmarks":
        return RobotEKF(1, .5, [1.1, .01], .1, .015)
    elif model['name'] == "EggLandscape":
        return EggLandscape()
    elif model['name'] == "BearingsOnlyTracking":
        return BearingsOnlyTracking()
    elif model['name'] == "LinearToy":
        f = KalmanFilter(model['state_cnt'], model['measurement_models'][0]['meas_cnt'])
        f.Q = np.array([
            [0, 0, 0, 0],
            [0, 0, 0, 0],
            [0, 0, 0.01, 0],
            [0, 0, 0, 0.01],
        ])
        f.H = np.array([
            [-1, 0, 0, 0],
            [0, 1, 0, 0],
            [-.5, .5, 0, 0]
        ])
        f.F = np.array([
            [1, 0, 1, 0],
            [0, 1, 0, 1],
            [0, 0, 1, 0],
            [0, 0, 0, 1],
        ])
        return f


ellipse_step = kp['ellipse_step']
run_every = kp['run_every']

Zs = matrices['Z']

f = create_filter(model)

f.x = np.array([Xf[0]]).transpose()
f.P = Ps[0]

Rs = matrices['Rs']

Xs = []

ax = plt.gca()  # you first need to get the axis handle
ax.set_aspect(1)  # sets the height to width ratio to 1.5.

dZs = []
z_idx = 0
Xs.append(f.x)
pError = 0

for i in range(0, GT.shape[0]):
    if i % run_every != 0:
        continue

    f.predict()

    if i % ellipse_step == 0 and False:
        plot_covariance(
            (f.x[0], f.x[1]), cov=f.P[0:2, 0:2],
            std=1, facecolor='k', alpha=0.3)

    z = Zs[z_idx]
    z_idx += 1

    if len(model['measurement_models']) == 1:
        R = Rs[0]
        f.update(z[0].reshape(-1, 1), Rs[0])
    else:
        for meas_idx in range(len(model['measurement_models'])):
            R = Rs[meas_idx]
            if hasattr(f, 'predict_meas'):
                ZGt = f.predict_meas(GT[i].reshape(-1, 1), meas_idx)
                dZ = z[meas_idx].reshape(-1) - ZGt.reshape(-1)
                dZs.append(dZ)
            f.update(z[meas_idx].reshape(-1, 1), Rs[meas_idx], meas_idx)

    dP = f.P - Ps[z_idx]
    pError = np.linalg.norm(dP) ** 2

    Xs.append(f.x)

    if i % ellipse_step == 0:
        plot_covariance(
            (f.x[0], f.x[1]), cov=f.P[0:2, 0:2],
            std=1, facecolor='g', alpha=0.8)

Xs = np.array(Xs)

plt.plot(Xs[:, 0], Xs[:, 1], label="pyfilter")
plt.plot(Xf[:, 0], Xf[:, 1], '--', label="cnkalman")
plt.plot(GT[:, 0], GT[:, 1], '-.', label="GT")
plt.legend()

if '--show' in sys.argv:
    plt.show()

err = np.linalg.norm(Xf - Xs.reshape(Xf.shape)) / Xf.shape[0]
print(err, math.sqrt(pError) / Xf.shape[0])

has_error = math.sqrt(pError) / Xf.shape[0] > 1e-3 or err > 1e-3

sys.exit(-1 if has_error else 0)
