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

gyroMeasError = np.radians(3 / 3600)
beta = 0 * 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 = 0 * 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] + a[1]*b[0] + a[2]*b[3] - a[3]*b[2]
    q2 = a[0]*b[2] + a[2]*b[0] + a[3]*b[1] - a[1]*b[3]
    q3 = a[0]*b[3] + a[3]*b[0] + a[1]*b[2] - a[2]*b[1]
    return [q0, q1, q2, q3]


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


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

    norm = np.sqrt(a_x ** 2 + a_y ** 2 + a_z ** 2)
    a_x /= norm
    a_y /= norm
    a_z /= 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__':
    dq = []
    A = [0, 0, 0]
    nu_x = []
    g_new = []
    E_aa_g = []
    qq = []
    up = []
    norm_ = []
    om_u = []
    q0_my, q1_my, q2_my, q3_my = rotation_imu()
    T = np.arange(0, len(q0_data[:210 * 250]), 1) * deltat

    for i in range(len(T)):
        dq.append(deltaq([q0_my[i], q1_my[i], q2_my[i], q3_my[i]],
                         [q0_data[i], q1_data[i], q2_data[i], q3_data[i]]))

    a_E, a_N, a_U = np.empty(250 * 210), np.empty(250 * 210), np.empty(250 * 210)

    for i in range(len(T)):
        if dq[i][0] < 0:
            dq[i][1] = -dq[i][1]
            dq[i][2] = -dq[i][2]
            dq[i][3] = -dq[i][3]

    a_E[0], a_N[0], a_U[0] = dq[0][1] * 2, dq[0][2] * 2, dq[0][3] * 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]])))

        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


# if __name__ == '__main__':
#     q_my = rotation_imu()
#     T = np.arange(0, len(q0_data), 1) * deltat
#     A = np.array([np.zeros(len(T)), np.zeros(len(T)), np.zeros(len(T))])
#
#     a_E, a_N, a_U = np.empty(len(T)), np.empty(len(T)), np.empty(len(T))
#     dq = np.array(deltaq(q_my, [q0_data, q1_data, q2_data, q3_data]))
#     dq[:, dq[0] < 0] *= -1
#     a_E[0], a_N[0], a_U[0] = dq[1][0] * 2, dq[2][0] * 2, dq[3][0] * 2
#
#     g_new = quat_mult(quat_mult(inverse_quat(np.array(q_my)),
#                                 [np.zeros(len(T)), np.zeros(len(T)), np.zeros(len(T)), np.full(len(T), g0)]), q_my)[1:]  # ?
#     g_a = np.array([g0 * a_N, -g0 * a_E, np.full(len(T), g0)])  # ?
#
#     Aa = np.cross(A.T, np.array([a_E, a_N, a_U]).T).T  # ?
#     E_aa_g = np.insert((A + g_a - Aa), 0, np.zeros(len(T)), axis=0)  # ?
#     qq = np.array(quat_mult(quat_mult(inverse_quat(np.array(q_my)), E_aa_g), q_my))[1:] + np.array([np.full(len(T), df1), np.full(len(T), df2), np.full(len(T), df3)])  # ?
#     up = np.cross(qq.T, np.array(g_new).T).T  # ?
#     norm_ = np.linalg.norm(up)  # ?
#     om_u = -2 * beta * up / norm_  # ?
#     nu_omu = np.array([np.full(len(T), nu_1), np.full(len(T), nu_2), np.full(len(T), nu_3)]) - np.array([om_u[0], om_u[1], om_u[2]])  # ?
#     nu_omu_0 = np.insert(nu_omu, 0, np.zeros(len(T)), axis=0)
#
#     for i in range(len(T)-1):
#         a_E[i + 1] = a_E[i] + (uE + nu_omu_0[1][i]) * deltat
#         a_N[i + 1] = a_N[i] + (uN + nu_omu_0[2][i]) * deltat
#         a_U[i + 1] = a_U[i] + (uU + nu_omu_0[3][i]) * deltat