import torch, numpy as np

from .utils import euler_angles_to_matrix



def project3DPoints(Ps_, points_3D_):
	NUM_VIEWS = len(Ps_)
	points_3D = points_3D_ if points_3D_.shape[0] == 3 else points_3D_.T
	Ps = np.concatenate((Ps_)).reshape(NUM_VIEWS, 3, 4)
	images_points = Ps[:,:3,:3] @ points_3D + Ps[:,:,-1].reshape(NUM_VIEWS, 3, 1)
	images_points = images_points[:, :2] / images_points[:, 2].reshape(NUM_VIEWS, 1, -1)
	return images_points.transpose(0,2,1)



def projectBatchP(P, poses):
	# P is B x 3 x 4
	# pose is B x PTS x 3
	# out is B x PTS x 2
	poses_2d = torch.matmul(P[..., :3, :3], poses.transpose(-1,-2)) + P[..., :,  [-1]]
	return (poses_2d[...,:2,:] / poses_2d[...,[2],:]).transpose(-1,-2)

def projectBatchK(K, poses):
	# K is ... x 3 x 3
	# pose is ... x PTS x 3
	# out is ... x PTS x 2
	assert K.shape[-1] == K.shape[-2] == 3
	poses_2d = torch.matmul(K[..., :3, :3], poses.transpose(-1,-2))
	return (poses_2d[...,:2,:] / (1e-10 + poses_2d[...,[2],:])).transpose(-1,-2)


def projectP(P, poses, with_depth=False):
	# P is CAM x 3 x 4
	# pose is N x PTS x 3
	# out is N x CAM x PTS x 2
	# while P.ndim < poses.ndim: P = P[None]
	# poses_2d = torch.matmul(P[..., :3, :3], poses.transpose(-1,-2)[:, None]) + P[..., :,  [-1]]

	if poses.ndim == 3:
		if P.ndim == 3:
			poses_2d = torch.matmul(P[None, :, :3, :3], poses.transpose(-1,-2)[:, None]) + P[None, :, :,  [-1]]
		elif P.ndim == 4:
			poses_2d = torch.matmul(P[:, :, :3, :3], poses.transpose(-1,-2)[:, None]) + P[:, :, :,  [-1]]
		else: assert False
	elif poses.ndim == 4:
		poses_2d = torch.matmul(P[None, None, :, :3, :3], poses.transpose(-1,-2)[..., None, :, :]) + P[None, None, :, :,  [-1]]
	else: assert False
	depths = poses_2d[...,[2],:]
	pts2d = (poses_2d[...,:2,:] / (depths + 1e-10)).transpose(-1,-2)
	if with_depth:
		return pts2d, depths.squeeze(-2)
	return pts2d

def project(pose, K, rigid_angles, T):
	# K	: 3 x 3
	# rigid_angles : C x 3
	# T : C x 3 x 1
	# pose : B x N x 3
	# return : B x C x N x 2
	# K [R t] * X = K [RX + t]
	poses_2d_est = K @ (euler_angles_to_matrix(rigid_angles, "ZXY") @ pose.transpose(1,2)[:, None] + T)
	return (poses_2d_est[:,:,:2] / poses_2d_est[:,:,[2]]).transpose(2,3)


def rigidTransform(pts, R, t):
	"""
	pts : B x N x 3
	R : 1 x 3 x 3 or B x 3 x 3
	t: 1 x 1 x 3 or B x 1 x 3
	or

	pts : B x 1 x N x 3
	R : 1 x K x 3 x 3
	t: 1 x K x 1 x 3
	"""
	pts_new = (R @ pts.transpose(-1, -2)).transpose(-1,-2)
	if t is not None:
		pts_new = pts_new + t
	return pts_new

def rigidTransformRT(pts, Rt):
	"""
	pts : ... x N x 3
	Rt : ... x 3 x 4
	"""
	# assert pts.ndim == Rt.ndim
	return (Rt[...,:3] @ pts.transpose(-1, -2) + Rt[...,[-1]]).transpose(-1,-2)


import math
def intrinsic_matrix_from_field_of_view(fov_degrees, imshape, to_torch=False):
	# imshape is (height, width)
	fov_radians = fov_degrees * np.pi / 180.
	larger_side = max(imshape[0], imshape[1])
	focal_length = larger_side / (math.tan(fov_radians / 2.) * 2.)
	K = np.array([[focal_length, 0, imshape[1] / 2.], [0, focal_length, imshape[0] / 2.], [0, 0, 1]])
	if to_torch:
		K = torch.from_numpy(K).float()
	return K

def nearest_rotation_batch(A: torch.Tensor) -> torch.Tensor:
    """
    Enforce each 3x3 matrix in batch A to be a valid rotation via SVD.
    A: (..., 3, 3)
    Returns R with det(R)=+1 and R^T R = I
    """
    U, _, Vt = torch.linalg.svd(A)
    R = U @ Vt
    det = torch.det(R)
    # Flip last column if det is negative
    mask = det < 0
    if mask.any():
        U[mask, :, -1] *= -1
        R = U @ Vt
    return R

def reoptimizeRT(X_world, x_proj, R, T, K, lr=1e-3, steps=100, with_R=False, batch_size=4096, device='cuda', conv='ZYX', min_lr=1e-6, PRINT_FREQ=10):
	"""
	X_world : B x N x 3
	x_proj : B x N x 2 (original projection)
	R: B x 3 x 3 or 1 x 3 x 3 or in euler B x 3 or 1 x 3
	T: B x 3 x 1 or 1 x 3 x 1
	K: B x 3 x 3 or 1 x 3 x 3
	returns:
	R_opt : B x 3 x 3 or 1 x 3 x 3
	T_opt : B x 3 x 1 or 1 x 3 x 1
	"""
	assert X_world.ndim == 3 and x_proj.ndim == 3
	assert X_world.shape[:2] == x_proj.shape[:2]
	assert X_world.shape[2] == 3 and x_proj.shape[2] == 2
	if R.ndim == 2:
		is_input_euler = True
		if len(R) == 1:
			assert R.shape[-1] == 3
		else:
			assert len(R) == len(X_world)
	else:
		is_input_euler = False
		assert R.ndim == 3 and R.shape[1] == 3 and R.shape[2] == 3
	assert T.ndim == 3 and T.shape[1] == 3 and T.shape[2] == 1
	assert K.ndim == 3 and K.shape[1] == 3 and K.shape[2] == 3

	num_chunks = (len(X_world) + batch_size - 1) // batch_size
	X_world = torch.chunk(X_world.detach().cpu() if num_chunks > 1 else X_world, num_chunks)
	x_proj = torch.chunk(x_proj.detach().cpu() if num_chunks > 1 else x_proj, num_chunks)

	T_opt = T.clone().detach().to(device).requires_grad_(True)
	is_single_RT = len(R) == 1
	if is_single_RT:
		assert len(T) == len(K) == 1
	if not is_single_RT:
		T_opt = torch.chunk(T_opt.detach(), num_chunks) # keep on GPU, its impact negiligible
		T_opt = [T_opt[i].detach().requires_grad_(True) for i in range(num_chunks)] # set again to ensure requires_grad is True for each chunk
	
	opt_list = [{'params': T_opt, 'lr': lr, 'name': 'T'}]
	from .utils import matrix_to_euler_angles, euler_angles_to_matrix
	if with_R:
		if is_input_euler:
			R_opt = R.clone().detach().to(device).requires_grad_(True)
		else:
			R_opt = matrix_to_euler_angles(R.clone().detach().to(device), conv).requires_grad_(True)
		if not is_single_RT:
			R_opt = torch.chunk(R_opt.detach(), num_chunks) # keep on GPU, its impact negiligible
			R_opt = [R_opt[i].detach().requires_grad_(True) for i in range(num_chunks)] # set again to ensure requires_grad is True for each chunk
		opt_list.append({'params': R_opt, 'lr': lr, 'name': 'R'})
	else:
		if not is_single_RT:
			R = torch.chunk(R, num_chunks)

	K = K.to(device) # keep on GPU, its impact negiligible
	if not is_single_RT:
		K = torch.chunk(K, num_chunks) # offload to cpu

	optimizer = torch.optim.Adam(opt_list, lr=lr)
	scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, factor=0.8, patience=0, min_lr=min_lr, verbose=False)
	best_loss = 1e10
	R_best, T_best = R, T
	for it in range(steps):
		loss_sum = 0.0
		for ci in range(num_chunks):
			optimizer.zero_grad()
			if with_R:
				R_curr = euler_angles_to_matrix(R_opt if is_single_RT else R_opt[ci], conv)
			else:
				R_curr = R if is_single_RT else R[ci]

			if is_single_RT:
				x_proj_est = projectBatchP(torch.cat([K @ R_curr, K.to(device) @ T_opt], dim=-1), X_world[ci].to(device))
			else:
				x_proj_est = projectBatchP(torch.cat([K @ R_curr, K[ci] @ T_opt[ci]], dim=-1), X_world[ci].to(device))
			loss = (x_proj_est - x_proj[ci].to(device)).abs().mean()
			loss.backward()
			if it % PRINT_FREQ == 0:
				print('Iter %04d, Loss: %.6f' % (it, loss.item()), 'lr %.6f' % optimizer.param_groups[0]['lr'])
			optimizer.step()
			loss_sum += loss.item()

		scheduler.step(loss_sum)
		if loss_sum < best_loss:
			best_loss = loss_sum
			if with_R:
				R_best = (R_opt if is_single_RT else torch.cat(R_opt)).clone().detach().cpu()
			T_best = (T_opt if is_single_RT else torch.cat(T_opt)).clone().detach().cpu()
		if abs(optimizer.param_groups[0]['lr'] - min_lr) < 1e-9:
			break

	R_best_euler = None
	if with_R:
		R_best_euler = R_best.clone().detach().cpu()
		R_best = euler_angles_to_matrix(R_best.to(device), conv).detach().cpu() # convert euler angles back to rotation matrix

	return R_best, T_best, R_best_euler

