import json

import cv2
import h5py
import numpy as np
import torch
import tqdm

from detrsmpl.models.body_models.builder import build_body_model
from detrsmpl.models.body_models.utils import batch_transform_to_camera_frame


class SMCReader:
    def __init__(self, file_path, body_model=None):
        """Read SenseMocapFile endswith ".smc", see: https://github.com/open-
        mmlab/detrsmpl/blob/main/docs/smc.md.

        Args:
            file_path (str):
                Path to an SMC file.
            body_model (nn.Module or dict):
                Only needed for SMPL transformation to device frame
                if nn.Module: a body_model instance
                if dict: a body_model config
        """
        self.smc = h5py.File(file_path, 'r')
        self.__calibration_dict__ = None
        self.action_id = self.smc.attrs['action_id']
        self.actor_id = self.smc.attrs['actor_id']
        self.datetime_str = self.smc.attrs['datetime_str']  # .decode()
        self.kinect_num_frames = self.smc['Kinect'].attrs['num_frame']
        self.num_kinects = self.smc['Kinect'].attrs['num_device']
        self.kinect_color_resolution = self.get_kinect_color_resolution(0)
        self.kinect_depth_resolution = self.get_kinect_depth_resolution(0)
        self.iphone_exists = 'iPhone' in self.smc.keys()
        self.num_iphones = 1
        if self.iphone_exists:
            self.iphone_num_frames = self.smc['iPhone'].attrs['num_frame']
            self.iphone_color_resolution = \
                self.smc['iPhone'].attrs['color_resolution']  # vertical
            self.iphone_depth_resolution = \
                self.smc['iPhone'].attrs['depth_resolution']  # vertical
        self.keypoint_exists = 'Keypoints3D' in self.smc.keys()
        if self.keypoint_exists:
            self.keypoints_num_frames = self.smc['Keypoints3D'].attrs[
                'num_frame']
            self.keypoints_convention = self.smc['Keypoints3D'].attrs[
                'convention']
            self.keypoints_created_time = self.smc['Keypoints3D'].attrs[
                'created_time']
        self.smpl_exists = 'SMPL' in self.smc.keys()
        if self.smpl_exists:
            self.smpl_num_frames = self.smc['SMPL'].attrs['num_frame']
            self.smpl_created_time = self.smc['SMPL'].attrs['created_time']

            # initialize body model
            if isinstance(body_model, torch.nn.Module):
                self.body_model = body_model
            elif isinstance(body_model, dict):
                self.body_model = build_body_model(body_model)
            else:
                # in most cases, SMCReader is instantiated for image reading
                # only. Hence, it is wasteful to initialize a body model until
                # really needed in get_smpl()
                self.body_model = None
                self.default_body_model_config = dict(
                    type='SMPL',
                    gender='neutral',
                    num_betas=10,
                    keypoint_src='smpl_45',
                    keypoint_dst='smpl_45',
                    model_path='data/body_models/smpl',
                    batch_size=1,
                )

    def get_kinect_color_extrinsics(self, kinect_id, homogeneous=True):
        """Get extrinsics(cam2world) of a kinect RGB camera by kinect id.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.
            homogeneous (bool, optional):
                If true, returns rotation and translation in
                one 4x4 matrix. Defaults to True.

        Returns:
            homogeneous is True
                ndarray: A 4x4 matrix of rotation and translation(cam2world).
            homogeneous is False
                dict: A dict of rotation and translation,
                        keys are R and T,
                        each value is an ndarray.
        """
        R = np.asarray(self.calibration_dict[str(kinect_id * 2)]['R']).reshape(
            3, 3)
        T = np.asarray(self.calibration_dict[str(kinect_id *
                                                 2)]['T']).reshape(3)
        if homogeneous:
            extrinsics = np.identity(4, dtype=float)
            extrinsics[:3, :3] = R
            extrinsics[:3, 3] = T
            return extrinsics
        else:
            return {'R': R, 'T': T}

    @property
    def calibration_dict(self):
        """Get the dict of calibration.

        Returns:
            dict:
                A dict of calibrated extrinsics.
        """
        if self.__calibration_dict__ is not None:
            return self.__calibration_dict__
        else:
            return json.loads(self.smc['Extrinsics'][()])

    def get_kinect_depth_extrinsics(self, kinect_id, homogeneous=True):
        """Get extrinsics(cam2world) of a kinect depth camera by kinect id.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.
            homogeneous (bool, optional):
                If true, returns rotation and translation in
                one 4x4 matrix. Defaults to True.

        Returns:
            homogeneous is True
                ndarray: A 4x4 matrix of rotation and translation(cam2world).
            homogeneous is False
                dict: A dict of rotation and translation,
                        keys are R and T,
                        each value is an ndarray.
        """
        R = np.asarray(self.calibration_dict[str(kinect_id * 2 +
                                                 1)]['R']).reshape(3, 3)
        T = np.asarray(self.calibration_dict[str(kinect_id * 2 +
                                                 1)]['T']).reshape(3)
        if homogeneous:
            extrinsics = np.identity(4, dtype=float)
            extrinsics[:3, :3] = R
            extrinsics[:3, 3] = T
            return extrinsics
        else:
            return {'R': R, 'T': T}

    def get_kinect_color_intrinsics(self, kinect_id):
        """Get intrinsics of a kinect RGB camera by kinect id.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.

        Returns:
            ndarray: A 3x3 matrix.
        """
        kinect_dict = self.smc['Kinect'][str(kinect_id)]
        intrinsics = \
            kinect_dict['Calibration']['Color']['Intrinsics'][()]
        cx, cy, fx, fy = intrinsics[:4]
        intrinsics = \
            np.asarray([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
        return intrinsics

    def get_kinect_color_resolution(self, kinect_id):
        """Get resolution of a kinect RGB camera by kinect id.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.

        Returns:
            ndarray:
                An ndarray of (width, height), shape=[2, ].
        """
        kinect_dict = self.smc['Kinect'][str(kinect_id)]
        resolution = \
            kinect_dict['Calibration']['Color']['Resolution'][()]
        return resolution

    def get_kinect_depth_resolution(self, kinect_id):
        """Get resolution of a kinect depth camera by kinect id.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.

        Returns:
            ndarray:
                An ndarray of (width, height), shape=[2, ].
        """
        kinect_dict = self.smc['Kinect'][str(kinect_id)]
        resolution = \
            kinect_dict['Calibration']['Depth']['Resolution'][()]
        return resolution

    def get_kinect_depth_intrinsics(self, kinect_id):
        """Get intrinsics of a kinect depth camera by kinect id.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.

        Returns:
            ndarray: A 3x3 matrix.
        """
        kinect_dict = self.smc['Kinect'][str(kinect_id)]
        intrinsics = \
            kinect_dict['Calibration']['Depth']['Intrinsics'][()]
        cx, cy, fx, fy = intrinsics[:4]
        intrinsics = \
            np.asarray([[fx, 0, cx], [0, fy, cy], [0, 0, 1]])
        return intrinsics

    def get_iphone_intrinsics(self, iphone_id=0, frame_id=0, vertical=True):
        """Get intrinsics of an iPhone RGB camera by iPhone id.

        Args:
            iphone_id (int, optional):
                ID of an iPhone, starts from 0.
                Defaults to 0.
            frame_id (int, optional):
                int: frame id of one selected frame
                Defaults to 0.
            vertical (bool, optional):
                iPhone assumes landscape orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            ndarray: A 3x3 matrix.
        """
        camera_info = self.smc['iPhone'][str(iphone_id)]['CameraInfo'][str(
            frame_id)]
        camera_info = json.loads(camera_info[()])
        intrinsics = np.asarray(camera_info['cameraIntrinsics']).transpose()

        # Intrinsics have to be adjusted to achieve rotation
        #   1. swapping fx, fy
        #   2. cx -> image height - cy; cy -> cx
        if vertical:
            fx, fy = intrinsics[0, 0], intrinsics[1, 1]
            cx, cy = intrinsics[0, 2], intrinsics[1, 2]
            W, H = self.get_iphone_color_resolution(vertical=False)
            intrinsics = np.eye(3)
            intrinsics[0, 0], intrinsics[1, 1] = fy, fx
            intrinsics[0, 2], intrinsics[1, 2] = H - cy, cx

        return intrinsics

    def get_iphone_extrinsics(self,
                              iphone_id=0,
                              homogeneous=True,
                              vertical=True):
        """Get extrinsics(cam2world) of an iPhone RGB camera by iPhone id.

        Args:
            iphone_id (int, optional):
                ID of an iPhone, starts from 0.
                Defaults to 0.
            homogeneous (bool, optional):
                If true, returns rotation and translation in
                one 4x4 matrix. Defaults to True.
            vertical (bool, optional):
                iPhone assumes landscape orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            homogeneous is True
                ndarray: A 4x4 transformation matrix(cam2world).
            homogeneous is False
                dict: A dict of rotation and translation,
                    keys are R and T,
                    each value is an ndarray.
        """
        if iphone_id != 0:
            raise KeyError('Currently only one iPhone.')
        R = np.asarray(self.calibration_dict['iPhone']['R']).reshape(3, 3)
        T = np.asarray(self.calibration_dict['iPhone']['T']).reshape(3)

        # cam2world
        extrinsics = np.identity(4, dtype=float)
        extrinsics[:3, :3] = R
        extrinsics[:3, 3] = T

        # Extrinsics have to be adjusted to achieve rotation
        # A rotation matrix is applied on the extrinsics
        if vertical:
            # 90-degree clockwise rotation around z-axis
            R = np.eye(4)
            R[:2, :2] = np.array([[0, -1], [1, 0]])
            # Note the extrinsics is cam2world
            # world2cam_adjusted = R @ world2cam
            # => cam2world_adjusted = cam2world @ inv(R)
            extrinsics = extrinsics @ np.linalg.inv(R)
            R = extrinsics[:3, :3]
            T = extrinsics[:3, 3]

        if homogeneous:
            return extrinsics
        else:
            return {'R': R, 'T': T}

    def get_iphone_color_resolution(self, iphone_id=0, vertical=True):
        """Get color image resolution of an iPhone RGB camera by iPhone id.

        Args:
            iphone_id (int, optional):
                ID of an iPhone, starts from 0.
                Defaults to 0.
            vertical (bool, optional):
                iPhone assumes landscape orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            ndarray:get_iphone_keypoints2d
                An ndarray of (width, height), shape=[2, ].
        """
        if iphone_id != 0:
            raise KeyError('Currently only one iPhone.')
        if vertical:
            W_horizontal, H_horizontal = self.iphone_color_resolution
            W_vertical, H_vertical = H_horizontal, W_horizontal
            return np.array([W_vertical, H_vertical])
        else:
            return self.iphone_color_resolution

    def get_kinect_color(self, kinect_id, frame_id=None, disable_tqdm=True):
        """Get several frames captured by a kinect RGB camera.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            disable_tqdm (bool, optional):
                Whether to disable the entire progressbar wrapper.
                Defaults to True.

        Returns:
            ndarray:
                An ndarray in shape [frame_number, height, width, channels].
        """
        frames = []
        if frame_id is None:
            frame_list = range(self.get_kinect_num_frames())
        elif isinstance(frame_id, list):
            frame_list = frame_id
        elif isinstance(frame_id, int):
            assert frame_id < self.get_kinect_num_frames(),\
                'Index out of range...'
            frame_list = [frame_id]
        else:
            raise TypeError('frame_id should be int, list or None.')
        for i in tqdm.tqdm(frame_list, disable=disable_tqdm):
            frames.append(
                self.__read_color_from_bytes__(
                    self.smc['Kinect'][str(kinect_id)]['Color'][str(i)][()]))
        return np.stack(frames, axis=0)

    def get_kinect_rgbd(self,
                        kinect_id,
                        frame_id,
                        mode='color2depth',
                        threshold=0):
        if mode == 'color2depth':
            mapped_color = \
                self.__map_color_to_depth__(
                    kinect_id, frame_id, threshold=threshold
                )
            depth = self.get_kinect_depth(kinect_id, frame_id)[0]
            return mapped_color, depth
        else:
            print('Model {} is not supported...'.format(mode))

    def get_kinect_depth(self, kinect_id, frame_id=None, disable_tqdm=True):
        """Get several frames captured by a kinect depth camera.

        Args:
            kinect_id (int):
                ID of a kinect, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            disable_tqdm (bool, optional):
                Whether to disable the entire progressbar wrapper.
                Defaults to True.

        Returns:
            ndarray:
                An ndarray in shape [frame_number, height, width, channels].
        """
        frames = []
        frame_list = []
        if frame_id is None or type(frame_id) == list:
            frame_list = range(self.get_kinect_num_frames())
            if frame_id:
                frame_list = frame_id
        else:
            assert frame_id < self.get_kinect_num_frames(),\
                'Index out of range...'
            frame_list.append(frame_id)
        for i in tqdm.tqdm(frame_list, disable=disable_tqdm):
            frames.append(
                self.smc['Kinect'][str(kinect_id)]['Depth'][str(i)][()])
        return np.stack(frames, axis=0)

    def __read_color_from_bytes__(self, color_array):
        """Decode an RGB image from an encoded byte array."""
        return cv2.cvtColor(cv2.imdecode(color_array, cv2.IMREAD_COLOR),
                            cv2.COLOR_BGR2RGB)

    def get_num_kinect(self):
        """Get the number of Kinect devices.

        Returns:
            int:
                Number of Kinect devices.
        """
        return self.num_kinects

    def get_kinect_num_frames(self):
        """Get the number of frames recorded by one Kinect RGB camera.

        Returns:
            int:
                Number of frames.
        """
        return self.kinect_num_frames

    def get_iphone_num_frames(self):
        """Get the number of frames recorded by one iPhone RGB camera.

        Returns:
            int:
                Number of frames.
        """
        return self.iphone_num_frames

    def get_depth_mask(self, device_id, frame_id):
        return self.smc['Kinect'][str(device_id)]['Mask'][str(frame_id)][()]

    def get_kinect_mask(self, device_id, frame_id):
        kinect_dict = self.smc['Kinect'][str(device_id)]
        return kinect_dict['Mask_k4abt'][str(frame_id)][()]

    def get_num_iphone(self):
        """Get the number of iPhone devices.

        Returns:
            int:
                Number of iPhone devices.
        """
        return self.num_iphones

    def get_iphone_color(self,
                         iphone_id=0,
                         frame_id=None,
                         disable_tqdm=True,
                         vertical=True):
        """Get several frames captured by an iPhone RGB camera.

        Args:
            iphone_id (int):
                ID of an iPhone, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            disable_tqdm (bool, optional):
                Whether to disable the entire progressbar wrapper.
                Defaults to True.
            vertical (bool, optional):
                iPhone assumes horizontal orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            frames:
                An ndarray in shape [frame_number, height, width, channels].
        """
        frames = []
        if frame_id is None:
            frame_list = range(self.get_iphone_num_frames())
        elif isinstance(frame_id, list):
            frame_list = frame_id
        elif isinstance(frame_id, int):
            assert frame_id < self.get_iphone_num_frames(),\
                'Index out of range...'
            frame_list = [frame_id]
        else:
            raise TypeError('frame_id should be int, list or None.')
        for i in tqdm.tqdm(frame_list, disable=disable_tqdm):
            frame = self.__read_color_from_bytes__(
                self.smc['iPhone'][str(iphone_id)]['Color'][str(i)][()])
            if vertical:
                frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
            frames.append(frame)
        return np.stack(frames, axis=0)

    def get_iphone_depth(self,
                         iphone_id=0,
                         frame_id=None,
                         disable_tqdm=True,
                         vertical=True):
        """Get several frames captured by an iPhone RGB camera.

        Args:
            iphone_id (int):
                ID of an iPhone, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            disable_tqdm (bool, optional):
                Whether to disable the entire progressbar wrapper.
                Defaults to True.
            vertical (bool, optional):
                iPhone assumes horizontal orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            frames:
                An ndarray in shape [frame_number, height, width, channels].
        """
        frames = []
        if frame_id is None:
            frame_list = range(self.get_iphone_num_frames())
        elif isinstance(frame_id, list):
            frame_list = frame_id
        elif isinstance(frame_id, int):
            assert frame_id < self.get_iphone_num_frames(),\
                'Index out of range...'
            frame_list = [frame_id]
        else:
            raise TypeError('frame_id should be int, list or None.')
        for i in tqdm.tqdm(frame_list, disable=disable_tqdm):
            frame = self.smc['iPhone'][str(iphone_id)]['Depth'][str(i)][()]
            if vertical:
                frame = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
            frames.append(frame)
        return np.stack(frames, axis=0)

    def get_kinect_transformation_depth_to_color(self, device_id):
        """Get transformation matrix from depth to color from a single kinect.

        Args:
            kinect_id (int, optional):
                ID of a Kinect, starts from 0.

        Returns:
            ndarray: A 4x4 transformation matrix.
        """
        return np.linalg.inv(self.get_kinect_color_extrinsics(
            device_id)) @ self.get_kinect_depth_extrinsics(device_id)

    def get_kinect_transformation_color_to_depth(self, device_id):
        """Get transformation matrix from color to depth from a single kinect.

        Args:
            kinect_id (int, optional):
                ID of a Kinect, starts from 0.

        Returns:
            ndarray: A 4x4 transformation matrix.
        """
        return np.linalg.inv(self.get_kinect_depth_extrinsics(
            device_id)) @ self.get_kinect_color_extrinsics(device_id)

    def __map_color_to_depth__(self, device_id, frame_id, threshold=100):
        color_image = self.get_kinect_color(device_id, frame_id)[0]
        depth_image = self.get_kinect_depth(device_id, frame_id)[0]
        color_intrinsic = self.get_kinect_color_intrinsics(device_id)
        depth_intrinsic = self.get_kinect_depth_intrinsics(device_id)

        mask = self.get_depth_mask(device_id, frame_id)

        Td2c = self.get_kinect_transformation_depth_to_color(device_id)

        colidx = np.arange(depth_image.shape[1])
        rowidx = np.arange(depth_image.shape[0])
        colidx_map, rowidx_map = np.meshgrid(colidx, rowidx)
        col_indices = colidx_map[mask >= threshold]
        row_indices = rowidx_map[mask >= threshold]

        homo_padding = \
            np.ones((col_indices.shape[0], 1), dtype=np.float32)
        homo_indices = \
            np.concatenate(
                (col_indices[..., None], row_indices[..., None], homo_padding),
                axis=1
            )

        depth_intrinsic_inv = np.linalg.inv(depth_intrinsic)
        normalized_points = \
            depth_intrinsic_inv[None, ...] @ homo_indices[..., None]

        z_values = (depth_image / 1000)[mask >= threshold]
        valid_points = \
            normalized_points.squeeze() * z_values[..., None]

        R = Td2c[:3, :3]
        T = Td2c[:3, 3]
        valid_points = \
            R[None, ...] @ valid_points[..., None] + T[None, ..., None]
        valid_uvs = \
            color_intrinsic[None, ...] @\
            valid_points / valid_points[:, 2][..., None]
        valid_uvs = np.int32(valid_uvs.squeeze()[..., :2] + 0.5)
        valid_uvs[:, 0] = np.clip(valid_uvs[:, 0], 0, color_image.shape[1] - 1)
        valid_uvs[:, 1] = np.clip(valid_uvs[:, 1], 0, color_image.shape[0] - 1)
        mapped_color = np.ones((depth_image.shape[0], depth_image.shape[1], 3),
                               dtype=np.uint8) * 255
        mapped_color[mask >= threshold] = \
            color_image[valid_uvs[:, 1], valid_uvs[:, 0]]

        if threshold == 1:
            return valid_uvs
        return mapped_color

    def get_kinect_skeleton_3d(self, device_id, frame_id):
        """Get the 3D skeleton key points from a certain kinect.

        Args:
            device_id (int):
                ID of a kinect, starts from 0.

        Returns:
            list:
                A list with 3D keypoints
        """
        kinect_dict = self.smc['Kinect'][str(device_id)]
        return json.loads(kinect_dict['Skeleton_k4abt'][str(frame_id)][()])

    def get_depth_floor(self, device_id: int) -> dict:
        """Get the floor plane defined by a normal vector and a center point
        from a certain kinect.

        Args:
            device_id (int):
                ID of a kinect, starts from 0.

        Raises:
            KeyError:
                Key 'floor' not in ID of a kinect.

        Returns:
            dict:
                A dict with 'center', 'normal' and 'pnum'.
        """
        device_dict = self.calibration_dict[str(device_id * 2 + 1)]
        if 'floor' in device_dict:
            return device_dict['floor']
        else:
            raise KeyError(f'Kinect {device_id} has no floor data.')

    def get_keypoints2d(self, device, device_id, frame_id=None, vertical=True):
        """Get keypoints2d projected from keypoints3d.

        Args:
            device (str):
                Device name, should be Kinect or iPhone.
            device_id (int):
                ID of a device, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            vertical (bool, optional):
                Only applicable to iPhone as device
                iPhone assumes horizontal orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            Tuple[np.ndarray, np.ndarray]:
                keypoints2d (N, J, 3) and its mask (J, )
        """
        assert device in {
            'Kinect', 'iPhone'
        }, f'Undefined device: {device}, should be "Kinect" or "iPhone"'
        assert device_id >= 0

        kps2d_dict = self.smc['Keypoints2D'][device][str(device_id)]
        keypoints2d = kps2d_dict['keypoints2d'][...]
        keypoints2d_mask = kps2d_dict['keypoints2d_mask'][...]

        if frame_id is None:
            frame_list = range(self.get_keypoints_num_frames())
        elif isinstance(frame_id, list):
            frame_list = frame_id
        elif isinstance(frame_id, int):
            assert frame_id < self.get_keypoints_num_frames(),\
                'Index out of range...'
            frame_list = [frame_id]
        else:
            raise TypeError('frame_id should be int, list or None.')

        keypoints2d = keypoints2d[frame_list, ...]

        if device == 'iPhone' and vertical:
            # rotate keypoints 2D clockwise by 90 degrees
            W, H = self.get_iphone_color_resolution(vertical=False)
            xs, ys, conf = \
                keypoints2d[..., 0], keypoints2d[..., 1], keypoints2d[..., 2]
            xs, ys = H - ys, xs  # horizontal -> vertical
            keypoints2d[..., 0], keypoints2d[..., 1] = xs.copy(), ys.copy()
            keypoints2d[conf == 0.0] = 0.0

        return keypoints2d, keypoints2d_mask

    def get_kinect_keypoints2d(self, device_id, frame_id=None):
        """Get Kinect 2D keypoints.

        Args:
            device_id (int):
                ID of Kinect, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.

        Returns:
            Tuple[np.ndarray, np.ndarray]:
                keypoints2d (N, J, 3) and its mask (J, )
        """
        assert self.num_kinects > device_id >= 0
        return self.get_keypoints2d('Kinect', device_id, frame_id)

    def get_iphone_keypoints2d(self,
                               device_id=0,
                               frame_id=None,
                               vertical=True):
        """Get iPhone 2D keypoints.

        Args:
            device_id (int):
                ID of iPhone, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            vertical (bool, optional):
                iPhone assumes horizontal orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            Tuple[np.ndarray, np.ndarray]:
                keypoints2d (N, J, 3) and its mask (J, )
        """
        assert device_id >= 0
        return self.get_keypoints2d('iPhone',
                                    device_id,
                                    frame_id,
                                    vertical=vertical)

    def get_color(self,
                  device,
                  device_id,
                  frame_id=None,
                  disable_tqdm=True,
                  vertical=True):
        """Get RGB image(s) from Kinect RGB or iPhone RGB camera.

        Args:
            device (str):
                Device name, should be Kinect or iPhone.
            device_id (int):
                Device ID, starts from 0.
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            disable_tqdm (bool, optional):
                Whether to disable the entire progressbar wrapper.
                Defaults to True.
            vertical (bool, optional):
                Only applicable to iPhone as device
                iPhone assumes horizontal orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            img (ndarray):
                An ndarray in shape [frame_number, height, width, channels].
        """

        assert device in {
            'Kinect', 'iPhone'
        }, f'Undefined device: {device}, should be "Kinect" or "iPhone"'

        if device == 'Kinect':
            img = self.get_kinect_color(device_id, frame_id, disable_tqdm)
        else:
            img = self.get_iphone_color(device_id,
                                        frame_id,
                                        disable_tqdm,
                                        vertical=vertical)

        return img

    def get_keypoints_num_frames(self):
        return self.keypoints_num_frames

    def get_keypoints_convention(self):
        return self.keypoints_convention

    def get_keypoints_created_time(self):
        return self.keypoints_created_time

    def get_keypoints3d(self,
                        device=None,
                        device_id=None,
                        frame_id=None,
                        vertical=True):
        """Get keypoints3d (world coordinate) computed by mocap processing
        pipeline.

        Args:
            device (str):
                Device name, should be Kinect or iPhone.
                None: world coordinate
                Defaults to None.
            device_id (int):
                ID of a device, starts from 0.
                None: world coordinate
                Defaults to None
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            vertical (bool, optional):
                Only applicable to iPhone as device
                iPhone assumes horizontal orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            Tuple[np.ndarray, np.ndarray]:
                keypoints3d (N, J, 4) and its mask (J, )
        """
        assert (device is None and device_id is None) or \
            (device is not None and device_id is not None), \
            'device and device_id should be both None or both not None.'
        if device is not None:
            assert device in {
                'Kinect', 'iPhone'
            }, f'Undefined device: {device}, should be "Kinect" or "iPhone"'
        if device_id is not None:
            assert device_id >= 0

        if frame_id is None:
            frame_list = range(self.get_keypoints_num_frames())
        elif isinstance(frame_id, list):
            frame_list = frame_id
        elif isinstance(frame_id, int):
            assert frame_id < self.get_keypoints_num_frames(),\
                'Index out of range...'
            frame_list = [frame_id]
        else:
            raise TypeError('frame_id should be int, list or None.')

        kps3d_dict = self.smc['Keypoints3D']

        # keypoints3d are in world coordinate system
        keypoints3d_world = kps3d_dict['keypoints3d'][...]
        keypoints3d_world = keypoints3d_world[frame_list, ...]
        keypoints3d_mask = kps3d_dict['keypoints3d_mask'][...]

        # return keypoints3d in world coordinate system
        if device is None:
            return keypoints3d_world, keypoints3d_mask

        # return keypoints3d in device coordinate system
        else:
            if device == 'Kinect':
                cam2world = self.get_kinect_color_extrinsics(
                    kinect_id=device_id, homogeneous=True)
            else:
                cam2world = self.get_iphone_extrinsics(iphone_id=device_id,
                                                       vertical=vertical)

            xyz, conf = keypoints3d_world[..., :3], keypoints3d_world[..., [3]]
            xyz_homogeneous = np.ones([*xyz.shape[:-1], 4])
            xyz_homogeneous[..., :3] = xyz
            world2cam = np.linalg.inv(cam2world)
            keypoints3d = np.einsum('ij,kmj->kmi', world2cam, xyz_homogeneous)
            keypoints3d = np.concatenate([keypoints3d[..., :3], conf], axis=-1)

            return keypoints3d, keypoints3d_mask

    def get_smpl_num_frames(self):
        return self.smpl_num_frames

    def get_smpl_created_time(self):
        return self.smpl_created_time

    def get_smpl(self,
                 device=None,
                 device_id=None,
                 frame_id=None,
                 vertical=True):
        """Get SMPL (world coordinate) computed by mocap processing pipeline.

        Args:
            device (str):
                Device name, should be Kinect or iPhone.
                None: world coordinate
                Defaults to None.
            device_id (int):
                ID of a device, starts from 0.
                None: world coordinate
                Defaults to None
            frame_id (int, list or None, optional):
                int: frame id of one selected frame
                list: a list of frame id
                None: all frames will be returned
                Defaults to None.
            vertical (bool, optional):
                Only applicable to iPhone as device
                iPhone assumes horizontal orientation
                if True, convert data to vertical orientation
                Defaults to True.

        Returns:
            dict:
                'global_orient': np.ndarray of shape (N, 3)
                'body_pose': np.ndarray of shape (N, 69)
                'transl': np.ndarray of shape (N, 3)
                'betas': np.ndarray of shape (N, 10)
        """
        smpl_dict = self.smc['SMPL']
        global_orient = smpl_dict['global_orient'][...]
        body_pose = smpl_dict['body_pose'][...]
        transl = smpl_dict['transl'][...]
        betas = smpl_dict['betas'][...]

        if frame_id is None:
            frame_list = range(self.get_smpl_num_frames())
        elif isinstance(frame_id, list):
            frame_list = frame_id
        elif isinstance(frame_id, int):
            assert frame_id < self.get_keypoints_num_frames(),\
                'Index out of range...'
            frame_list = [frame_id]
        else:
            raise TypeError('frame_id should be int, list or None.')

        body_pose = body_pose[frame_list, ...]
        global_orient = global_orient[frame_list, ...]
        transl = transl[frame_list, ...]

        # return SMPL parameters in world coordinate system
        if device is None:
            smpl_dict = dict(global_orient=global_orient,
                             body_pose=body_pose,
                             transl=transl,
                             betas=betas)

            return smpl_dict

        # return SMPL parameters in device coordinate system
        else:

            if self.body_model is None:
                self.body_model = \
                    build_body_model(self.default_body_model_config)
            torch_device = self.body_model.global_orient.device

            assert device in {
                'Kinect', 'iPhone'
            }, f'Undefined device: {device}, should be "Kinect" or "iPhone"'
            assert device_id >= 0

            if device == 'Kinect':
                T_cam2world = self.get_kinect_color_extrinsics(
                    kinect_id=device_id, homogeneous=True)
            else:
                T_cam2world = self.get_iphone_extrinsics(iphone_id=device_id,
                                                         vertical=vertical)

            T_world2cam = np.linalg.inv(T_cam2world)

            output = self.body_model(
                global_orient=torch.tensor(global_orient, device=torch_device),
                body_pose=torch.tensor(body_pose, device=torch_device),
                transl=torch.tensor(transl, device=torch_device),
                betas=torch.tensor(betas, device=torch_device))
            joints = output['joints'].detach().cpu().numpy()
            pelvis = joints[:, 0, :]

            new_global_orient, new_transl = batch_transform_to_camera_frame(
                global_orient=global_orient,
                transl=transl,
                pelvis=pelvis,
                extrinsic=T_world2cam)

            smpl_dict = dict(global_orient=new_global_orient,
                             body_pose=body_pose,
                             transl=new_transl,
                             betas=betas)

            return smpl_dict