text: # -*- coding: utf-8 -*-
  #
  # Copyright (C) 2019 Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG),
  # acting on behalf of its Max Planck Institute for Intelligent Systems and the
  # Max Planck Institute for Biological Cybernetics. All rights reserved.
  #
  # Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is holder of all proprietary rights
  # on this computer program. You can only use this computer program if you have closed a license agreement
  # with MPG or you get the right to use the computer program from someone who is authorized to grant you that right.
  # Any use of the computer program without a valid license is prohibited and liable to prosecution.
  # Contact: ps-license@tuebingen.mpg.de
  #
  #
  # If you use this code in a research publication please consider citing the following:
  #
  # Expressive Body Capture: 3D Hands, Face, and Body from a Single Image <https://arxiv.org/abs/1904.05866>
  #
  # Code Developed by:
  # Vassilis Choutas <https://ps.is.tuebingen.mpg.de/employees/vchoutas>
  # For the original and better implementation please refer to : https://github.com/vchoutas/smplx
  #
  # 2018.01.02
  
  from __future__ import absolute_import
  from __future__ import print_function
  from __future__ import division
  
  import time
  import torch
  import torch.nn.functional as F
  
  # TODO: Create merged c++ and CUDA kernel
  # TODO: Add argument to choose between python impl and c++/CUDA merged op
  def lbs(betas, pose, v_template, clothed_v_template, shapedirs, posedirs, J_regressor, parents,
          lbs_weights, num_joints=23, dtype=torch.float32):
      ''' Performs Linear Blend Skinning with the given shape and pose parameters
  
          Parameters
          ----------
          betas : torch.tensor BxNB
              The tensor of shape parameters
          pose : torch.tensor Bx(J + 1) * 3
              The pose parameters in axis-angle format
          v_template torch.tensor BxVx3
              The template mesh that will be deformed
          shapedirs : torch.tensor 1xNB
              The tensor of PCA shape displacements
          posedirs : torch.tensor Px(V * 3)
              The pose PCA coefficients
          J_regressor : torch.tensor JxV
              The regressor array that is used to calculate the joints from
              the position of the vertices
          parents: torch.tensor J
              The array that describes the kinematic tree for the model
          lbs_weights: torch.tensor N x V x (J + 1)
              The linear blend skinning weights that represent how much the
              rotation matrix of each part affects each vertex
          pose2rot: bool, optional
              Flag on whether to convert the input pose tensor to rotation
              matrices. The default value is True. If False, then the pose tensor
              should already contain rotation matrices and have a size of
              Bx(J + 1)x9
          num_joints : int, optional
              The number of joints of the model. The default value is equal
              to the number of joints of the SMPL body model
          dtype: torch.dtype, optional
  
          Returns
          -------
          verts: torch.tensor BxVx3
              The vertices of the mesh after applying the shape and pose
              displacements.
          joints: torch.tensor BxJx3
              The joints of the model
      '''
  
      batch_size = betas.shape[0]
      device = betas.device
  
      # Add shape contribution
      v_shaped = v_template + blend_shapes(betas, shapedirs)
  
      # Get the joints
      # NxJx3 array
      J = vertices2joints(J_regressor, v_shaped)
      if clothed_v_template is not None:
          v_shaped = clothed_v_template
  
      rot_mats = batch_rodrigues(pose.view(-1, 3)).view([batch_size, -1, 3, 3])
  
      if posedirs is not None:
          # 3. Add pose blend shapes
          # N x J x 3 x 3
          ident = torch.eye(3, dtype=dtype, device=device)
          pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
  
          # (N x P) x (P, V * 3) -> N x V x 3
          pose_offsets = torch.matmul(pose_feature, posedirs).view(batch_size, -1, 3)
          v_posed = pose_offsets + v_shaped
      else:
          v_posed = v_shaped
  
      # 4. Get the global joint location
      J_transformed, A, abs_A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)
  
      # 5. Do skinning:
      # W is N x V x (J + 1)
      W = lbs_weights.unsqueeze(dim=0).repeat([batch_size, 1, 1])
      num_joints = J_regressor.shape[0]
      # (N x V x (J + 1)) x (N x (J + 1) x 16)
      T = torch.matmul(W, A.view(batch_size, num_joints, 16)).view(batch_size, -1, 4, 4)
  
      homogen_coord = torch.ones([batch_size, v_posed.shape[1], 1], dtype=dtype, device=device)
      v_posed_homo = torch.cat([v_posed, homogen_coord], dim=2)
      v_homo = torch.matmul(T, torch.unsqueeze(v_posed_homo, dim=-1))
  
      verts = v_homo[:, :, :3, 0]
  
      return verts, J_transformed, J, A, abs_A, v_posed
  
  
  def vertices2joints(J_regressor, vertices):
      ''' Calculates the 3D joint locations from the vertices
  
      Parameters
      ----------
      J_regressor : torch.tensor JxV
          The regressor array that is used to calculate the joints from the
          position of the vertices
      vertices : torch.tensor BxVx3
          The tensor of mesh vertices
  
      Returns
      -------
      torch.tensor BxJx3
          The location of the joints
      '''
  
      return torch.einsum('bik,ji->bjk', [vertices, J_regressor])
  
  
  def blend_shapes(betas, shape_disps):
      ''' Calculates the per vertex displacement due to the blend shapes
  
  
      Parameters
      ----------
      betas : torch.tensor Bx(num_betas)
          Blend shape coefficients
      shape_disps: torch.tensor Vx3x(num_betas)
          Blend shapes
  
      Returns
      -------
      torch.tensor BxVx3
          The per-vertex displacement due to shape deformation
      '''
  
      # Displacement[b, m, k] = sum_{l} betas[b, l] * shape_disps[m, k, l]
      # i.e. Multiply each shape displacement by its corresponding beta and
      # then sum them.
      blend_shape = torch.einsum('bl,mkl->bmk', [betas, shape_disps])
      return blend_shape
  
  
  def batch_rodrigues(aa_rots):
      '''
      convert batch of rotations in axis-angle representation to matrix representation
      :param aa_rots: Nx3
      :return: mat_rots: Nx3x3
      '''
  
      dtype = aa_rots.dtype
      device = aa_rots.device
  
      batch_size = aa_rots.shape[0]
  
      angle = torch.norm(aa_rots + 1e-8, dim=1, keepdim=True)
      rot_dir = aa_rots / angle
  
      cos = torch.unsqueeze(torch.cos(angle), dim=1)
      sin = torch.unsqueeze(torch.sin(angle), dim=1)
  
      # Bx1 arrays
      rx, ry, rz = torch.split(rot_dir, 1, dim=1)
  
      zeros = torch.zeros((batch_size, 1), dtype=dtype, device=device)
      K = torch.cat([zeros, -rz, ry, rz, zeros, -rx, -ry, rx, zeros], dim=1) \
          .view((batch_size, 3, 3))
  
      ident = torch.eye(3, dtype=dtype, device=device).unsqueeze(dim=0)
      rot_mat = ident + sin * K + (1 - cos) * torch.bmm(K, K)
      return rot_mat
  
  
  def transform_mat(R, t):
      ''' Creates a batch of transformation matrices
          Args:
              - R: Bx3x3 array of a batch of rotation matrices
              - t: Bx3x1 array of a batch of translation vectors
          Returns:
              - T: Bx4x4 Transformation matrix
      '''
      # No padding left or right, only add an extra row
      return torch.cat([F.pad(R, [0, 0, 0, 1]), F.pad(t, [0, 0, 0, 1], value=1)], dim=2)
  
  
  def batch_rigid_transform(rot_mats, joints, parents, dtype=torch.float32):
      """
      Applies a batch of rigid transformations to the joints
  
      Parameters
      ----------
      rot_mats : torch.tensor BxNx3x3
          Tensor of rotation matrices
      joints : torch.tensor BxNx3
          Locations of joints
      parents : torch.tensor BxN
          The kinematic tree of each object
      dtype : torch.dtype, optional:
          The data type of the created tensors, the default is torch.float32
  
      Returns
      -------
      posed_joints : torch.tensor BxNx3
          The locations of the joints after applying the pose rotations
      rel_transforms : torch.tensor BxNx4x4
          The relative (with respect to the root joint) rigid transformations
          for all the joints
      """
  
      batch_size = rot_mats.shape[0]
      num_joints = joints.shape[1]
      device = rot_mats.device
  
      joints = torch.unsqueeze(joints, dim=-1)
  
      rel_joints = joints.clone()
      rel_joints[:, 1:] -= joints[:, parents[1:]]
  
      transforms_mat = transform_mat(
          rot_mats.view(-1, 3, 3),
          rel_joints.reshape(-1, 3, 1)).view(-1, joints.shape[1], 4, 4)
  
      transform_chain = [transforms_mat[:, 0]]
      for i in range(1, parents.shape[0]):
          # Subtract the joint location at the rest pose
          # No need for rotation, since it's identity when at rest
          curr_res = torch.matmul(transform_chain[parents[i]],
                                  transforms_mat[:, i])
          transform_chain.append(curr_res)
  
      transforms = torch.stack(transform_chain, dim=1)
  
      # The last column of the transformations contains the posed joints
      posed_joints = transforms[:, :, :3, 3]
  
      joints_homogen = torch.cat([joints, torch.zeros([batch_size, num_joints, 1, 1], dtype=dtype, device=device)],dim=2)
      init_bone = torch.matmul(transforms, joints_homogen)
      init_bone = F.pad(init_bone, [3, 0, 0, 0, 0, 0, 0, 0])
      rel_transforms = transforms - init_bone
  
      return posed_joints, rel_transforms, transforms
  