import torch
import numpy as np
import matplotlib.pyplot as plt
import pdb

from pykeops.torch import LazyTensor

class PrintColors:
	# https://stackoverflow.com/questions/287871/how-do-i-print-colored-text-to-the-terminal
	HEADER = '\033[95m'
	OKBLUE = '\033[94m'
	OKCYAN = '\033[96m'
	OKGREEN = '\033[92m'
	WARNING = '\033[93m'
	FAIL = '\033[91m'
	ENDC = '\033[0m'
	BOLD = '\033[1m'
	UNDERLINE = '\033[4m'

def getTemporalIndicesBasedKeys(sorted_keys, diff=2):
	"""
	for diff = 2:
		sorted_keys_: N, sorted, e.g., 0, 2, 3, 10, 12, 13, 14, 20, 21, 22, 23, 24, 30, 33, 34
		output: sorted temporal indices K x 2, [[1, 2], [4, 5], [5, 6], [7, 8], [8, 9], [9, 10], [10, 11], [13, 14]]
	"""
	assert sorted_keys.ndim == 1
	assert isinstance(sorted_keys, np.ndarray)
	j = 0
	temp_matrix = []
	diff_m1 = diff - 1
	while j < len(sorted_keys)-diff_m1:
		if (sorted_keys[j+diff_m1] - sorted_keys[j])+1 == diff:
			temp_matrix.append(np.arange(j, j+diff_m1+1))
		j += 1
	temp_matrix = np.stack(temp_matrix)
	assert temp_matrix.min() >= 0 and temp_matrix.max() <= len(sorted_keys)-1
	assert ((sorted_keys[temp_matrix][:,-1] - sorted_keys[temp_matrix][:,0]) == diff_m1).all()
	for kk in range(diff_m1): assert ((sorted_keys[temp_matrix][:,kk+1] - sorted_keys[temp_matrix][:,kk]) == 1).all()
	return temp_matrix

def createBatchTemporalIndices(sorted_keys, NUM_ELEMENTS, batch_size, diff=2):
	batched_temp_indices = []
	for bs in range(0, NUM_ELEMENTS, batch_size):
		if sorted_keys is None:
			print(PrintColors.WARNING + 'keys are not provided! Assuming the difference between keys is 1' + PrintColors.ENDC)
			max_ = min(batch_size, NUM_ELEMENTS-bs)
			pdb.set_trace()
			if diff == 2:
				batched_temp_indices.append(np.stack([np.arange(0, max_-1), np.arange(1, max_)]))
			else:
				assert False
		else:
			batched_temp_indices.append(getTemporalIndicesBasedKeys(sorted_keys[bs:bs+batch_size], diff))
	return batched_temp_indices

def createChunkedTemporalIndices(sorted_keys, NUM_ELEMENTS, num_chunks, diff=2):
	"""
	Splits the data into 'num_chunks' segments (similar to torch.chunk) 
	and computes temporal indices for each chunk.
	"""
	batched_temp_indices = []

	assert num_chunks > 0
	chunk_size = (NUM_ELEMENTS + num_chunks - 1) // num_chunks

	# Iterate through the number of chunks
	for i in range(num_chunks):
		# Calculate start and end indices for the current chunk
		start_idx = i * chunk_size
		end_idx = min(start_idx + chunk_size, NUM_ELEMENTS)
		
		# Stop if we have run out of elements (e.g., requested more chunks than elements)
		if start_idx >= NUM_ELEMENTS:
			break

		current_chunk_len = end_idx - start_idx

		if sorted_keys is None:
			# Replicated logic from original: assume diff between keys is 1 if keys are missing
			print(PrintColors.WARNING + 'keys are not provided! Assuming the difference between keys is 1' + PrintColors.ENDC)
			pdb.set_trace() # Optional: keep if you need debugging
			
			if diff == 2:
				# Create indices relative to this chunk
				# [0, 1], [1, 2], ... up to current_chunk_len
				batched_temp_indices.append(
					np.stack([np.arange(0, current_chunk_len - 1), 
								np.arange(1, current_chunk_len)])
				)
			else:
				assert False, "Logic for diff != 2 with None keys is not implemented."
		else:
			batched_temp_indices.append(getTemporalIndicesBasedKeys(sorted_keys[start_idx:end_idx], diff))
			
	return batched_temp_indices

def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
	# https://pytorch3d.readthedocs.io/en/latest/_modules/pytorch3d/transforms/rotation_conversions.html#euler_angles_to_matrix
	"""
	Return the rotation matrices for one of the rotations about an axis
	of which Euler angles describe, for each value of the angle given.

	Args:
		axis: Axis label "X" or "Y or "Z".
		angle: any shape tensor of Euler angles in radians

	Returns:
		Rotation matrices as tensor of shape (..., 3, 3).
	"""
	cos = torch.cos(angle)
	sin = torch.sin(angle)
	one = torch.ones_like(angle)
	zero = torch.zeros_like(angle)

	if   axis == "X": R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
	elif axis == "Y": R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
	elif axis == "Z": R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
	else: raise ValueError("letter must be either X, Y or Z.")
	return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))

def euler_angles_to_matrix(euler_angles: torch.Tensor, convention: str) -> torch.Tensor:
	# assert len(convention) == euler_angles.numel()
	# euler_angles must be tensor of size: [..., NUM_ANGLES] or scalar
	if euler_angles.numel() == 1 or (len(euler_angles.shape) == 1 and len(euler_angles) > 3):
		NUM_ANGLES = 1
	else: NUM_ANGLES = euler_angles.shape[-1]
	assert NUM_ANGLES >= 1 and NUM_ANGLES <= 3 and len(convention) == NUM_ANGLES
	if NUM_ANGLES == 1: 
		if euler_angles.ndim > 1: euler_angles = euler_angles[...,0] # flatten last dimension
		return _axis_angle_rotation(convention[0], euler_angles)
	matrices = [
		_axis_angle_rotation(c, e)
		for c, e in zip(convention, torch.unbind(euler_angles, -1))
	]
	if len(matrices) == 2: return torch.matmul(matrices[0], matrices[1])
	else: return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])


def _angle_from_tan(axis: str, other_axis: str, data, horizontal: bool, tait_bryan: bool) -> torch.Tensor:
	"""
	Extract the first or third Euler angle from the two members of
	the matrix which are positive constant times its sine and cosine.

	Args:
		axis: Axis label "X" or "Y or "Z" for the angle we are finding.
		other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
			convention.
		data: Rotation matrices as tensor of shape (..., 3, 3).
		horizontal: Whether we are looking for the angle for the third axis,
			which means the relevant entries are in the same row of the
			rotation matrix. If not, they are in the same column.
		tait_bryan: Whether the first and third axes in the convention differ.

	Returns:
		Euler Angles in radians for each matrix in data as a tensor
		of shape (...).
	"""

	i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
	if horizontal: i2, i1 = i1, i2
	even = (axis + other_axis) in ["XY", "YZ", "ZX"]
	if horizontal == even: return torch.atan2(data[..., i1], data[..., i2])
	if tait_bryan: return torch.atan2(-data[..., i2], data[..., i1])
	return torch.atan2(data[..., i2], -data[..., i1])

def _index_from_letter(letter: str) -> int:
	if letter == "X": return 0
	if letter == "Y": return 1
	if letter == "Z": return 2
	raise ValueError("letter must be either X, Y or Z.")

def matrix_to_euler_angles(matrix: torch.Tensor, convention: str) -> torch.Tensor:

	"""
	Convert rotations given as rotation matrices to Euler angles in radians.

	Args:
		matrix: Rotation matrices as tensor of shape (..., 3, 3).
		convention: Convention string of three uppercase letters.

	Returns:
		Euler angles in radians as tensor of shape (..., 3).
	"""
	if len(convention) != 3:
		raise ValueError("Convention must have 3 letters.")
	if convention[1] in (convention[0], convention[2]):
		# raise ValueError(f"Invalid convention {convention}.")
		raise ValueError("Invalid convention "+convention)
	for letter in convention:
		if letter not in ("X", "Y", "Z"):
			raise ValueError("Invalid letter "+letter+" in convention string.")
	if matrix.size(-1) != 3 or matrix.size(-2) != 3:
		raise ValueError("Invalid rotation matrix shape "+str(matrix.shape))
	i0 = _index_from_letter(convention[0])
	i2 = _index_from_letter(convention[2])
	tait_bryan = i0 != i2
	if tait_bryan:
		# central_angle = torch.asin(matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0))
		central_angle = torch.asin((matrix[..., i0, i2] * (-1.0 if i0 - i2 in [-1, 2] else 1.0)).clamp(-1.0, 1.0))
	else:
		# central_angle = torch.acos(matrix[..., i0, i0])
		central_angle = torch.acos(matrix[..., i0, i0].clamp(-1.0, 1.0))
	o = (
		_angle_from_tan(convention[0], convention[1], matrix[..., i2], False, tait_bryan),
		central_angle,
		_angle_from_tan(convention[2], convention[1], matrix[..., i0, :], True, tait_bryan),
	)
	return torch.stack(o, -1)


def convertMMPoseScores2WeightsOpenPose133(weights, face_idxs, hand_idxs, body_idxs, low_q_face=0.05, high_q_face=0.95, low_q_body=0.05, high_q_body=0.95, low_q_hand=0.05, high_q_hand=0.95, verbose=False, normalize_min_body=None, normalize_min_hand=None):
	assert weights.ndim == 3 or weights.ndim == 4, 'weights should have shape (B, T, 133) or (B, T, 133, 1)'
	weights[:,:,face_idxs] = convertMMPoseScores2Weights(weights[:,:,face_idxs], low_q=low_q_face, high_q=high_q_face, verbose=verbose)
	weights[:,:,hand_idxs] = convertMMPoseScores2Weights(weights[:,:,hand_idxs], low_q=low_q_hand, high_q=high_q_hand, verbose=verbose, normalize_min=normalize_min_hand)
	weights[:,:,body_idxs] = convertMMPoseScores2Weights(weights[:,:,body_idxs], low_q=low_q_body, high_q=high_q_body, verbose=verbose, normalize_min=normalize_min_body)
	return weights

def convertMMPoseScores2Weights(weights, normalize_min=None, normalize_max=None, low_q=0.05, high_q=0.95, gamma=1.0, verbose=False):
	original_min = weights.min().item()
	original_max = weights.max().item()
	show_plot = False
	# show_plot = True
	if show_plot:
		_, axs = plt.subplots(1, 2, figsize=(12,5))
		axs[0].set_title('before convertMMPoseScores2Weights')
		histogram = np.histogram(weights.cpu().numpy().flatten(), bins=100)
		axs[0].bar(histogram[1][:-1], histogram[0], width=histogram[1][1]-histogram[1][0])

	num_weights = weights.numel()
	if normalize_min is None:
		normalize_min = torch.sort(weights.flatten()).values[int(num_weights * low_q)].item()
	if normalize_max is None:
		normalize_max = torch.sort(weights.flatten()).values[int(num_weights * high_q)].item()
	assert normalize_max > normalize_min, 'normalize_max must be greater than normalize_min'
	weights = (weights - normalize_min) / (normalize_max - normalize_min)
	weights = weights.clamp(0, 1)
	weights = weights.pow(gamma)
	num_weights_0 = (weights < 1e-6).sum().item()
	num_weights_1 = (weights > 0.9999).sum().item()
	median_weight = weights.median().item()
	if show_plot:
		histogram = np.histogram(weights.cpu().numpy().flatten(), bins=100)
		axs[1].bar(histogram[1][:-1], histogram[0], width=histogram[1][1]-histogram[1][0])
		axs[1].set_title('after convertMMPoseScores2Weights')
		plt.show()

	if verbose:
		print('orig min %.4f' % original_min, 'orig max %.4f' % original_max,  'median_weight %.4f' % median_weight, 'ratio_0 %.4f' % (num_weights_0/weights.numel()), 'ratio_1 %.4f' % (num_weights_1/weights.numel()), 'used min %.4f' % normalize_min, 'used max %.4f' % normalize_max)

	# weights_zero = (weights < 1e-7).sum(-2)
	# for i in range(weights.shape[1]):
	# 	if weights_zero[i].sum().item() == weights.shape[-1]:
	# 		weights[i] = 1.0 # if all weights are zero, set them to 1
	return weights


def calculateInitialLogits(x, constraints, eps=1e-7):
	"""
	x : B x D
	constraints: D x 2
	x = sigmoid(y) * (constraints[:,1] - constraints[:,0]) + constraints[:,0]
	given x and constraints, calculate y
	"""
	assert x.ndim == 2 and constraints.ndim == 2
	assert constraints.shape[0] == x.shape[1] and constraints.shape[1] == 2

	min_v = constraints[:, 0].view(1, -1)
	diff_v = (constraints[:, 1] - constraints[:, 0]).view(1, -1)

	# Normalize x to [0, 1] based on constraints
	x_normalized = (x - min_v) / (diff_v + eps)
	x_normalized = torch.clamp(x_normalized, eps, 1.0 - eps)

	# Calculate initial logits using inverse sigmoid
	initial_logits = torch.log(x_normalized / (1.0 - x_normalized))

	### test ###
	# y = torch.sigmoid(initial_logits) * diff_v + min_v
	# assert torch.allclose(y, x, atol=1e-5), f"calculated initial logits are incorrect"
	# print('test is ok!!!')
	############

	return initial_logits

def fastKNN_batched(x, y, k, exclude_self=False):
    """
    x: B x N x D
    y: B x M x D
    output: B x N x k (indices), B x N x k (distances)
    """
    B, N, D = x.shape
    _, M, _ = y.shape

    # Create LazyTensors: shape (B, N, 1, D) and (B, 1, M, D)
    x_i = LazyTensor(x[:, :, None, :])  # (B, N, 1, D)
    y_j = LazyTensor(y[:, None, :, :])  # (B, 1, M, D)

    # Squared Euclidean distance: (B, N, M)
    dist_ij = ((x_i - y_j) ** 2).sum(-1)

    # Find k nearest neighbors
    dists_all, knn = dist_ij.Kmin_argKmin(k, dim=2)  # shape: (B, N, k)

    if exclude_self:
        dists_all = dists_all[:, :, 1:]
        knn = knn[:, :, 1:]

    dists_all = dists_all.sqrt()  # convert from squared distance

    return knn, dists_all