import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.transform import Rotation

gyroMeasError = np.radians(3 / 3600)
beta = np.sqrt(3. / 4.) * gyroMeasError

nu_1 = np.radians(0 * 1 / 3600)
nu_2 = np.radians(0 * 1 / 3600)
nu_3 = np.radians(0 * 1 / 3600)
df1 = 1e-2
df2 = 0 * 1e-2
df3 = 0 * 1e-2

phi = np.radians(55.8)
u_0 = 7.292115e-5
uE, uN, uU = 0, np.cos(phi) * u_0, np.sin(phi) * u_0

deltat = 0.004

directory = '/home/andrei/CLionProjects/imitators/cmake-build-release/tt3axis.imu'


def quat_mult(a, b):
    q0 = a[0]*b[0] - a[1]*b[1] - a[2]*b[2] - a[3]*b[3]
    q1 = a[0]*b[1] + b[0]*a[1] + a[2]*b[3] - a[3]*b[2]
    q2 = a[0]*b[2] + b[0]*a[2] + a[3]*b[1] - a[1]*b[3]
    q3 = a[0]*b[3] + b[0]*a[3] + a[1]*b[2] - a[2]*b[1]
    return np.array([q0, q1, q2, q3], dtype=np.float64)


def inverse_quat(q):
    return [q[0], *[-i for i in q[1:]]]


def deltaq(q1, q0):
    q0_inv = inverse_quat(q0)
    return quat_mult(q1, q0_inv)


def get_all_data(path):
    data_1 = np.loadtxt(path, skiprows=1)
    w_1, w_2, w_3 = data_1[:, 0], data_1[:, 1], data_1[:, 2]
    f_1, f_2, f_3 = data_1[:, 3], data_1[:, 4], data_1[:, 5]
    q0, q1, q2, q3 = data_1[:, 6], data_1[:, 7], data_1[:, 8], data_1[:, 9]
    return w_1, w_2, w_3, f_1, f_2, f_3, q0, q1, q2, q3


data_file = get_all_data(directory)
nu = np.array([nu_1, nu_2, nu_3])
df = np.array([df1, df2, df3])
w_init = np.radians(data_file[0:3]).T.tolist()
f_init = (data_file[3:6] + df[:, np.newaxis]).T.tolist()
w = np.radians(data_file[:3]) + nu[:, np.newaxis]
f = data_file[3:6] + df[:, np.newaxis]
q_data = data_file[6:]
w1, w2, w3 = w[0], w[1], w[2]
f1, f2, f3 = f[0], f[1], f[2]
q0_data, q1_data, q2_data, q3_data = q_data[0], q_data[1], q_data[2], q_data[3]

def initial_alignment(f_init, w_init):
    f_z = np.mean(f_init[:30*250], axis=0)
    w_z = np.mean(w_init[:30*250], axis=0)
    L_xz = np.array([np.cross(w_z, f_z) / np.linalg.norm(np.cross(w_z, f_z)),
                     np.cross(f_z, np.cross(w_z, f_z)) / np.linalg.norm(np.cross(f_z, np.cross(w_z, f_z))),
                     f_z / np.linalg.norm(f_z)])

    r = Rotation.from_matrix(L_xz)
    r_ = r.as_quat()[[3, 0, 1, 2]]
    return L_xz, r_


L_xz, r_ = initial_alignment(f_init, w_init)
qmy_0, qmy_1, qmy_2, qmy_3 = r_


def f_data_initial():
    g0_ = np.sqrt(np.mean(f_init[:30*250], axis=0)[0] ** 2 +
                  np.mean(f_init[:30*250], axis=0)[1] ** 2 +
                  np.mean(f_init[:30*250], axis=0)[2] ** 2)
    return g0_

g0 = f_data_initial()

def rotation_imu():
    f = []
    q0, q1, q2, q3 = [], [], [], []
    for m in range(len(data_file[0])):
        f.append(FilterUpdate_My(w1[m], w2[m], w3[m], f1[m], f2[m], f3[m]))
        q0.append(float(f[m][0]))
        q1.append(float(f[m][1]))
        q2.append(float(f[m][2]))
        q3.append(float(f[m][3]))
    return q0, q1, q2, q3


# our filter
def FilterUpdate_My(w_x, w_y, w_z, a_x, a_y, a_z):
    global qmy_0, qmy_1, qmy_2, qmy_3

    a = np.array([a_x, a_y, a_z])
    norm = np.linalg.norm(a)
    a /= norm

    dp0, dp1, dp2, dp3 = quat_mult([qmy_0 / 2, qmy_1 / 2, qmy_2 / 2, qmy_3 / 2], [0, w_x, w_y, w_z])
    p = qmy_0 + dp0 * deltat / 2, qmy_1 + dp1 * deltat / 2, qmy_2 + dp2 * deltat / 2, qmy_3 + dp3 * deltat / 2,

    u = np.cross([a_x, a_y, a_z], quat_mult(quat_mult(inverse_quat(p), [0, 0, 0, g0]), p)[1:])
    norm_u = np.linalg.norm(u)
    if norm_u > 0:
        u_1, u_2, u_3 = (-2 * beta * u[:3]) / norm_u
    else:
        u_0, u_1, u_2, u_3 = 0, 0, 0, 0
    qq_omega_0, qq_omega_1, qq_omega_2, qq_omega_3 = quat_mult([0.5 * qmy_0, 0.5 * qmy_1, 0.5 * qmy_2, 0.5 * qmy_3],
                                                               [0, w_x - u_1, w_y - u_2, w_z - u_3])

    qmy_0 += qq_omega_0 * deltat
    qmy_1 += qq_omega_1 * deltat
    qmy_2 += qq_omega_2 * deltat
    qmy_3 += qq_omega_3 * deltat

    norm = np.linalg.norm([qmy_0, qmy_1, qmy_2, qmy_3])
    qmy_0 /= norm
    qmy_1 /= norm
    qmy_2 /= norm
    qmy_3 /= norm
    return [qmy_0, qmy_1, qmy_2, qmy_3]


if __name__ == '__main__':
    A = [0, 0, 0]
    nu_x = []
    g_new = []
    E_aa_g = []
    qq = []
    t = []
    up, norm_ = [], []
    om_u = []
    q_df, q_f, q_f_cross = [], [], []
    q0_my, q1_my, q2_my, q3_my = rotation_imu()
    T = np.arange(0, len(q0_data), 1) * deltat

    dq = np.array(deltaq([q0_my, q1_my, q2_my, q3_my], [q0_data, q1_data, q2_data, q3_data]))
    dq[:, dq[0] < 0] *= -1

    x_E, x_N, x_U = np.empty(250 * 210), np.empty(250 * 210), np.empty(250 * 210)
    v_E, v_N, v_U = np.empty(250 * 210), np.empty(250 * 210), np.empty(250 * 210)
    a_E, a_N, a_U = np.empty(250 * 210), np.empty(250 * 210), np.empty(250 * 210)
    x_E[0], x_N[0], x_U[0] = 0, 0, 0
    v_E[0], v_N[0], v_U[0] = 0, 0, 0
    a_E[0], a_N[0], a_U[0] = dq[1][0] * 2, dq[2][0] * 2, dq[3][0] * 2


    for i in range(len(T) - 1):
        E_aa_g.append([0,
                       A[0] - np.cross(A, [a_E[i], a_N[i], a_U[i]])[0] + g0 * a_N[i],
                       A[1] - np.cross(A, [a_E[i], a_N[i], a_U[i]])[1] - g0 * a_E[i],
                       A[2] - np.cross(A, [a_E[i], a_N[i], a_U[i]])[2] + g0])
        qq.append(quat_mult(quat_mult(inverse_quat([q0_my[i], q1_my[i], q2_my[i], q3_my[i]]), E_aa_g[i]),
                            [q0_my[i], q1_my[i], q2_my[i], q3_my[i]])[1:] + df)
        g_new.append(quat_mult(quat_mult(inverse_quat([q0_my[i], q1_my[i], q2_my[i], q3_my[i]]), [0, 0, 0, g0]),
                               [q0_my[i], q1_my[i], q2_my[i], q3_my[i]])[1:])
        up.append(np.cross(qq[i], g_new[i]))
        norm_.append(np.linalg.norm(up[i]))
        om_u.append(-2 * beta * up[i] / norm_[i])

        nu_x.append(quat_mult(quat_mult([q0_my[i], q1_my[i], q2_my[i], q3_my[i]],
                                        [0, nu_1 - om_u[i][0], nu_2 - om_u[i][1], nu_3 - om_u[i][2]]),
                              inverse_quat([q0_my[i], q1_my[i], q2_my[i], q3_my[i]])))

        q_df.append(quat_mult(quat_mult([q0_my[i], q1_my[i], q2_my[i], q3_my[i]],
                                        [0, df1, df2, df3]),
                              inverse_quat([q0_my[i], q1_my[i], q2_my[i], q3_my[i]])))
        q_f.append(quat_mult(quat_mult([q0_my[i], q1_my[i], q2_my[i], q3_my[i]],
                                        [0, f1[i], f2[i], f3[i]]),
                              inverse_quat([q0_my[i], q1_my[i], q2_my[i], q3_my[i]])))

        x_E[i + 1] = x_E[i] + v_E[i] * deltat
        x_N[i + 1] = x_N[i] + v_N[i] * deltat
        x_U[i + 1] = x_U[i] + v_U[i] * deltat
        v_E[i + 1] = v_E[i] + (q_f[:][i][1] - q_df[:][i][1]) * deltat
        v_N[i + 1] = v_N[i] + (q_f[:][i][2] - q_df[:][i][1]) * deltat
        v_U[i + 1] = v_U[i] + (q_f[:][i][3] - q_df[:][i][1]) * deltat
        a_E[i + 1] = a_E[i] + (uE + nu_x[:][i][1]) * deltat
        a_N[i + 1] = a_N[i] + (uN + nu_x[:][i][2]) * deltat
        a_U[i + 1] = a_U[i] + (uU + nu_x[:][i][3]) * deltat