import numpy as np, torch
import pytorch3d
import pytorch3d.transforms
from pytorch3d.transforms import matrix_to_axis_angle, axis_angle_to_matrix

from .utils import matrix_to_euler_angles


def distanceL2(T1, T2, weights=None):
	# T1 and T2: ... x DIM
	# output is: ...
	# weights is ...
	# return torch.sqrt(torch.sum((T1 - T2)**2,-1))
	dists = torch.linalg.norm(T1 - T2, dim=-1)
	if weights is not None:
		if weights.shape[-1] == 1:
			dists *= weights.squeeze(-1)
		else:
			dists *= weights
	return dists
def avgL2distance(T1, T2, weights=None):
	if weights is None:
		return distanceL2(T1, T2, weights).mean()
	else:
		weights_sum = weights.sum()
		if weights_sum == 0:
			return torch.tensor(0.0, dtype=T1.dtype, device=T1.device)
		return distanceL2(T1, T2, weights).sum() / weights_sum

def medianL2distance(T1, T2, weights=None):
	errors = distanceL2(T1, T2, weights)
	if weights is None:
		return errors.median()
	else:
		pdb.set_trace()
		werrs = errors / (weights + 1e-10)
		return werrs.median()

def getL2distance(T1, T2, weights=None):
	dist = distanceL2(T1, T2, weights)
	return torch.stack([dist.mean(), dist.median()])


def getEulerToAxisAngles(angles):
	"""
	angles: B x (NUM_JOINTS*3)
	# warning: conversion to and from is never exact!
	"""
	return euler2axis_direct(angles.reshape(-1, 3)).reshape(len(angles), -1)


def euler2axis_direct(euler_angles, eps=1e-6):
	# assert False # TEST!!!
	"""
	backpropagation will fail for euler_angles = [0, 0, 0], which means w == 1, and norm == 0, in this case add small eps to euler_angles!!!
	euler_angles: B x 3
	https://www.euclideanspace.com/maths/geometry/rotations/conversions/eulerToAngle/index.htm
	"""
	# heading, attitude, bank = euler_angles[...,0], euler_angles[...,1], euler_angles[...,2]
	heading, attitude, bank = euler_angles[...,1], euler_angles[...,0], euler_angles[...,2]
	c1 = torch.cos(heading * 0.5)
	s1 = torch.sin(heading * 0.5)
	c2 = torch.cos(attitude * 0.5)
	s2 = torch.sin(attitude * 0.5)
	c3 = torch.cos(bank * 0.5)
	s3 = torch.sin(bank * 0.5)
	c1c2 = c1 * c2
	s1s2 = s1 * s2
	w = c1c2 * c3 - s1s2 * s3
	x = c1c2 * s3 + s1s2 * c3
	y = s1 * c2 * c3 + c1 * s2 * s3
	z = c1 * s2 * c3 - s1 * c2 * s3
	# norm = torch.sqrt(x * x + y * y + z * z)
	norm = torch.sqrt(x * x + y * y + z * z + 1e-8) # add small noise to prevent norm == 0 and division by zero

	min_range, max_range = -1.0 + eps, 1.0 - eps
	change = torch.zeros((len(w)), device=w.device)
	lower = w <= min_range
	upper = w >= max_range
	change[lower] = w[lower].detach() - min_range # < 0, but - (-) = +
	change[upper] = w[upper].detach() - max_range # > 0, but - (+) = -
	w = w - change

	# with torch.no_grad():
	#	print(w.min().item(), w.max().item())
	# 	w.clamp_(-1.0+eps, 1.0-eps)  # Clamp to avoid numerical instability

	angle = 2 * torch.acos(w)
	axis_angles = torch.stack((x, y, z), -1) / (norm[...,None] + 1e-10)
	final_axis_angles = axis_angles * angle[...,None]
	return final_axis_angles

def getAxisToEulerAngles_deprecated(angles, convention='ZYX'):
	"""
	angles: B x (NUM_JOINTS*3)
	"""
	rot_mat = axis_angle_to_matrix(angles.reshape(len(angles), -1, 3))
	euler_angles = matrix_to_euler_angles(rot_mat, convention=convention)
	out_euler = euler_angles.reshape(len(angles), -1)
	########### test ############
	# axis_back_from_rot_mat = matrix_to_axis_angle(rot_mat).reshape(len(angles), -1)
	# assert torch.allclose(angles, axis_back_from_rot_mat, atol=1e-4, rtol=1e-4), 'conversion back failed' # this holds
	# axis_back = getEulerToAxisAngles(out_euler)
	# diff = (angles - axis_back).abs()
	# diff_b = diff.sum(-1)
	# diff_non_zero = torch.where(diff_b > 1e-5)[0]
	# assert torch.allclose(angles, axis_back, rtol=1e-2, atol=1e-2), 'conversion back failed'
	# pdb.set_trace()
	#############################
	return out_euler

def getAxisToEulerAngles(angles):
	"""
	angles: B x (NUM_JOINTS*3)
	# warning: conversion to and from is never exact!
	"""
	euler_angles = axis2euler_direct(angles.reshape(-1, 3)).reshape(len(angles), -1)
	############ test #############
	# axis_back = euler2axis_direct(euler_angles.reshape(-1,3))
	# assert torch.allclose(angles, axis_back.reshape(len(angles), -1), atol=1e-3, rtol=1e-3), 'conversion back failed' # this holds
	# euler_angles_deprecated = getAxisToEulerAngles_deprecated(angles.reshape(-1, 3)).reshape(len(angles), -1)
	# diff_euler = (euler_angles - euler_angles_deprecated).reshape(-1, 3).abs().sum(-1)
	# pdb.set_trace()
	###############################
	return euler_angles


def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
    """
    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
    using Gram--Schmidt orthogonalization per Section B of [1].
    Args:
        d6: 6D rotation representation, of size (*, 6)

    Returns:
        batch of rotation matrices of size (*, 3, 3)

    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
    On the Continuity of Rotation Representations in Neural Networks.
    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
    Retrieved from http://arxiv.org/abs/1812.07035
    """

    a1, a2 = d6[..., :3], d6[..., 3:]
    b1 = torch.nn.functional.normalize(a1, dim=-1)
    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
    b2 = torch.nn.functional.normalize(b2, dim=-1)
    b3 = torch.cross(b1, b2, dim=-1)
    return torch.stack((b1, b2, b3), dim=-2)

def matrix_to_axis(x):
	"""
	x: ... x 3 x 3
	"""
	return pytorch3d.transforms.matrix_to_axis_angle(x)


def axis_to_euler_pytorch3d(axis_angles, convention='ZYX'):
	"""
	axis_angles: B x (NUM_JOINTS*3) or B x N x 3
	"""
	num_dim = axis_angles.ndim
	assert num_dim == 2 or num_dim == 3, f"Expected axis_angles to have 2 or 3 dimensions, but got {num_dim}"
	rot_mat = pytorch3d.transforms.axis_angle_to_matrix(axis_angles.reshape(len(axis_angles), -1, 3))
	euler_angles = pytorch3d.transforms.matrix_to_euler_angles(rot_mat, convention=convention)
	if num_dim == 2:
		euler_angles = euler_angles.reshape(len(axis_angles), -1)
	return euler_angles

def axis_to_6d_pytorch3d(axis_angles):
	"""
	axis_angles: B x (NUM_JOINTS*3) or B x N x 3
	"""
	num_dim = axis_angles.ndim
	assert num_dim == 2 or num_dim == 3, f"Expected axis_angles to have 2 or 3 dimensions, but got {num_dim}"
	rot_mat = pytorch3d.transforms.axis_angle_to_matrix(axis_angles.reshape(len(axis_angles), -1, 3))
	d6 = matrix_to_rotation_6d(rot_mat)
	if num_dim == 2:
		d6 = d6.reshape(len(axis_angles), -1)
	return d6


from .smpl_fx.utils_3d import matrix_to_rotation_6d, rotation_6d_to_matrix
def euler_to_6d_pytorch3d(euler_angles, convention='ZYX'):
	"""
	euler_angles: B x (NUM_JOINTS*3)
	"""
	rot_mat = pytorch3d.transforms.euler_angles_to_matrix(euler_angles.reshape(len(euler_angles), -1, 3), convention=convention)
	d6 = matrix_to_rotation_6d(rot_mat).flatten(1, -1)
	return d6

def six_d_to_euler_pytorch(rot6d, convention='ZYX'):
	"""
	rot6d: B x (NUM_JOINTS*6)
	"""
	rot_mat = rotation_6d_to_matrix(rot6d.reshape(len(rot6d), -1, 6))
	# assert rot_mat.isnan().any() == False
	# assert rot_mat.isinf().any() == False
	eulers = pytorch3d.transforms.matrix_to_euler_angles(rot_mat, convention=convention).flatten(1, -1)
	# assert eulers.isnan().any() == False
	# assert eulers.isinf().any() == False

	return eulers
