# Copyright (c) OpenMMLab. All rights reserved.
import cv2
import mmcv
import numpy as np
import torch

from mmpose.core.post_processing import (affine_transform, fliplr_joints,
                                         get_affine_transform)
from mmpose.datasets.builder import PIPELINES


def _flip_smpl_pose(pose):
    """Flip SMPL pose parameters horizontally.

    Args:
        pose (np.ndarray([72])): SMPL pose parameters

    Returns:
        pose_flipped
    """

    flippedParts = [
        0, 1, 2, 6, 7, 8, 3, 4, 5, 9, 10, 11, 15, 16, 17, 12, 13, 14, 18, 19,
        20, 24, 25, 26, 21, 22, 23, 27, 28, 29, 33, 34, 35, 30, 31, 32, 36, 37,
        38, 42, 43, 44, 39, 40, 41, 45, 46, 47, 51, 52, 53, 48, 49, 50, 57, 58,
        59, 54, 55, 56, 63, 64, 65, 60, 61, 62, 69, 70, 71, 66, 67, 68
    ]
    pose_flipped = pose[flippedParts]
    # Negate the second and the third dimension of the axis-angle
    pose_flipped[1::3] = -pose_flipped[1::3]
    pose_flipped[2::3] = -pose_flipped[2::3]
    return pose_flipped


def _flip_iuv(iuv, uv_type='BF'):
    """Flip IUV image horizontally.

    Note:
        IUV image height: H
        IUV image width: W

    Args:
        iuv np.ndarray([H, W, 3]): IUV image
        uv_type (str): The type of the UV map.
            Candidate values:
                'DP': The UV map used in DensePose project.
                'SMPL': The default UV map of SMPL model.
                'BF': The UV map used in DecoMR project.
            Default: 'BF'

    Returns:
        iuv_flipped np.ndarray([H, W, 3]): Flipped IUV image
    """
    assert uv_type in ['DP', 'SMPL', 'BF']
    if uv_type == 'BF':
        iuv_flipped = iuv[:, ::-1, :]
        iuv_flipped[:, :, 1] = 255 - iuv_flipped[:, :, 1]
    else:
        # The flip of other UV map is complex, not finished yet.
        raise NotImplementedError(
            f'The flip of {uv_type} UV map is not implemented yet.')

    return iuv_flipped


def _construct_rotation_matrix(rot, size=3):
    """Construct the in-plane rotation matrix.

    Args:
        rot (float): Rotation angle (degree).
        size (int): The size of the rotation matrix.
            Candidate Values: 2, 3. Defaults to 3.

    Returns:
        rot_mat (np.ndarray([size, size]): Rotation matrix.
    """
    rot_mat = np.eye(size, dtype=np.float32)
    if rot != 0:
        rot_rad = np.deg2rad(rot)
        sn, cs = np.sin(rot_rad), np.cos(rot_rad)
        rot_mat[0, :2] = [cs, -sn]
        rot_mat[1, :2] = [sn, cs]

    return rot_mat


def _rotate_joints_3d(joints_3d, rot):
    """Rotate the 3D joints in the local coordinates.

    Note:
        Joints number: K

    Args:
        joints_3d (np.ndarray([K, 3])): Coordinates of keypoints.
        rot (float): Rotation angle (degree).

    Returns:
        joints_3d_rotated
    """
    # in-plane rotation
    # 3D joints are rotated counterclockwise,
    # so the rot angle is inversed.
    rot_mat = _construct_rotation_matrix(-rot, 3)

    joints_3d_rotated = np.einsum('ij,kj->ki', rot_mat, joints_3d)
    joints_3d_rotated = joints_3d_rotated.astype('float32')
    return joints_3d_rotated


def _rotate_smpl_pose(pose, rot):
    """Rotate SMPL pose parameters. SMPL (https://smpl.is.tue.mpg.de/) is a 3D
    human model.

    Args:
        pose (np.ndarray([72])): SMPL pose parameters
        rot (float): Rotation angle (degree).

    Returns:
        pose_rotated
    """
    pose_rotated = pose.copy()
    if rot != 0:
        rot_mat = _construct_rotation_matrix(-rot)
        orient = pose[:3]
        # find the rotation of the body in camera frame
        per_rdg, _ = cv2.Rodrigues(orient)
        # apply the global rotation to the global orientation
        res_rot, _ = cv2.Rodrigues(np.dot(rot_mat, per_rdg))
        pose_rotated[:3] = (res_rot.T)[0]

    return pose_rotated


def _flip_joints_3d(joints_3d, joints_3d_visible, flip_pairs):
    """Flip human joints in 3D space horizontally.

    Note:
        num_keypoints: K

    Args:
        joints_3d (np.ndarray([K, 3])): Coordinates of keypoints.
        joints_3d_visible (np.ndarray([K, 1])): Visibility of keypoints.
        flip_pairs (list[tuple()]): Pairs of keypoints which are mirrored
            (for example, left ear -- right ear).

    Returns:
        joints_3d_flipped, joints_3d_visible_flipped
    """

    assert len(joints_3d) == len(joints_3d_visible)

    joints_3d_flipped = joints_3d.copy()
    joints_3d_visible_flipped = joints_3d_visible.copy()

    # Swap left-right parts
    for left, right in flip_pairs:
        joints_3d_flipped[left, :] = joints_3d[right, :]
        joints_3d_flipped[right, :] = joints_3d[left, :]

        joints_3d_visible_flipped[left, :] = joints_3d_visible[right, :]
        joints_3d_visible_flipped[right, :] = joints_3d_visible[left, :]

    # Flip horizontally
    joints_3d_flipped[:, 0] = -joints_3d_flipped[:, 0]
    joints_3d_flipped = joints_3d_flipped * joints_3d_visible_flipped

    return joints_3d_flipped, joints_3d_visible_flipped


@PIPELINES.register_module()
class LoadIUVFromFile:
    """Loading IUV image from file."""

    def __init__(self, to_float32=False):
        self.to_float32 = to_float32
        self.color_type = 'color'
        # channel relations: iuv->bgr
        self.channel_order = 'bgr'

    def __call__(self, results):
        """Loading image from file."""
        has_iuv = results['has_iuv']
        use_iuv = results['ann_info']['use_IUV']
        if has_iuv and use_iuv:
            iuv_file = results['iuv_file']
            iuv = mmcv.imread(iuv_file, self.color_type, self.channel_order)
            if iuv is None:
                raise ValueError(f'Fail to read {iuv_file}')
        else:
            has_iuv = 0
            iuv = None

        results['has_iuv'] = has_iuv
        results['iuv'] = iuv
        return results


@PIPELINES.register_module()
class IUVToTensor:
    """Transform IUV image to part index mask and uv coordinates image. The 3
    channels of IUV image means: part index, u coordinates, v coordinates.

    Required key: 'iuv', 'ann_info'.
    Modifies key: 'part_index', 'uv_coordinates'.

    Args:
        results (dict): contain all information about training.
    """

    def __call__(self, results):
        iuv = results['iuv']
        if iuv is None:
            H, W = results['ann_info']['iuv_size']
            part_index = torch.zeros([1, H, W], dtype=torch.long)
            uv_coordinates = torch.zeros([2, H, W], dtype=torch.float32)
        else:
            part_index = torch.LongTensor(iuv[:, :, 0])[None, :, :]
            uv_coordinates = torch.FloatTensor(iuv[:, :, 1:]) / 255
            uv_coordinates = uv_coordinates.permute(2, 0, 1)
        results['part_index'] = part_index
        results['uv_coordinates'] = uv_coordinates
        return results


@PIPELINES.register_module()
class MeshRandomChannelNoise:
    """Data augmentation with random channel noise.

    Required keys: 'img'
    Modifies key: 'img'

    Args:
        noise_factor (float): Multiply each channel with
         a factor between``[1-scale_factor, 1+scale_factor]``
    """

    def __init__(self, noise_factor=0.4):
        self.noise_factor = noise_factor

    def __call__(self, results):
        """Perform data augmentation with random channel noise."""
        img = results['img']

        # Each channel is multiplied with a number
        # in the area [1-self.noise_factor, 1+self.noise_factor]
        pn = np.random.uniform(1 - self.noise_factor, 1 + self.noise_factor,
                               (1, 3))
        img = cv2.multiply(img, pn)

        results['img'] = img
        return results


@PIPELINES.register_module()
class MeshRandomFlip:
    """Data augmentation with random image flip.

    Required keys: 'img', 'joints_2d','joints_2d_visible', 'joints_3d',
    'joints_3d_visible', 'center', 'pose', 'iuv' and 'ann_info'.
    Modifies key: 'img', 'joints_2d','joints_2d_visible', 'joints_3d',
    'joints_3d_visible', 'center', 'pose', 'iuv'.

    Args:
        flip_prob (float): Probability of flip.
    """

    def __init__(self, flip_prob=0.5):
        self.flip_prob = flip_prob

    def __call__(self, results):
        """Perform data augmentation with random image flip."""
        if np.random.rand() > self.flip_prob:
            return results

        img = results['img']
        joints_2d = results['joints_2d']
        joints_2d_visible = results['joints_2d_visible']
        joints_3d = results['joints_3d']
        joints_3d_visible = results['joints_3d_visible']
        pose = results['pose']
        center = results['center']

        img = img[:, ::-1, :]
        pose = _flip_smpl_pose(pose)

        joints_2d, joints_2d_visible = fliplr_joints(
            joints_2d, joints_2d_visible, img.shape[1],
            results['ann_info']['flip_pairs'])

        joints_3d, joints_3d_visible = _flip_joints_3d(
            joints_3d, joints_3d_visible, results['ann_info']['flip_pairs'])
        center[0] = img.shape[1] - center[0] - 1

        if 'iuv' in results.keys():
            iuv = results['iuv']
            if iuv is not None:
                iuv = _flip_iuv(iuv, results['ann_info']['uv_type'])
            results['iuv'] = iuv

        results['img'] = img
        results['joints_2d'] = joints_2d
        results['joints_2d_visible'] = joints_2d_visible
        results['joints_3d'] = joints_3d
        results['joints_3d_visible'] = joints_3d_visible
        results['pose'] = pose
        results['center'] = center
        return results


@PIPELINES.register_module()
class MeshGetRandomScaleRotation:
    """Data augmentation with random scaling & rotating.

    Required key: 'scale'. Modifies key: 'scale' and 'rotation'.

    Args:
        rot_factor (int): Rotating to ``[-2*rot_factor, 2*rot_factor]``.
        scale_factor (float): Scaling to ``[1-scale_factor, 1+scale_factor]``.
        rot_prob (float): Probability of random rotation.
    """

    def __init__(self, rot_factor=30, scale_factor=0.25, rot_prob=0.6):
        self.rot_factor = rot_factor
        self.scale_factor = scale_factor
        self.rot_prob = rot_prob

    def __call__(self, results):
        """Perform data augmentation with random scaling & rotating."""
        s = results['scale']

        sf = self.scale_factor
        rf = self.rot_factor

        s_factor = np.clip(np.random.randn() * sf + 1, 1 - sf, 1 + sf)
        s = s * s_factor

        r_factor = np.clip(np.random.randn() * rf, -rf * 2, rf * 2)
        r = r_factor if np.random.rand() <= self.rot_prob else 0

        results['scale'] = s
        results['rotation'] = r

        return results


@PIPELINES.register_module()
class MeshAffine:
    """Affine transform the image to get input image. Affine transform the 2D
    keypoints, 3D kepoints and IUV image too.

    Required keys: 'img', 'joints_2d','joints_2d_visible', 'joints_3d',
    'joints_3d_visible', 'pose', 'iuv', 'ann_info','scale',  'rotation' and
    'center'. Modifies key: 'img', 'joints_2d','joints_2d_visible',
    'joints_3d',  'pose', 'iuv'.
    """

    def __call__(self, results):
        image_size = results['ann_info']['image_size']

        img = results['img']
        joints_2d = results['joints_2d']
        joints_2d_visible = results['joints_2d_visible']
        joints_3d = results['joints_3d']
        pose = results['pose']

        c = results['center']
        s = results['scale']
        r = results['rotation']
        trans = get_affine_transform(c, s, r, image_size)

        img = cv2.warpAffine(
            img,
            trans, (int(image_size[0]), int(image_size[1])),
            flags=cv2.INTER_LINEAR)

        for i in range(results['ann_info']['num_joints']):
            if joints_2d_visible[i, 0] > 0.0:
                joints_2d[i] = affine_transform(joints_2d[i], trans)

        joints_3d = _rotate_joints_3d(joints_3d, r)
        pose = _rotate_smpl_pose(pose, r)

        results['img'] = img
        results['joints_2d'] = joints_2d
        results['joints_2d_visible'] = joints_2d_visible
        results['joints_3d'] = joints_3d
        results['pose'] = pose

        if 'iuv' in results.keys():
            iuv = results['iuv']
            if iuv is not None:
                iuv_size = results['ann_info']['iuv_size']
                iuv = cv2.warpAffine(
                    iuv,
                    trans, (int(iuv_size[0]), int(iuv_size[1])),
                    flags=cv2.INTER_NEAREST)
            results['iuv'] = iuv

        return results