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

from datetime import datetime
start_time = datetime.now()

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

# errors
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
nu = np.array([nu_1, nu_2, nu_3])
df = np.array([df1, df2, df3])

# ang Earth velocity
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'

# quaternion product
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])


# inverse quaternion
def inverse_quat(q):
    q_inv = np.array(q, dtype=np.float64)
    q_inv[1:] *= -1
    return q_inv


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


# data from file
def get_all_data(path):
    data_1 = np.loadtxt(path, skiprows=1)
    w, f, q = np.transpose(data_1[:, :3]), np.transpose(data_1[:, 3:6]), data_1[:, 6:]
    return w[0], w[1], w[2], f[0], f[1], f[2], q[:, 0], q[:, 1], q[:, 2], q[:, 3]


data_file = get_all_data(directory)
w_init = np.radians(data_file[:3]).T
f_init = (data_file[3:6] + df[:, None]).T
w = np.radians(data_file[:3]) + nu[:, None]
f = data_file[3:6] + df[:, None]

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 = data_file[6:]


# initial alignment to find q(0)
def initial_alignment(f_init, w_init, n_samples=30*250):
    f_z = np.mean(f_init[:n_samples], axis=0)
    w_z = np.mean(w_init[:n_samples], 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)

# initial conditions
qmy_0, qmy_1, qmy_2, qmy_3 = r_
vmy_0, vmy_1, vmy_2 = 0, 0, 0
xmy_0, xmy_1, xmy_2 = 0, 0, 0


def g0_find():
    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 = g0_find()


def rotation_imu():
    new = [v_my(w1[m], w2[m], w3[m], f1[m], f2[m], f3[m])[2] for m in range(len(data_file[0]))]
    q = [[float(new[m][i]) for i in range(4)] for m in range(len(data_file[0]))]
    q0, q1, q2, q3 = zip(*q)
    return q0, q1, q2, q3


def x_imu():
    new = [v_my(w1[m], w2[m], w3[m], f1[m], f2[m], f3[m])[0] for m in range(len(data_file[0]))]
    x = [[float(new[m][i]) for i in range(3)] for m in range(len(data_file[0]))]
    x0, x1, x2 = zip(*x)
    return x0, x1, x2


def v_imu():
    new = [v_my(w1[m], w2[m], w3[m], f1[m], f2[m], f3[m])[1] for m in range(len(data_file[0]))]
    v = [[float(new[m][i]) for i in range(3)] for m in range(len(data_file[0]))]
    v0, v1, v2 = zip(*v)
    return v0, v1, v2


def v_my(w_x, w_y, w_z, a_x, a_y, a_z):
    global xmy_0, xmy_1, xmy_2
    global vmy_0, vmy_1, vmy_2
    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
    g_quat = [0, 0, 0, g0]
    a_quat = [0, a_x, a_y, a_z]

    v_new = (quat_mult(quat_mult([qmy_0, qmy_1, qmy_2, qmy_3], a_quat), inverse_quat([qmy_0, qmy_1, qmy_2, qmy_3])) + g_quat)[1:]

    vmy_0 += v_new[0] * deltat
    vmy_1 += v_new[1] * deltat
    vmy_2 += v_new[2] * deltat

    xmy_0 += vmy_0 * deltat
    xmy_1 += vmy_1 * deltat
    xmy_2 += vmy_2 * deltat

    return [xmy_0, xmy_1, xmy_2], [vmy_0, vmy_1, vmy_2], [qmy_0, qmy_1, qmy_2, qmy_3]


if __name__ == '__main__':
    q0_my, q1_my, q2_my, q3_my = rotation_imu()
    v0, v1, v2 = v_imu()
    x0, x1, x2 = x_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
    fig, (ax1, ax2, ax3) = plt.subplots(1, 3, layout='constrained')
    fig.suptitle(fr'$\beta={round(beta * 3600 * (180 / np.pi), 1)},\ \Delta f_1={df1},\ \Delta f_2={df2},\ \Delta f_3={df3},\ \nu_1={nu_1 * 3600 * (180 / np.pi)},\ \nu_2={nu_2 * 3600 * (180 / np.pi)},\ \nu_3={nu_3 * 3600 * (180 / np.pi)}$', fontsize=14)

    ax3.plot(T, dq[1] * (180 / np.pi) * 2, "-", label=r'$\alpha_1$')
    ax3.plot(T, dq[2] * (180 / np.pi) * 2, "-", label=r'$\alpha_2$')
    ax3.plot(T, dq[3] * (180 / np.pi) * 2, "-", label=r'$\alpha_3$')

    ax2.plot(T, v0, "-", label=r'$v_1$')
    ax2.plot(T, v1, "-", label=r'$v_2$')
    ax2.plot(T, v2, "-", label=r'$v_3$')

    ax1.plot(T, x0, "-", label=r'$x_1$')
    ax1.plot(T, x1, "-", label=r'$x_2$')
    ax1.plot(T, x2, "-", label=r'$x_3$')

    plt.show()


end_time = datetime.now()
print('Duration: {}'.format(end_time - start_time))