# -*- coding: utf-8 -*-
#
# Max-Planck-Gesellschaft zur Förderung der Wissenschaften e.V. (MPG) is
# holder of all proprietary rights on this computer program.
# Copyright 2019 Max-Planck-Gesellschaft zur Förderung
# der Wissenschaften e.V. (MPG). acting on behalf of its Max Planck Institute
# for Intelligent Systems. All rights reserved.
#
# Contact: ps-license@tuebingen.mpg.de

from __future__ import annotations

import numpy as np
import torch
import torch.nn.functional as F


def rot_mat_to_euler(rot_mats: torch.Tensor) -> torch.Tensor:
    """Convert rotation matrices to euler angles (Y-axis rotation only).

    Args:
        rot_mats: Rotation matrices (B, 3, 3).

    Returns:
        Y-axis euler angles (B,).
    """
    sy = torch.sqrt(
        rot_mats[:, 0, 0] * rot_mats[:, 0, 0]
        + rot_mats[:, 1, 0] * rot_mats[:, 1, 0]
    )
    return torch.atan2(-rot_mats[:, 2, 0], sy)


def vertices2landmarks(
    vertices: torch.Tensor,
    faces: torch.Tensor,
    lmk_faces_idx: torch.Tensor,
    lmk_bary_coords: torch.Tensor,
) -> torch.Tensor:
    """Calculate landmarks by barycentric interpolation.

    Args:
        vertices: Mesh vertices (B, V, 3).
        faces: Mesh faces (F, 3).
        lmk_faces_idx: Landmark face indices (B, L).
        lmk_bary_coords: Barycentric coordinates (B, L, 3).

    Returns:
        Landmarks (B, L, 3).
    """
    batch_size, num_verts = vertices.shape[:2]
    device = vertices.device

    lmk_faces = torch.index_select(faces, 0, lmk_faces_idx.view(-1)).view(
        batch_size, -1, 3
    )
    lmk_faces += (
        torch.arange(batch_size, dtype=torch.long, device=device).view(-1, 1, 1)
        * num_verts
    )
    lmk_vertices = vertices.view(-1, 3)[lmk_faces].view(batch_size, -1, 3, 3)
    landmarks = torch.einsum("blfi,blf->bli", [lmk_vertices, lmk_bary_coords])
    return landmarks


def vertices2joints(
    J_regressor: torch.Tensor, vertices: torch.Tensor
) -> torch.Tensor:
    """Calculate 3D joint locations from vertices.

    Args:
        J_regressor: Joint regressor (J, V).
        vertices: Mesh vertices (B, V, 3).

    Returns:
        Joint locations (B, J, 3).
    """
    return torch.einsum("bik,ji->bjk", [vertices, J_regressor])


def blend_shapes(
    betas: torch.Tensor, shape_disps: torch.Tensor
) -> torch.Tensor:
    """Calculate per-vertex displacement from blend shapes.

    Args:
        betas: Blend shape coefficients (B, num_betas).
        shape_disps: Shape displacements (V, 3, num_betas).

    Returns:
        Per-vertex displacement (B, V, 3).
    """
    return torch.einsum("bl,mkl->bmk", [betas, shape_disps])


def batch_rodrigues(
    rot_vecs: torch.Tensor,
    epsilon: float = 1e-8,
    dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
    """Convert batch of axis-angle vectors to rotation matrices.

    Args:
        rot_vecs: Axis-angle vectors (N, 3).
        epsilon: Small value for numerical stability.
        dtype: Output data type.

    Returns:
        Rotation matrices (N, 3, 3).
    """
    batch_size = rot_vecs.shape[0]
    device = rot_vecs.device

    angle = torch.norm(rot_vecs + 1e-8, dim=1, keepdim=True)
    rot_dir = rot_vecs / angle

    cos = torch.unsqueeze(torch.cos(angle), dim=1)
    sin = torch.unsqueeze(torch.sin(angle), dim=1)

    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: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
    """Create batch of 4x4 transformation matrices.

    Args:
        R: Rotation matrices (B, 3, 3).
        t: Translation vectors (B, 3, 1).

    Returns:
        Transformation matrices (B, 4, 4).
    """
    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: torch.Tensor,
    joints: torch.Tensor,
    parents: torch.Tensor,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Apply batch of rigid transformations to joints.

    Args:
        rot_mats: Rotation matrices (B, N, 3, 3).
        joints: Joint locations (B, N, 3).
        parents: Kinematic tree (N,).
        dtype: Data type.

    Returns:
        Tuple of (posed_joints (B, N, 3), rel_transforms (B, N, 4, 4)).
    """
    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)
    ).reshape(-1, joints.shape[1], 4, 4)

    transform_chain = [transforms_mat[:, 0]]
    for i in range(1, parents.shape[0]):
        curr_res = torch.matmul(transform_chain[parents[i]], transforms_mat[:, i])
        transform_chain.append(curr_res)

    transforms = torch.stack(transform_chain, dim=1)
    posed_joints = transforms[:, :, :3, 3]

    joints_homogen = F.pad(joints, [0, 0, 0, 1])
    rel_transforms = transforms - F.pad(
        torch.matmul(transforms, joints_homogen), [3, 0, 0, 0, 0, 0, 0, 0]
    )

    return posed_joints, rel_transforms


def lbs(
    betas: torch.Tensor,
    pose: torch.Tensor,
    v_template: torch.Tensor,
    shapedirs: torch.Tensor,
    posedirs: torch.Tensor,
    J_regressor: torch.Tensor,
    parents: torch.Tensor,
    lbs_weights: torch.Tensor,
    pose2rot: bool = True,
    dtype: torch.dtype = torch.float32,
) -> tuple[torch.Tensor, torch.Tensor]:
    """Perform Linear Blend Skinning.

    Args:
        betas: Shape parameters (B, NB).
        pose: Pose parameters in axis-angle (B, (J+1)*3).
        v_template: Template mesh (B, V, 3).
        shapedirs: Shape displacements (V, 3, NB).
        posedirs: Pose PCA coefficients (P, V*3).
        J_regressor: Joint regressor (J, V).
        parents: Kinematic tree (J,).
        lbs_weights: Skinning weights (V, J+1).
        pose2rot: Whether to convert axis-angle to rotation matrices.
        dtype: Data type.

    Returns:
        Tuple of (vertices (B, V, 3), joints (B, J, 3)).
    """
    batch_size = max(betas.shape[0], pose.shape[0])
    device = betas.device

    v_shaped = v_template + blend_shapes(betas, shapedirs)
    J = vertices2joints(J_regressor, v_shaped)

    ident = torch.eye(3, dtype=dtype, device=device)
    if pose2rot:
        rot_mats = batch_rodrigues(pose.view(-1, 3), dtype=dtype).view(
            [batch_size, -1, 3, 3]
        )
        pose_feature = (rot_mats[:, 1:, :, :] - ident).view([batch_size, -1])
        pose_offsets = torch.matmul(pose_feature, posedirs).view(batch_size, -1, 3)
    else:
        pose_feature = pose[:, 1:].view(batch_size, -1, 3, 3) - ident
        rot_mats = pose.view(batch_size, -1, 3, 3)
        pose_offsets = torch.matmul(pose_feature.view(batch_size, -1), posedirs).view(
            batch_size, -1, 3
        )

    v_posed = pose_offsets + v_shaped

    J_transformed, A = batch_rigid_transform(rot_mats, J, parents, dtype=dtype)

    W = lbs_weights.unsqueeze(dim=0).expand([batch_size, -1, -1])
    num_joints = J_regressor.shape[0]
    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
