import torch
import numpy as np

from tqdm import tqdm

def get_edges_from_faces(faces):
    """
    Extract unique edges from faces.
    """
    edges = torch.cat([
        faces[:, [0, 1]],
        faces[:, [1, 2]],
        faces[:, [2, 0]]
    ], dim=0)
    edges = torch.sort(edges, dim=1).values  # Sort to ensure uniqueness
    edges = torch.unique(edges, dim=0)
    return edges

def runBatchedSMPLX(params, smplx_model, BS=500, with_vertices=False, device='cuda'):
    params = {k: v.to(device) for k, v in params.items()}
    out = {'joints': []}
    if with_vertices:
        out['vertices'] = []
    num_data = 1
    for key in params.keys():
        num_data = max(num_data, params[key].shape[0])
    if BS > num_data:
        out_ = smplx_model(**params)
        out['joints'] = out_.joints.detach().cpu()
        if with_vertices:
            out['vertices'] = out_.vertices.detach().cpu()
    else:
        for i in tqdm(range(0, num_data, BS)):
            batch_params = {k: v if v.shape[0] == 1 else v[i:i+BS] for k, v in params.items()}
            if 'J_regressor' in params:
                batch_params['J_regressor'] = params['J_regressor']
            if 'lbs_weights' in params:
                batch_params['lbs_weights'] = params['lbs_weights']
            out_ = smplx_model(**batch_params)
            if with_vertices:
                out['vertices'] += [out_['vertices'].detach().cpu()]
            out['joints'] += [out_['joints'].detach().cpu()]

        for key in out.keys():
            out[key] = torch.cat(out[key], 0)
    return out

def getModelOut(model, params, BS=512, with_vertices=True):
    with torch.no_grad():
        out = runBatchedSMPLX(params, model, BS=BS, with_vertices=with_vertices)
    return out

def getSMPLXOut(smplx_kwargs, device, params, BS=512, run_with_verts=True):
    from .smpl_fx.body_models import SMPLFX
    model = SMPLFX(**smplx_kwargs).to(device)
    out = runBatchedSMPLX(params, model, BS=BS, with_vertices=run_with_verts, device=device)
    return out, model

def compute_normals_diff(vertices, faces, eps=1e-6):
    """
    Differentiable computation of vertex and face normals.
    Args:
        vertices: (V, 3)
        faces: (F, 3) long
    Returns:
        vertex_normals: (V, 3)
        face_normals: (F, 3)
    """
    v0 = vertices[faces[:, 0]]  # (F, 3)
    v1 = vertices[faces[:, 1]]
    v2 = vertices[faces[:, 2]]

    face_normals = torch.cross(v1 - v0, v2 - v0, dim=1)  # (F, 3)
    face_normals = face_normals / (face_normals.norm(dim=1, keepdim=True) + eps)

    # Prepare indices to scatter into vertex normals
    F = faces.shape[0]
    V = vertices.shape[0]

    idx = faces.reshape(-1)          # (3*F,)
    face_normals_repeat = face_normals.repeat_interleave(3, dim=0)  # (3*F, 3)

    vertex_normals = torch.zeros_like(vertices)
    vertex_normals = vertex_normals.scatter_add(0, idx[:, None].expand(-1, 3), face_normals_repeat)
    vertex_normals = vertex_normals / (vertex_normals.norm(dim=1, keepdim=True) + eps)

    return vertex_normals, face_normals

def convertSMPLX2SMPLFXKwargs(kwargs, num_betas_flame=100, v_mask_path_flame=None, corr_fname=None):
	kwargs['gender'] = 'neutral'
	kwargs.setdefault('ext', 'pkl')
	if v_mask_path_flame is None:
		raise ValueError("v_mask_path_flame must be provided for SMPLFX model")
	kwargs['v_mask_path_flame'] = v_mask_path_flame
	if corr_fname is None:
		raise ValueError("corr_fname must be provided for SMPLFX model")
	kwargs['corr_fname'] = corr_fname
	kwargs['use_eyelids'] = True
	kwargs['add_teeth'] = True
	kwargs['num_betas_flame'] = num_betas_flame
	kwargs['with_mouth_interior'] = True
	kwargs['mount_interior_resolution_u'] = 10
	kwargs['mount_interior_resolution_v'] = 20
	return kwargs


def flipPoseRot(poses):
    # poses ... x 3 x 3
    assert poses.shape[-1] == 3 and poses.shape[-2] == 3
    flip_mask = torch.tensor([
        [ 1, -1, -1],
        [-1,  1,  1],
        [-1,  1,  1]
    ], device=poses.device, dtype=poses.dtype)
    while flip_mask.ndim < poses.ndim:
        flip_mask = flip_mask[None]
    flipped_poses = poses * flip_mask
    return flipped_poses


import pytorch3d
def remapIndependentHandToBody(right_hand_orient, left_hand_orient, body_pose, global_orient, hands_format, body_format, conv='ZYX'):
    """
    Given hand poses extracted separately (e.g., HaMeR/WiLor) this code remaps their more accurate hand orientation to the body pose.
    """
    assert len(right_hand_orient) == len(left_hand_orient) == len(body_pose) == len(global_orient)

    if body_format == 'euler':
        assert body_pose.ndim == 2
        assert global_orient.ndim == 2
        body_pose_mat = pytorch3d.transforms.euler_angles_to_matrix(body_pose.reshape(-1, 21, 3), conv)
        global_orient_mat = pytorch3d.transforms.euler_angles_to_matrix(global_orient, conv)
    elif body_format == 'rot':
        assert body_pose.ndim == 4
        assert global_orient.ndim == 3
        body_pose_mat = body_pose
        global_orient_mat = global_orient
    else:
        raise NotImplementedError
    
    R_global_elbow_l = global_orient_mat @ body_pose_mat[:,2] @ body_pose_mat[:,5] @ body_pose_mat[:,8] @ body_pose_mat[:,11] @ body_pose_mat[:,15] @ body_pose_mat[:,17]
    R_global_elbow_r = global_orient_mat @ body_pose_mat[:,2] @ body_pose_mat[:,5] @ body_pose_mat[:,8] @ body_pose_mat[:,11] @ body_pose_mat[:,16] @ body_pose_mat[:,18]

    if hands_format == 'euler':
        assert right_hand_orient.ndim == 2
        assert left_hand_orient.ndim == 2
        right_hand_orient_mat = pytorch3d.transforms.euler_angles_to_matrix(right_hand_orient, conv)
        left_hand_orient_mat = pytorch3d.transforms.euler_angles_to_matrix(left_hand_orient, conv)
    elif hands_format == 'rot':
        assert right_hand_orient.ndim == 3
        assert left_hand_orient.ndim == 3
        right_hand_orient_mat = right_hand_orient
        left_hand_orient_mat = left_hand_orient
    else:
        raise NotImplementedError

    body_elbow_r = R_global_elbow_r.transpose(-1,-2) @ right_hand_orient_mat
    body_elbow_l = R_global_elbow_l.transpose(-1,-2) @ left_hand_orient_mat

    body_pose = body_pose.clone()
    if body_format == 'euler':
        body_elbow_r = pytorch3d.transforms.matrix_to_euler_angles(body_elbow_r, conv)
        body_elbow_l = pytorch3d.transforms.matrix_to_euler_angles(body_elbow_l, conv)
        body_pose[...,19*3:19*3+3] = body_elbow_l # left hand
        body_pose[...,20*3:20*3+3] = body_elbow_r # right hand
    elif body_format == 'rot':
        body_pose[...,19,:,:] = body_elbow_l
        body_pose[...,20,:,:] = body_elbow_r

    return body_pose


def cross_product_indices(idxs1, idxs2):
    all_indices = []
    for i in range(len(idxs1)):
        for j in range(len(idxs2)):
            all_indices.append([idxs1[i], idxs2[j]])
    return np.array(all_indices)

class Contacts():
    def __init__(self, on_joints=True, is_1to1_connections=True):
        self.on_joints = on_joints
        self.is_1to1_connections = is_1to1_connections
        self.right_hand_joints = np.array([21, 52, 53, 54, 71, 40, 41, 42, 72, 43, 44, 45, 73, 49, 50, 51, 74, 46, 47, 48, 75])
        self.left_hand_joints = np.array([20, 37, 38, 39, 66, 25, 26, 27, 67, 28, 29, 30, 68, 34, 35, 36, 69, 31, 32, 33, 70])
        self.face_indices = np.array(np.arange(95, 107, 1).tolist() + np.arange(107, 119, 1).tolist())
        if on_joints:
            if is_1to1_connections:
                self.contact_connections = np.concatenate((self.right_hand_joints[:,None], self.left_hand_joints[:,None]), axis=-1)
            else:
                self.contact_connections = cross_product_indices(self.right_hand_joints, self.left_hand_joints)

            self.contact_connections = np.concatenate((self.contact_connections, cross_product_indices(self.right_hand_joints, self.face_indices)), axis=0)
            self.contact_connections = np.concatenate((self.contact_connections, cross_product_indices(self.left_hand_joints, self.face_indices)), axis=0)            
        else:
            raise NotImplementedError()

    def get_relative_contacts_joints(self, joints):
        return joints[:, self.contact_connections[:,1]] - joints[:, self.contact_connections[:,0]]

    def get_relative_contacts(self, output):
        if self.on_joints:
            return self.get_relative_contacts_joints(output.joints)
        else:
            raise NotImplementedError()

    

def getContactWeights(relative_contacts, max_dist_contact):
    contact_distances = torch.linalg.norm(relative_contacts, dim=-1, keepdim=True)
    mask = (contact_distances < max_dist_contact).float()
    return mask



class SMPLX2DJointMapper():
	def __init__(self, with_hips=False, with_legs=True, with_top_feet=True, with_mp_main_face=False, with_hand_palm=False, with_shoulders=True, detector_name='rtmp', use_face_contour=False, is_full_face=False, use_full_mouth=False):
		if use_face_contour:
			assert is_full_face, "use_face_contour can be True only if is_full_face is also True"
		self.use_face_contour = use_face_contour
		if with_shoulders:
			smplx_body_idxs = [17, 19, 21, 16, 18, 20]
			self.edges = [[0,1], [1,2], [3,4], [4,5], [0,3]]

			if detector_name == 'rtmp':
				body_idxs = [6, 8, 10, 5, 7, 9]
			elif detector_name == 'openpose-25':
				body_idxs = [2, 3, 4, 5, 6, 7]
			else: raise ValueError('Unknown detector name: ' + detector_name)
		else:
			smplx_body_idxs = [19, 21, 18, 20]
			self.edges = [[0,1], [2,3]]

			if detector_name == 'rtmp':
				body_idxs = [8, 10, 7, 9]
			elif detector_name == 'openpose-25':
				body_idxs = [3, 4, 6, 7]
			else: raise ValueError('Unknown detector name: ' + detector_name)

		if with_hips:
			smplx_body_idxs += [2, 1]	# right hip, left hip
			self.edges += [[len(smplx_body_idxs)-2, len(smplx_body_idxs)-1]]

			if detector_name == 'rtmp':
				body_idxs += [12, 11] # right hip, left hip
			elif detector_name == 'openpose-25':
				body_idxs += [9, 12]
			else: raise ValueError('Unknown detector name: ' + detector_name)

		if with_legs:
			smplx_body_idxs += [5, 8, 4, 7]
			self.edges += [[len(smplx_body_idxs)-4, len(smplx_body_idxs)-3], [len(smplx_body_idxs)-2, len(smplx_body_idxs)-1]]
			if detector_name == 'rtmp':
				body_idxs += [14, 16, 13, 15]
			elif detector_name == 'openpose-25':
				assert False
			else: raise ValueError('Unknown detector name: ' + detector_name)

			if with_top_feet:
				smplx_body_idxs += [63, 60]

				if detector_name == 'rtmp':
					body_idxs += [20, 17]
				elif detector_name == 'openpose-25':
					assert False
				else: raise ValueError('Unknown detector name: ' + detector_name)

		if with_mp_main_face:
			smplx_body_idxs += [95, 104] + [107, 113]

			if detector_name == 'rtmp':
				body_idxs += [59, 68] + [71, 77]
			elif detector_name == 'openpose-25':
				assert False
			else: raise ValueError('Unknown detector name: ' + detector_name)

		if with_hand_palm:
			npts = len(smplx_body_idxs)
			smplx_body_idxs += [40, 46] + [25, 31]
			self.edges += [[2, npts  ], [2, npts+1], [npts  , npts+1]]
			self.edges += [[5, npts+2], [5, npts+3], [npts+2, npts+3]]

			if detector_name == 'rtmp':
				body_idxs += [117, 129] + [96, 108]
			elif detector_name == 'openpose-25':
				assert False
			else: raise ValueError('Unknown detector name: ' + detector_name)

		self.smplx_rhand_idxs = [21, 52, 53, 54, 71, 40, 41, 42, 72, 43, 44, 45, 73, 49, 50, 51, 74, 46, 47, 48, 75]
		self.smplx_lhand_idxs = [20, 37, 38, 39, 66, 25, 26, 27, 67, 28, 29, 30, 68, 34, 35, 36, 69, 31, 32, 33, 70]

		# smplx_face_idxs_no_contour = [76,77,78,79,80] + [81,82,83,84,85] + [86,87,88,89] + [90,91,92,93,94] + [95,96,97,98,99,100] + [104,103,102,101,106,105] + [107,108,109,110,111,112,113,114,115,116,117,118] + [119,120,121,122,123,124,125,126] + [24, 23]

		smplx_face_idxs_no_contour = np.arange(76, 127).tolist() + [24, 23]

		npts = len(body_idxs) + len(self.smplx_rhand_idxs) + len(self.smplx_lhand_idxs)
		self.face_edges = []
		if detector_name == 'rtmp':
			self.rhand_idxs = np.arange(133-21, 133).tolist()
			self.lhand_idxs = np.arange(133-42, 133-21).tolist()
			if is_full_face:
				if use_face_contour:
					self.face_idxs = np.arange(23, 23+68).tolist() + [2, 1]
					self.smplx_face_idxs = np.arange(127, 127+17).tolist() + smplx_face_idxs_no_contour
					self.mouth_face_subidxs = None
				else:
					self.face_idxs = np.arange(40, 23+68).tolist() + [2, 1]
					self.smplx_face_idxs = smplx_face_idxs_no_contour
					self.mouth_face_subidxs = None
			else:
				if use_full_mouth:
					self.face_idxs = (np.arange(59, 71).tolist() + np.arange(71, 91, 1).tolist())
					self.smplx_face_idxs = np.arange(95, 107, 1).tolist() + np.arange(107, 127, 1).tolist()
					self.face_edges += getLoopedEdgesPairs(6) + getLoopedEdgesPairs(6,6) + getLoopedEdgesPairs(12,12) + getLoopedEdgesPairs(8,24)
					self.mouth_face_subidxs = np.arange(len(self.face_idxs)-20, len(self.face_idxs), 1)
				else:
					self.face_idxs = (np.arange(59, 71).tolist() + np.arange(71, 91, 1).tolist())[:24]
					self.smplx_face_idxs = np.arange(95, 107, 1).tolist() + np.arange(107, 119, 1).tolist()
					self.mouth_face_subidxs = np.arange(len(self.face_idxs)-12, len(self.face_idxs), 1)
	
		elif detector_name == 'openpose-25':
			self.rhand_idxs = np.arange(25, 25+21).tolist()
			self.lhand_idxs = np.arange(25+21, 25+42).tolist()
			if is_full_face:
				# right eyebrow, left eyebrow, nose bridge, nose bottom, right eye, left eye, outer lips, inner lips, pupils
				if use_face_contour:
					self.face_idxs = np.arange(25+42, 25+42+70).tolist()
					self.smplx_face_idxs = np.arange(127, 127+17).tolist() + smplx_face_idxs_no_contour
					# openpose scheme: https://www.researchgate.net/figure/Openpose-detected-body-hand-and-face-keypoints_fig2_353668783
					self.mouth_face_subidxs = None
				else:
					self.face_idxs = np.arange(25+42+17, 25+42+70).tolist()
					self.smplx_face_idxs = smplx_face_idxs_no_contour
					self.mouth_face_subidxs = None
			else:
				if use_full_mouth:
					self.face_idxs = (np.arange(25+42, 25+42+12).tolist() + np.arange(25+42+12, 25+42+32, 1).tolist())
					self.smplx_face_idxs = np.arange(95, 107, 1).tolist() + np.arange(107, 127, 1).tolist()
					self.mouth_face_subidxs = np.arange(len(self.face_idxs)-20, len(self.face_idxs), 1)
				else:
					self.face_idxs = (np.arange(25+42, 25+42+12).tolist() + np.arange(25+42+12, 25+42+12+12, 1).tolist())
					self.smplx_face_idxs = np.arange(95, 107, 1).tolist() + np.arange(107, 119, 1).tolist()
					self.mouth_face_subidxs = np.arange(len(self.face_idxs)-12, len(self.face_idxs), 1)
		else:
			raise ValueError('Unknown detector name: ' + detector_name)

		if len(self.face_edges) > 0:
			self.edges = self.edges + (npts + np.array(self.face_edges)).tolist()

		self.smplx_hand_idxs = self.smplx_rhand_idxs + self.smplx_lhand_idxs
		self.hand_idxs = self.rhand_idxs + self.lhand_idxs

		edges_one_hand = [[0, 1], [1, 2], [2, 3], [3, 4], [0, 5], [5, 6], [6, 7], [7, 8], [5, 9], [9, 10], [10, 11], [11, 12], [9, 13], [13, 14], [14, 15], [15, 16], [13, 17], [17, 18], [18, 19], [19, 20], [17, 0]]
		self.edges += (len(smplx_body_idxs) + np.array(edges_one_hand)).tolist() + (len(smplx_body_idxs) + len(self.smplx_rhand_idxs) + np.array(edges_one_hand)).tolist()
		self.smplx_idxs = smplx_body_idxs + self.smplx_rhand_idxs + self.smplx_lhand_idxs + self.smplx_face_idxs

		self.body_idxs = body_idxs
		self.smplx_body_idxs = smplx_body_idxs

		self.idxs = body_idxs + self.rhand_idxs + self.lhand_idxs + self.face_idxs
		self.face_subidxs = np.arange(len(self.idxs)-len(self.face_idxs), len(self.idxs), 1) # same for 2D detections and SMPLX
		self.body_subidxs = np.arange(0, len(self.body_idxs), 1)
		self.hand_subidxs = np.arange(len(self.body_idxs), len(self.body_idxs)+len(self.hand_idxs), 1)

		self.smplx_idxs_no_face = self.smplx_idxs[:-len(self.smplx_face_idxs)]
		self.idxs_no_face = self.idxs[:-len(self.face_idxs)]

		self.smplx_body_mask_no_face = torch.zeros(len(self.smplx_idxs) - len(self.smplx_face_idxs), dtype=torch.bool)
		self.smplx_body_mask_no_face[:len(self.smplx_body_idxs)] = True
		self.smplx_hand_mask_no_face = torch.zeros(len(self.smplx_idxs) - len(self.smplx_face_idxs), dtype=torch.bool)
		self.smplx_hand_mask_no_face[len(self.smplx_body_idxs):] = True

		assert len(self.idxs) == len(self.smplx_idxs)
		assert np.array(self.edges).max() < len(self.idxs)
