# partly adapted from https://github.com/microsoft/human-pose-estimation.pytorch

import torch
import numpy as np
import cv2
import kornia

def transform_anchor_points(A, *argv):
    """
    perform matrix multiplication A*anchor_point for each body part and anchor_point
    Args:
        A: torch.tensor (batch, num_parts, 3, 3) of transformation matrices
        *args: tensors with shape (batch, num_parts, num_anchors, 3)
    """

    num_parts = 0
    for arg in argv:
        num_parts += arg.shape[1]
    assert num_parts == A.shape[1], "number of matrices should match number of parts!"

    index = 0
    transformed = []
    for arg in argv:
        num_parts = arg.shape[1]
        num_anchors = arg.shape[2]
        # repeat matrix num_anchors times
        A_ = A[:, index:index+num_parts].unsqueeze(2).repeat(1, 1, num_anchors, 1, 1)
        tr = torch.matmul(A_, arg)
        transformed.append(tr)
        index += num_parts

    return transformed


def get_3rd_point(a, b):
    direct = a - b
    return b + np.array([-direct[1], direct[0]], dtype=np.float32)


def get_dir(src_point, rot_rad):
    sn, cs = np.sin(rot_rad), np.cos(rot_rad)

    src_result = [0, 0]
    src_result[0] = src_point[0] * cs - src_point[1] * sn
    src_result[1] = src_point[0] * sn + src_point[1] * cs

    return src_result


def get_affine_transform(center,
                         scale,
                         rot,
                         output_size,
                         shift=np.array([0, 0], dtype=np.float32),
                         inv=0):
    if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
        print(scale)
        scale = np.array([scale, scale])

    scale_tmp = scale * 200.0
    src_w = scale_tmp[0]
    dst_w = output_size[0]
    dst_h = output_size[1]

    rot_rad = np.pi * rot / 180
    src_dir = get_dir([0, src_w * -0.5], rot_rad)
    dst_dir = np.array([0, dst_w * -0.5], np.float32)

    src = np.zeros((3, 2), dtype=np.float32)
    dst = np.zeros((3, 2), dtype=np.float32)
    src[0, :] = center + scale_tmp * shift
    src[1, :] = center + src_dir + scale_tmp * shift
    dst[0, :] = [dst_w * 0.5, dst_h * 0.5]
    dst[1, :] = np.array([dst_w * 0.5, dst_h * 0.5]) + dst_dir

    src[2:, :] = get_3rd_point(src[0, :], src[1, :])
    dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])

    if inv:
        trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
    else:
        trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))

    return trans


def transform_template(input, params):
    size = input.shape[2]
    # scale up translation
    params[..., -1] = params[..., -1] * size
    return kornia.geometry.warp_affine(input, params, dsize=(size, size))

