'''
pkl2mot.py

Convert joint tracking pickle recorded in pybullet (.pkl) to OpenIm motion (.mot) file

input:
The .pkl file to convert

output:
The name of the .mot file to export

usage:

pkl files stored in data/joint tracking
python pkl2mot.py -i 'benchmark_grooming.mot' -o 'benchmark_grooming.mot'
'''
import argparse
import numpy as np
import pandas as pd
from scipy.interpolate import interp1d

def interpolate_data(data, inc=10, kind='cubic'):
    """ Interpolate data to increase the number of points """
    original_x = data.shape[0]
    original_indices = np.linspace(0, original_x - 1, original_x)
    new_indices = np.linspace(0, original_x - 1, original_x * inc)

    # Perform interpolation
    f = interp1d(original_indices, data, axis=0, kind=kind)
    return f(new_indices)

parser = argparse.ArgumentParser(description='Translate joint tracking pickle to osim motion')
parser.add_argument('-i', '--input', required=True)
parser.add_argument('-o', '--output', required=True)
parser.add_argument('-s', '--start', required=False, default=0)
parser.add_argument('-e', '--end', required=False, default=None)
args = parser.parse_args()

data_path = args.input
filename = args.output

names_equivalence = {
    'Angle_RF_ThC_pitch': 'RFCoxa_pitch',
    'Angle_RF_ThC_yaw': 'RFCoxa_yaw',
    'Angle_RF_ThC_roll': 'RFCoxa_roll',
    'Angle_RF_CTr_pitch': 'RFFemur_pitch',
    'Angle_RF_CTr_roll': 'RFFemur_roll',
    'Angle_RF_FTi_pitch': 'RFTibia_pitch',
    'Angle_RF_TiTa_pitch': 'RFTarsus1_pitch',
    'Angle_LF_ThC_pitch': 'LFCoxa_pitch',
    'Angle_LF_ThC_yaw': 'LFCoxa_yaw',
    'Angle_LF_ThC_roll': 'LFCoxa_roll',
    'Angle_LF_CTr_pitch': 'LFFemur_pitch',
    'Angle_LF_CTr_roll': 'LFFemur_roll',
    'Angle_LF_FTi_pitch': 'LFTibia_pitch',
    'Angle_LF_TiTa_pitch': 'LFTarsus1_pitch'
}
converted_dict = {}

data = pd.read_pickle(data_path)
start = int(args.start)
end = int(args.end) if args.end else -1
# print(data.items())
old_step_time=1e-2
step_time=5e-4

for joints, val in data.items():
    new_name = 'joint_' + names_equivalence[joints]
    if old_step_time != step_time:
        val = interpolate_data(val[start:end], int(old_step_time/step_time))

    converted_dict[new_name] = val

d = converted_dict
angles = []
for k,v in d.items():
    steps = len(v)
    angles.append(v)
angle_count = len(angles)

with open(filename, 'w') as fo:
    fo.write("Coordinates\n")
    fo.write("version=1\n")
    fo.write("nRows={}\n".format(steps))
    fo.write("nColumns={}\n".format(angle_count+1))
    fo.write("inDegrees=yes\n")
    fo.write("endheader\n")
    fo.write("{:>16}\t".format('time'))
    for k in d.keys():
        c = k
        fo.write("{:>16}\t".format(k))
    fo.write("\n")
    for i in range(steps):
        fo.write("{:16.8f}\t".format(step_time * i))
        for j in range(angle_count):
            fo.write("{:16.8f}\t".format(angles[j][i]*180/np.pi))
        fo.write("\n")

