import argparse
import json
import transforms3d as tfs

def prepare_data(_data:list[float], save_path: str = "camera_pose_data.json", format: str = "quat_xyzw", mm2m: bool = False):
    new_list = []
    if format == "quat_wxyz":
        for i in range(0, len(_data), 7):
            if not mm2m:
                new = [_data[i], _data[i + 1], _data[i + 2], _data[i + 3], _data[i + 4], _data[i + 5], _data[i + 6]]
            else:
                new = [_data[i]/1e3, _data[i + 1]/1e3, _data[i + 2]/1e3, _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5], _data[i + 6]]
            new_list.extend(new)
    elif format == "quat_xyzw":
        for i in range(0, len(_data), 7):
            if not mm2m:
                new = [_data[i], _data[i + 1], _data[i + 2], _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5]]
            else:
                new = [_data[i] / 1e3, _data[i + 1] / 1e3, _data[i + 2] / 1e3, _data[i + 6], _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5]]
            new_list.extend(new)

    elif format == "euler":
        for i in range(0, len(_data), 6):
            rw, rx, ry, rz = tfs.euler.euler2quat(_data[i + 3], _data[i + 4], _data[i + 5], axes='sxyz')
            if not mm2m:
                new = [_data[i], _data[i + 1], _data[i + 2], rw, rx, ry, rz]
            else:
                new = [_data[i]/1e3, _data[i + 1]/1e3, _data[i + 2]/1e3, rw, rx, ry, rz]
            new_list.extend(new)

    elif format == "rotvertor":
        for i in range(0, len(_data), 6):
            # tfs.quaternions.rotvec2quat
            if not mm2m:
                new = [_data[i], _data[i + 1], _data[i + 2], _data[i + 3], _data[i + 4], _data[i + 5]]
            else:
                new = [_data[i]/1e3, _data[i + 1]/1e3, _data[i + 2]/1e3, _data[i + 6], _data[i + 3], _data[i + 4], _data[i + 5]]
            new_list.extend(new)

    else: pass
    with open(save_path, 'w') as f: json.dump(new_list, f, indent=4)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--data_path", type=str, required=True)
    parser.add_argument("--save_path", type=str, required=True)
    parser.add_argument("--format", type=str, default="quat_xyzw")
    parser.add_argument("--mm2m", type=bool, default=False)
    args = parser.parse_args()

    with open(args.data_path, "r") as f:
        _data = json.load(f)

    prepare_data(_data, args.save_path, args.format, args.mm2m)
