import argparse
import time
import os
import json

import casadi as cs
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from scipy.spatial.transform import Rotation as R
from scipy.spatial.transform import Slerp
from scipy import interpolate
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import ExpSineSquared

from neural_control.trajectory.q_funcs import (
    q_dot_q, quaternion_inverse, quaternion_to_euler
)
# from utils.visualization import debug_plot, draw_poly
"""
Autor: Elia Kaufmann
Script for generating random and geometric quadrotor trajectories
"""

from neural_control.dynamics.quad_dynamics_base import Dynamics


class Quad(Dynamics):

    def __init__(self, max_thrust_per_motor):
        '''
        :param mass: mass of the quadrotor in [kg]
        :param max_thrust_per_motor: maximum thrust in [N] per motor
        '''
        super().__init__()
        self.max_thrust_per_motor = max_thrust_per_motor

        self.J = self.inertia_vector
        h = self.arm_length / np.sqrt(2.0)
        self.x_f = np.array([h, -h, -h, h])
        self.y_f = np.array([-h, -h, h, h])

        # For z thrust torque calculation
        self.c = 0.013  # m   (z torque generated by each motor)
        self.z_l_tau = np.array([-self.c, self.c, -self.c, self.c])


def check_trajectory(trajectory, inputs, tvec, plot=False):
    """
    @param trajectory:
    @param inputs:
    @param tvec:
    @param plot:
    @return:
    """

    print("Checking trajectory integrity...")

    dt = np.expand_dims(np.gradient(tvec, axis=0), axis=1)
    numeric_derivative = np.gradient(trajectory, axis=0) / dt

    errors = np.zeros((dt.shape[0], 3))

    num_bodyrates = []

    for i in range(dt.shape[0]):
        # 1) check if velocity is consistent with position
        numeric_velocity = numeric_derivative[i, 0:3]
        analytic_velocity = trajectory[i, 7:10]
        errors[i, 0] = np.linalg.norm(numeric_velocity - analytic_velocity)
        if not np.allclose(
            analytic_velocity, numeric_velocity, atol=0.05, rtol=0.05
        ):
            print("inconsistent linear velocity at i = %d" % i)
            print(numeric_velocity)
            print(analytic_velocity)
            return False

        # 2) check if attitude is consistent with acceleration
        gravity = 9.81
        numeric_thrust = numeric_derivative[i, 7:10] + np.array(
            [0.0, 0.0, gravity]
        )
        numeric_thrust = numeric_thrust / np.linalg.norm(numeric_thrust)
        analytic_attitude = trajectory[i, 3:7]
        if np.abs(np.linalg.norm(analytic_attitude) - 1.0) > 1e-6:
            print("quaternion does not have unit norm at i = %d" % i)
            print(analytic_attitude)
            print(np.linalg.norm(analytic_attitude))
            return False

        e_z = np.array([0.0, 0.0, 1.0])
        q_w = 1.0 + np.dot(e_z, numeric_thrust)
        q_xyz = np.cross(e_z, numeric_thrust)
        numeric_attitude = 0.5 * np.array([q_w] + q_xyz.tolist())
        numeric_attitude = numeric_attitude / np.linalg.norm(numeric_attitude)
        # the two attitudes can only differ in yaw --> check x,y component
        q_diff = q_dot_q(
            quaternion_inverse(analytic_attitude), numeric_attitude
        )
        errors[i, 1] = np.linalg.norm(q_diff[1:3])
        if not np.allclose(q_diff[1:3], np.zeros(2, ), atol=0.05, rtol=0.05):
            print("Attitude and acceleration do not match at i = %d" % i)
            print(analytic_attitude)
            print(numeric_attitude)
            print(q_diff)
            return False

        # 3) check if bodyrates agree with attitude difference
        numeric_bodyrates = 2.0 * q_dot_q(
            quaternion_inverse(trajectory[i, 3:7]), numeric_derivative[i, 3:7]
        )[1:]
        num_bodyrates.append(numeric_bodyrates)
        analytic_bodyrates = trajectory[i, 10:13]
        errors[i, 2] = np.linalg.norm(numeric_bodyrates - analytic_bodyrates)
        if not np.allclose(
            numeric_bodyrates, analytic_bodyrates, atol=0.05, rtol=0.05
        ):
            print("inconsistent angular velocity at i = %d" % i)
            print(numeric_bodyrates)
            print(analytic_bodyrates)
            return False

    print("Trajectory check successful")
    print("Maximum linear velocity error: %.5f" % np.max(errors[:, 0]))
    print("Maximum attitude error: %.5f" % np.max(errors[:, 1]))
    print("Maximum angular velocity error: %.5f" % np.max(errors[:, 2]))

    if plot:
        num_bodyrates = np.stack(num_bodyrates)
        plt.figure()
        for i in range(3):
            plt.subplot(3, 2, i * 2 + 1)
            plt.plot(numeric_derivative[:, i], label='numeric')
            plt.plot(trajectory[:, 7 + i], label='analytic')
            plt.ylabel('m/s')
            if i == 0:
                plt.title("Velocity check")
            plt.legend()

        for i in range(3):
            plt.subplot(3, 2, i * 2 + 2)
            plt.plot(num_bodyrates[:, i], label='numeric')
            plt.plot(trajectory[:, 10 + i], label='analytic')
            plt.ylabel('rad/s')
            if i == 0:
                plt.title("Body rate check")
            plt.legend()
        plt.suptitle('Integrity check of reference trajectory')
        plt.show()

    return True


def smooth(x, window_len=11, window='hanning'):
    """smooth the data using a window with requested size.

    This method is based on the convolution of a scaled window with the signal.
    The signal is prepared by introducing reflected copies of the signal
    (with the window size) in both ends so that transient parts are minimized
    in the begining and end part of the output signal.

    input:
        x: the input signal
        window_len: the dimension of the smoothing window; should be an odd int
        window: the type of window from 'flat', 'hanning', 'hamming',
            'bartlett', 'blackman'
            flat window will produce a moving average smoothing.

    output:
        the smoothed signal

    example:

    t=linspace(-2,2,0.1)
    x=sin(t)+randn(len(t))*0.1
    y=smooth(x)

    see also:

    numpy.hanning, numpy.hamming, numpy.bartlett, numpy.blackman,
    numpy.convolve
    scipy.signal.lfilter
    """

    if x.ndim != 1:
        raise ValueError("smooth only accepts 1 dimension arrays.")

    if x.size < window_len:
        raise ValueError("Input vector needs to be bigger than window size.")

    if window_len < 3:
        return x

    if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:
        raise ValueError(
            "Window is on of 'flat', 'hanning', 'hamming', 'bartlett','blackman'"
        )

    # s = np.r_[x[(window_len - 1) // 2:0:-1], x, x[-2:-(window_len - 1) // 2:-1]]

    x_start = np.repeat(x[0], (window_len - 1) // 2)
    x_end = np.repeat(x[-1], (window_len - 1) // 2)
    s = np.concatenate([x_start, x, x_end])

    if window == 'flat':  # moving average
        w = np.ones(window_len, 'd')
    else:
        w = eval('np.' + window + '(window_len)')

    y = np.convolve(w / w.sum(), s, mode='valid')
    return y


def compute_full_traj(quad, t_np, pos_np, vel_np, alin_np):
    len_traj = t_np.shape[0]
    dt = np.mean(np.diff(t_np))

    # Add gravity to accelerations
    gravity = 9.81
    thrust_np = alin_np + np.tile(
        np.array([[0, 0, 1]]), (len_traj, 1)
    ) * gravity
    # Compute body axes
    z_b = thrust_np / np.sqrt(np.sum(thrust_np**2, 1))[:, np.newaxis]
    # new way to compute attitude:
    # https://math.stackexchange.com/questions/2251214/calculate-quaternions-from-two-directional-vectors
    e_z = np.array([[0.0, 0.0, 1.0]])
    q_w = 1.0 + np.sum(e_z * z_b, axis=1)
    q_xyz = np.cross(e_z, z_b)
    att_np = 0.5 * np.concatenate([np.expand_dims(q_w, axis=1), q_xyz], axis=1)
    att_np = att_np / np.sqrt(np.sum(att_np**2, 1))[:, np.newaxis]

    rate_np = np.zeros_like(pos_np)
    f_t = np.zeros((len_traj, 1))

    # Use numerical differentiation of quaternions
    q_dot = np.gradient(att_np, axis=0) / dt
    w_int = np.zeros((len_traj, 3))
    for i in range(len_traj):
        w_int[
            i, :] = 2.0 * q_dot_q(quaternion_inverse(att_np[i, :]),
                                  q_dot[i])[1:]
        f_t[i, 0] = quad.mass * z_b[i].dot(thrust_np[i, :].T)
    rate_np[:, 0] = w_int[:, 0]
    rate_np[:, 1] = w_int[:, 1]
    rate_np[:, 2] = w_int[:, 2]

    minimize_yaw_rate = True
    n_iter_yaw_fix = 20
    if minimize_yaw_rate:
        for iter_yaw_fix in range(n_iter_yaw_fix):
            # print(
            #     "Maximum yawrate before adaption %d / %d: %.6f" %
            #     (iter_yaw_fix, n_iter_yaw_fix, np.max(np.abs(rate_np[:, 2])))
            # )
            q_new = att_np
            yaw_corr_acc = 0.0
            for i in range(1, len_traj):
                yaw_corr = -rate_np[i, 2] * dt
                yaw_corr_acc += yaw_corr
                q_corr = np.array(
                    [
                        np.cos(yaw_corr_acc / 2.0), 0.0, 0.0,
                        np.sin(yaw_corr_acc / 2.0)
                    ]
                )
                q_new[i, :] = q_dot_q(att_np[i, :], q_corr)
                w_int[i, :] = 2.0 * q_dot_q(
                    quaternion_inverse(att_np[i, :]), q_dot[i]
                )[1:]

            q_new_dot = np.gradient(q_new, axis=0) / dt
            for i in range(1, len_traj):
                w_int[i, :] = 2.0 * q_dot_q(
                    quaternion_inverse(q_new[i, :]), q_new_dot[i]
                )[1:]

            att_np = q_new
            rate_np[:, 0] = w_int[:, 0]
            rate_np[:, 1] = w_int[:, 1]
            rate_np[:, 2] = w_int[:, 2]
            # print(
            #     "Maximum yawrate after adaption: %.3f" %
            #     np.max(np.abs(rate_np[:, 2]))
            # )
            if np.max(np.abs(rate_np[:, 2])) < 0.005:
                break

    arot_np = np.gradient(rate_np, axis=0)
    trajectory = np.concatenate(
        [pos_np, att_np, vel_np, rate_np, alin_np, arot_np], axis=1
    )
    motor_inputs = np.zeros((pos_np.shape[0], 4))

    # Compute inputs
    rate_dot = np.gradient(rate_np, axis=0) / dt
    rate_x_Jrate = np.array(
        [
            (quad.J[2] - quad.J[1]) * rate_np[:, 2] * rate_np[:, 1],
            (quad.J[0] - quad.J[2]) * rate_np[:, 0] * rate_np[:, 2],
            (quad.J[1] - quad.J[0]) * rate_np[:, 1] * rate_np[:, 0]
        ]
    ).T

    tau = rate_dot * quad.J[np.newaxis, :] + rate_x_Jrate
    b = np.concatenate((tau, f_t), axis=-1)
    a_mat = np.concatenate(
        (
            quad.y_f[np.newaxis, :], -quad.x_f[np.newaxis, :],
            quad.z_l_tau[np.newaxis, :], np.ones_like(quad.z_l_tau
                                                      )[np.newaxis, :]
        ), 0
    )

    # for i in range(len_traj):
    #     motor_inputs[i, :] = np.linalg.solve(a_mat, b[i, :])

    return trajectory, np.zeros(4), t_np


def compute_random_trajectory(
    quad,
    arena_bound_max,
    arena_bound_min,
    freq_x,
    freq_y,
    freq_z,
    duration=30.0,
    dt=0.01,
    seed=0
):
    # print("Computing random trajectory!")
    assert dt == 0.01

    debug = False

    # kernel to map functions that repeat exactly
    # print("seed is: %d" % seed)
    kernel_y = ExpSineSquared(length_scale=freq_x, periodicity=17) \
               + ExpSineSquared(length_scale=3.0, periodicity=23) \
               + ExpSineSquared(length_scale=4.0, periodicity=51)
    kernel_x = ExpSineSquared(length_scale=freq_y, periodicity=37) \
               + ExpSineSquared(length_scale=3.0, periodicity=61) \
               + ExpSineSquared(length_scale=4.0, periodicity=13)
    kernel_z = ExpSineSquared(length_scale=freq_z, periodicity=19) \
               + ExpSineSquared(length_scale=3.0, periodicity=29) \
               + ExpSineSquared(length_scale=4.0, periodicity=53)

    gp_x = GaussianProcessRegressor(kernel=kernel_x)
    gp_y = GaussianProcessRegressor(kernel=kernel_y)
    gp_z = GaussianProcessRegressor(kernel=kernel_z)

    t_coarse = np.linspace(0.0, duration, int(duration / 0.1), endpoint=False)
    t_vec, dt = np.linspace(
        0.0, duration, int(duration / dt), endpoint=False, retstep=True
    )

    t = cs.MX.sym("t")
    # t_speed is a function starting at zero and ending at zero that
    # modulates time
    # casadi cannot do symbolic integration --> write down the integrand by
    # hand of 2.0*sin^2(t)
    # t_adj = 2.0 * (t / 2.0 - cs.sin(2.0 / duration * cs.pi * t) /
    # (4.0 * cs.pi / duration))
    tau = t / duration
    t_adj = 1.524 * duration * (
        -(
            8 * cs.cos(tau * cs.pi) * cs.constpow(cs.sin(tau * cs.pi), 5) +
            10 * cs.cos(tau * cs.pi) * cs.constpow(cs.sin(tau * cs.pi), 3) +
            39 * cs.sin(tau * cs.pi) * cs.cos(tau * cs.pi) +
            12 * cs.sin(2 * tau * cs.pi) * cs.cos(2 * tau * cs.pi) -
            63 * tau * cs.pi
        ) / (96 * cs.pi)
    )

    f_t_adj = cs.Function('t_adj', [t], [t_adj])
    scaled_time = f_t_adj(t_vec)

    # print("sampling x...")
    x_sample_hr = gp_x.sample_y(t_coarse[:, np.newaxis], 1, random_state=seed)
    # print("sampling y...")
    y_sample_hr = gp_y.sample_y(
        t_coarse[:, np.newaxis], 1, random_state=seed + 1
    )
    # print("sampling z...")
    z_sample_hr = gp_z.sample_y(
        t_coarse[:, np.newaxis], 1, random_state=seed + 2
    )

    pos_np = np.concatenate([x_sample_hr, y_sample_hr, z_sample_hr], axis=1)
    # scale to arena bounds
    max_traj = np.max(pos_np, axis=0)
    min_traj = np.min(pos_np, axis=0)
    pos_centered = pos_np - (max_traj + min_traj) / 2.0
    pos_scaled = pos_centered * (arena_bound_max -
                                 arena_bound_min) / (max_traj - min_traj)
    pos_arena = pos_scaled + (arena_bound_max + arena_bound_min) / 2.0

    if debug:
        plt.plot(pos_arena[:, 0], label="x")
        plt.plot(pos_arena[:, 1], label="y")
        plt.plot(pos_arena[:, 2], label="z")
        plt.legend()
        plt.show()

    # rescale time to get smooth start and end states
    pos_blub_x = interpolate.interp1d(
        t_coarse, pos_arena[:, 0], kind="cubic", fill_value="extrapolate"
    )
    pos_blub_y = interpolate.interp1d(
        t_coarse, pos_arena[:, 1], kind="cubic", fill_value="extrapolate"
    )
    pos_blub_z = interpolate.interp1d(
        t_coarse, pos_arena[:, 2], kind="cubic", fill_value="extrapolate"
    )
    pos_arena = np.concatenate(
        [
            pos_blub_x(scaled_time),
            pos_blub_y(scaled_time),
            pos_blub_z(scaled_time)
        ],
        axis=1
    )

    pos_arena = np.concatenate(
        [
            smooth(np.squeeze(pos_arena[:, 0]), window_len=11)[:, np.newaxis],
            smooth(np.squeeze(pos_arena[:, 1]), window_len=11)[:, np.newaxis],
            smooth(np.squeeze(pos_arena[:, 2]), window_len=11)[:, np.newaxis]
        ],
        axis=1
    )

    # compute numeric derivative & smooth things
    vel_arena = np.gradient(pos_arena, axis=0) / dt
    vel_arena = np.concatenate(
        [
            smooth(np.squeeze(vel_arena[:, 0]), window_len=11)[:, np.newaxis],
            smooth(np.squeeze(vel_arena[:, 1]), window_len=11)[:, np.newaxis],
            smooth(np.squeeze(vel_arena[:, 2]), window_len=11)[:, np.newaxis]
        ],
        axis=1
    )
    acc_arena = np.gradient(vel_arena, axis=0) / dt
    acc_arena = np.concatenate(
        [
            smooth(np.squeeze(acc_arena[:, 0]), window_len=11)[:, np.newaxis],
            smooth(np.squeeze(acc_arena[:, 1]), window_len=11)[:, np.newaxis],
            smooth(np.squeeze(acc_arena[:, 2]), window_len=11)[:, np.newaxis]
        ],
        axis=1
    )
    t_np = t_vec

    trajectory, motor_inputs, t_vec = compute_full_traj(
        quad, t_np, pos_arena, vel_arena, acc_arena
    )

    return trajectory, motor_inputs, t_vec


def compute_geometric_trajectory(quad, duration=30.0, dt=0.001):
    print("Computing geometric trajectory!")
    assert dt == 0.001

    debug = False

    # define position trajectory symbolically
    t = cs.MX.sym("t")
    # t_speed is a function starting at zero and ending at zero that
    # modulates time
    # casadi cannot do symbolic integration --> write down the integrand by
    # hand of 2.0*sin^2(t)
    # t_adj = 2.0 * (t / 2.0 - cs.sin(2.0 / duration * cs.pi * t) /
    # (4.0 * cs.pi / duration))
    tau = t / duration
    t_adj = 1.524 * duration * (
        -(
            8 * cs.cos(tau * cs.pi) * cs.constpow(cs.sin(tau * cs.pi), 5) +
            10 * cs.cos(tau * cs.pi) * cs.constpow(cs.sin(tau * cs.pi), 3) +
            39 * cs.sin(tau * cs.pi) * cs.cos(tau * cs.pi) +
            12 * cs.sin(2 * tau * cs.pi) * cs.cos(2 * tau * cs.pi) -
            63 * tau * cs.pi
        ) / (96 * cs.pi)
    )

    # sphere trajectory rotating around x-axis
    radius_x = 5.0
    radius_y = 3.5
    radius_z = 2.5

    # fast config
    # freq_slow = 0.009
    # freq_fast = 0.33
    # slow config
    freq_slow = 0.02
    freq_fast = 0.12
    pos_x = 3.0 + radius_x * (
        cs.sin(2.0 * cs.pi * freq_fast * t_adj) *
        cs.cos(2.0 * cs.pi * freq_slow * t_adj)
    )
    pos_y = 1.0 + radius_y * (cs.cos(2.0 * cs.pi * freq_fast * t_adj))
    pos_z = 3.5 + radius_z * (
        cs.sin(2.0 * cs.pi * freq_fast * t_adj) *
        cs.sin(2.0 * cs.pi * freq_slow * t_adj)
    )

    pos = cs.vertcat(pos_x, pos_y, pos_z)
    vel = cs.jacobian(pos, t)
    acc = cs.jacobian(vel, t)
    jerk = cs.jacobian(acc, t)
    snap = cs.jacobian(jerk, t)

    t_vec, dt = np.linspace(
        0.0, duration, int(duration / dt), endpoint=False, retstep=True
    )

    f_t_adj = cs.Function('t_adj', [t], [t_adj])
    f_pos = cs.Function('f_pos', [t], [pos])
    f_vel = cs.Function('f_vel', [t], [vel])
    f_acc = cs.Function('f_acc', [t], [acc])
    f_jerk = cs.Function('f_jerk', [t], [jerk])
    f_snap = cs.Function('f_snap', [t], [snap])

    # evaluation seems to only work for scalar inputs --> iterate over time vector
    pos_list = []
    vel_list = []
    alin_list = []
    t_adj_list = []
    for t_curr in t_vec:
        t_adj_list.append(f_t_adj(t_curr).full().squeeze())
        pos_list.append(f_pos(t_curr).full().squeeze())
        vel_list.append(f_vel(t_curr).full().squeeze())
        alin_list.append(f_acc(t_curr).full().squeeze())

    t_adj_np = np.array(t_adj_list)
    pos_np = np.array(pos_list)
    vel_np = np.array(vel_list)
    alin_np = np.array(alin_list)

    if debug:
        plt.plot(t_adj_np)
        plt.show()

    trajectory, motor_inputs, t_vec = compute_full_traj(
        quad, t_vec, pos_np, vel_np, alin_np
    )

    return trajectory, motor_inputs, t_vec


def load_prepare_trajectory(base_dir, dt, speed_factor, test=False):
    """
    speed factor: between 0 and 1, 0.6 would mean that it's going at 0.6 of the
    actual speed (but discrete steps! if dt=0.05 then speed_factor can only be
    0.8, 0.6, 0.4, etc)
    """
    # FOR TESTING: original version
    # quad = Quad(10.0)
    # arena_bound_max = np.array([6.5, 10, 10])  # np.array([8.0, 5.0, 5.0]) #
    # arena_bound_min = np.array([-6.5, -10, 0])
    # trajectory, _, _ = compute_random_trajectory(
    #     quad, arena_bound_max, arena_bound_min, .9, .7, .7,
    #     10, 0.01, seed=np.random.randint(10000)
    # )
    folder = 'test' if test else "train"
    data_list = os.listdir(os.path.join(base_dir, folder))
    rand_traj = np.random.choice(data_list)
    trajectory = np.load(os.path.join(base_dir, folder, rand_traj))

    # dt for trajectory generation is 0.01, then transform back
    take_every_nth = int(dt / 0.01 * speed_factor)
    assert np.isclose(take_every_nth, dt / 0.01 * speed_factor)
    taken_every = trajectory[::take_every_nth, :]

    # transform to euler angels
    quaternions = taken_every[:, 3:7]
    euler_angles = np.array([quaternion_to_euler(q) for q in quaternions])
    # # add in stacking below: euler_angles, taken_every[:, 16:19] (av)
    # stack position, euler angles, velocity, body rates

    # only use pos and vel
    transformed_ref = np.hstack(
        (
            taken_every[:, :3], euler_angles * speed_factor,
            taken_every[:, 7:10] * speed_factor * 2
        )
    )
    # print("transformed shape", transformed_ref.shape)
    return transformed_ref


config = {
    "duration": 10,
    "train_split": .9,
    "freq_x": 0.9,
    "freq_y": 0.7,
    "freq_z": 0.7,
    "out_dir": "data/traj_data_1/",
}


def export(trajectory, t_vec, take_every_nth, output_fn):
    # save trajectory to csv
    # trajectory = [pos_np, att_np, vel_np, rate_np, alin_np, arot_np]
    df_traj = pd.DataFrame()
    df_traj['t'] = t_vec[::take_every_nth]
    df_traj['p_x'] = trajectory[::take_every_nth, 0]
    df_traj['p_y'] = trajectory[::take_every_nth, 1]
    df_traj['p_z'] = trajectory[::take_every_nth, 2]

    df_traj['q_w'] = trajectory[::take_every_nth, 3]
    df_traj['q_x'] = trajectory[::take_every_nth, 4]
    df_traj['q_y'] = trajectory[::take_every_nth, 5]
    df_traj['q_z'] = trajectory[::take_every_nth, 6]

    df_traj['v_x'] = trajectory[::take_every_nth, 7]
    df_traj['v_y'] = trajectory[::take_every_nth, 8]
    df_traj['v_z'] = trajectory[::take_every_nth, 9]

    df_traj['w_x'] = trajectory[::take_every_nth, 10]
    df_traj['w_y'] = trajectory[::take_every_nth, 11]
    df_traj['w_z'] = trajectory[::take_every_nth, 12]

    # df_traj['a_lin_x'] = trajectory[::take_every_nth, 13]
    # df_traj['a_lin_y'] = trajectory[::take_every_nth, 14]
    # df_traj['a_lin_z'] = trajectory[::take_every_nth, 15]

    # df_traj['a_rot_x'] = trajectory[::take_every_nth, 16]
    # df_traj['a_rot_y'] = trajectory[::take_every_nth, 17]
    # df_traj['a_rot_z'] = trajectory[::take_every_nth, 18]

    # df_traj['u_1'] = 0
    # df_traj['u_2'] = 0
    # df_traj['u_3'] = 0
    # df_traj['u_4'] = 0

    # # Sihao needs jerk and snap
    # df_traj['jerk_x'] = 0
    # df_traj['jerk_y'] = 0
    # df_traj['jerk_z'] = 0

    # df_traj['snap_x'] = 0
    # df_traj['snap_y'] = 0
    # df_traj['snap_z'] = 0

    print("Saving trajectory to [%s]." % output_fn)
    df_traj.to_csv(output_fn, index=False)


def upsample(trajectory):
    # to twice the speed
    # np.save("data/orig_traj.npy", trajectory)
    new_trajectory = np.zeros((2 * len(trajectory), 13))
    # np.set_printoptions(suppress=True, precision=3)

    # quaternion
    # norm = np.linalg.norm(trajectory[:, 3:7], axis=1)
    # normed_quats = trajectory[:, 3:7] / np.expand_dims(norm, 1)

    # half_quats = np.sqrt(np.absolute(normed_quats))
    # half_norm = np.linalg.norm(half_quats, axis=1)
    # normed_half_quats = half_quats / np.expand_dims(half_norm, 1)

    # signed_half_quats = np.sign(normed_quats) * normed_half_quats
    from pyquaternion import Quaternion
    quats = [
        Quaternion(trajectory[i, 3:7])**(0.5) for i in range(len(trajectory))
    ]
    quats_as_array = np.array([list(quat) for quat in quats])

    key_rots = R.from_quat(quats_as_array)
    key_times = np.arange(len(key_rots))
    sample_times = np.arange(len(new_trajectory)) * .5
    slerp = Slerp(key_times, key_rots)
    interp_rots = slerp(sample_times[:-1])
    new_quats = interp_rots.as_quat()
    new_trajectory[:-1, 3:7] = new_quats
    new_trajectory[-1, 3:7] = new_quats[-1]

    for i in range(len(trajectory) - 1):
        # pos
        new_trajectory[i * 2, :3] = trajectory[i, :3]
        new_trajectory[
            i * 2 +
            1, :3] = (0.5 * trajectory[i, :3] + 0.5 * trajectory[i + 1, :3])
        # vel
        new_trajectory[i * 2, 7:10] = trajectory[i, 7:10] / 2
        new_trajectory[
            i * 2 + 1,
            7:10] = 0.25 * (trajectory[i, 7:10] + trajectory[i + 1, 7:10])
        # body rates
        new_trajectory[i * 2, 10:] = trajectory[i, 10:13] / 2
        new_trajectory[
            i * 2 + 1,
            10:] = 0.25 * (trajectory[i, 10:13] + trajectory[i + 1, 10:13])
    return new_trajectory


def make_csvs_rl(num=3, out_path="data/csvs_rl"):
    if not os.path.exists(out_path):
        os.makedirs(out_path)
    base_dir = "data/traj_data_1/test"
    use_files = os.listdir(base_dir)[:num]
    seeds = [int(fn[5:-4]) for fn in use_files]
    print("use seeds", seeds)

    quad = Quad(10.0)
    # the arena bounds
    arena_bound_max = np.array([6.5, 10, 10])
    arena_bound_min = np.array([-6.5, -10, 0])
    for seed in seeds:
        trajectory, _, t_vec = compute_random_trajectory(
            quad,
            arena_bound_max,
            arena_bound_min,
            config["freq_x"],
            config["freq_y"],
            config["freq_z"],
            config["duration"],
            0.01,
            seed=seed
        )
        upsampled_traj = upsample(trajectory)
        new_t_vec = np.linspace(0.0, 20, int(20 / 0.01), endpoint=False)
        export(
            upsampled_traj, new_t_vec, 1,
            os.path.join(out_path,
                         str(seed) + ".csv")
        )


def make_dataset():

    cutoff = int(10000 * config["train_split"])
    rand_nums = np.random.permutation(10000)
    train_rand_states = rand_nums[:cutoff]
    test_rand_states = rand_nums[cutoff:]

    quad = Quad(10.0)

    # the arena bounds
    arena_bound_max = np.array([6.5, 10, 10])  # np.array([8.0, 5.0, 5.0]) #
    arena_bound_min = np.array([-6.5, -10, 0])

    for out_dir in ["train", "test"]:
        out_path = os.path.join(config["out_dir"], out_dir)
        if not os.path.exists(out_path):
            os.makedirs(out_path)

    for rand_states, train_test_dir in zip(
        [train_rand_states, test_rand_states], ["train", "test"]
    ):
        out_path = os.path.join(config["out_dir"], train_test_dir)
        for rand in rand_states:
            # compute trajectory
            trajectory, _, _ = compute_random_trajectory(
                quad,
                arena_bound_max,
                arena_bound_min,
                config["freq_x"],
                config["freq_y"],
                config["freq_z"],
                config["duration"],
                0.01,
                seed=rand
            )
            np.save(
                os.path.join(out_path, f"traj_{rand}.npy"), trajectory[:, :10]
            )

    traj_len = len(trajectory)
    config["traj_len"] = traj_len

    with open(os.path.join(config["out_dir"], "config.json"), "w") as outfile:
        json.dump(config, outfile)


if __name__ == '__main__':
    make_csvs_rl(num=20)
    # make_dataset()
    # # Test loading function
    # load_prepare_trajectory("data/traj_data_1/", 0.1, 0.6, test=1)
