import torch
import smplx
import numpy as np
import time
import copy

from .vision3d import projectP
from .losses import SegmentTriangleMeshIntersectionLossDualHands, getLossFnc, compute_pearson_depth_loss
from .temporal_losses import getAccelerationStats, getVelocityStats, getJerkStats
from .bspline_kernel import NeuralSplineFilterDecoder
from .smpl_fx.body_models import SMPLFX
from .maths_utils import avgL2distance, euler_to_6d_pytorch3d
from .utils import PrintColors, createChunkedTemporalIndices, getTemporalIndicesBasedKeys, matrix_to_euler_angles, euler_angles_to_matrix, calculateInitialLogits
from .utils_smplx import getContactWeights

def is_shape(x):
    return x == 'shape' or x == 'betas'

class SMPLXWrapper():
    def __init__(self, NUM_POSES, is_smplfx, smplx_kwargs, smplx_params_init, smplx_cp_params_init, chunk_size, optimize_shape, optimize_hands, optimize_smplx_rt, are_control_points, device, GLOBAL_ORIENT_ZERO_X, GLOBAL_ORIENT_FIX_Z, right_hand_constraints, left_hand_constraints, body_zero_angles_idxs, LRs, USE_SIGMOID_ACTIVATION, spacing, apply_constraints_before_decode, OPTIMIZE_FLAME_SHAPE, body_constraints, is_single_smplx_rt, OPTIMIZE_FACES, IS_BODY_POSE_6D, OPTIMIZE_BODY_POSE):
        self.are_control_points = are_control_points
        self.USE_SIGMOID_ACTIVATION = USE_SIGMOID_ACTIVATION
        if USE_SIGMOID_ACTIVATION:
            print(PrintColors.WARNING + "Using Sigmoid Activation for constraints enforcement seems worse than clamping!" + PrintColors.ENDC)

        self.device = device
        if is_smplfx:
            self.model = SMPLFX(**smplx_kwargs).to(device)
        else:
            self.model = smplx.create(**smplx_kwargs).to(device)
        for param in self.model.parameters():
            param.requires_grad = False

        self.LRs = LRs
        self.is_single_smplx_rt = is_single_smplx_rt
        self.smplx_params = {}
        self.OPTIM_J_REGRESSOR = LRs.get('J_regressor', 0) > 0
        J_regressor_copy = self.model.J_regressor.detach().clone().to(device)
        J_regressor_bool = J_regressor_copy > 0
        self.J_regressor_zero_mask = J_regressor_bool.type(torch.float32)
        if 'J_regressor' in smplx_params_init:
            J_regressor_init = smplx_params_init['J_regressor']
        elif self.OPTIM_J_REGRESSOR:
            J_regressor_init = J_regressor_copy
        else:
            J_regressor_init = None
        if J_regressor_init is None:
            self.J_regressor_log = None
        else:
            self.J_regressor_log = torch.log(J_regressor_init.detach().clone().to(device) + 1e-9).clone().requires_grad_(self.OPTIM_J_REGRESSOR)

        self.constraints = {'eyelid_params': torch.tensor([[0.0, 1.0], [0.0, 1.0]], device=device)}
        if right_hand_constraints is not None:
            assert right_hand_constraints.shape == (15*3, 2)
            self.constraints['right_hand_pose'] = right_hand_constraints.to(device)
        if left_hand_constraints is not None:
            assert left_hand_constraints.shape == (15*3, 2)
            self.constraints['left_hand_pose'] = left_hand_constraints.to(device)
        if GLOBAL_ORIENT_ZERO_X or GLOBAL_ORIENT_FIX_Z:
            self.constraints['global_orient'] = torch.tensor([[-np.pi, np.pi], [-np.pi, np.pi], [-np.pi, np.pi]], device=device)
            if GLOBAL_ORIENT_ZERO_X:
                self.constraints['global_orient'][0,:] = 0.0
            if GLOBAL_ORIENT_FIX_Z:
                self.constraints['global_orient'][2,:] = np.pi
        self.body_zero_angles_idxs = body_zero_angles_idxs
        if body_zero_angles_idxs is not None:
            if body_constraints is not None:
                assert body_constraints.shape == (21*3, 2)
                self.constraints['body_pose'] = body_constraints.to(device)
            else:
                self.constraints['body_pose'] = torch.tensor([[-np.pi, np.pi]], device=device).repeat(21*3,1)
            for idx in body_zero_angles_idxs:
                self.constraints['body_pose'][idx*3:idx*3+3,:] = 0.0

        self.constraints_diff = {k: v[:,1] - v[:,0] for k,v in self.constraints.items()}

        params_dim = {'transl': 3, 'global_orient': 3, 'body_pose': 21*3, 'left_hand_pose': 15*3, 'right_hand_pose': 15*3, 'jaw_pose': 3, 'expression': smplx_kwargs.get('num_expression_coeffs', 10), 'eyelid_params': 2, 'betas': smplx_kwargs.get('num_betas', 10)}
        params_optim = {'transl': optimize_smplx_rt, 'global_orient': optimize_smplx_rt, 'body_pose': OPTIMIZE_BODY_POSE, 'left_hand_pose': optimize_hands, 'right_hand_pose': optimize_hands, 'jaw_pose': OPTIMIZE_FACES, 'expression': OPTIMIZE_FACES, 'eyelid_params': OPTIMIZE_FACES, 'betas': optimize_shape}
        self.is_single = lambda x : is_shape(x) or (x == 'transl' and is_single_smplx_rt) or (x == 'global_orient' and is_single_smplx_rt)
        get_poses = lambda x : 1 if self.is_single(x) else NUM_POSES
        
        if is_smplfx:
            params_dim['shape'] = self.model.FLAME.betas.shape[-1]
            params_optim['shape'] = OPTIMIZE_FLAME_SHAPE

        for key in params_dim.keys():
            if key in smplx_params_init:
                if self.is_single(key):
                    self.smplx_params[key] = smplx_params_init[key].to(device).clone().detach().median(dim=0).values[None].requires_grad_(params_optim[key])
                else:
                    self.smplx_params[key] = smplx_params_init[key].to(device).clone().detach().requires_grad_(params_optim[key])
                assert self.smplx_params[key].ndim == 2
            else:
                self.smplx_params[key] = torch.zeros((get_poses(key), params_dim[key]), device=device, requires_grad=params_optim[key])
            print('smplx params', key, self.smplx_params[key].shape, 'optim:', params_optim[key])
            assert self.smplx_params[key].isnan().any() == False, f'SMPLX param {key} has NaNs!'
            assert self.smplx_params[key].isinf().any() == False, f'SMPLX param {key} has Infs!'

        self.IS_BODY_POSE_6D = IS_BODY_POSE_6D
        if IS_BODY_POSE_6D:
            if not are_control_points:
                print(PrintColors.WARNING + "Using 6D representation is optimal for control points! " + PrintColors.ENDC)
            params_dim['body_pose'] = 21*6
            assert body_constraints is None
            self.constraints.pop('body_pose', None)
            self.smplx_params['body_pose'] = euler_to_6d_pytorch3d(self.smplx_params['body_pose'])
            assert self.smplx_params['body_pose'].isnan().any() == False, f'SMPLX param body_pose has NaNs after 6D conversion!'
            assert self.smplx_params['body_pose'].isinf().any() == False, f'SMPLX param body_pose has Infs after 6D conversion!'
        else:
            if are_control_points and OPTIMIZE_BODY_POSE:
                print(PrintColors.WARNING + "For control points it is better to use 6D representation!" + PrintColors.ENDC)

        with torch.no_grad():
            self.smplx_params['global_orient'][:,2] = np.pi

        self.params_optim = params_optim
        self.optimized_params = {}
        if are_control_points:
            for key in self.smplx_params.keys():
                if (not self.is_single(key)) and params_optim[key]:
                    self.smplx_params[key] = self.smplx_params[key].detach().requires_grad_(False)
                    self.optimized_params[key] = NeuralSplineFilterDecoder(NUM_POSES, params_dim[key], chunk_size, spacing=spacing, constraints=self.constraints.get(key, None) if self.USE_SIGMOID_ACTIVATION else None, init_cp_params=smplx_cp_params_init.get(key, None), init_params=self.smplx_params[key], apply_constraints_before_decode=apply_constraints_before_decode).to(device)
                    self.num_chunks = self.optimized_params[key].num_chunks
                    self.knots_per_chunk = self.optimized_params[key].knots_per_chunk
                    self.total_knots = self.optimized_params[key].total_knots
                    # todo: check the initial decoder estimates so they match the init params!

                    if key in self.constraints and (not self.USE_SIGMOID_ACTIVATION):
                        for chunk_idx in range(self.num_chunks):
                            with torch.no_grad():
                                self.optimized_params[key].chunks[chunk_idx].clamp_(min=self.constraints[key][:,0][None,:,None], max=self.constraints[key][:,1][None,:,None])
                                assert not torch.isnan(self.optimized_params[key].chunks[chunk_idx]).any(), f'NaNs in optimized param {key} chunk {chunk_idx} after clamping!'
                                assert not torch.isinf(self.optimized_params[key].chunks[chunk_idx]).any(), f'Infs in optimized param {key} chunk {chunk_idx} after clamping!'

                    for chunk_idx in range(self.num_chunks):
                        with torch.no_grad():
                            test_decoded, _ = self.optimized_params[key](chunk_idx, get_loss=False)
                            assert not torch.isnan(test_decoded).any(), f'NaNs in optimized param {key} chunk {chunk_idx} test_decoded!'
                            assert not torch.isinf(test_decoded).any(), f'Infs in optimized param {key} chunk {chunk_idx} test_decoded!'

                        diff_init = torch.linalg.norm(test_decoded - self.smplx_params[key][chunk_idx*self.knots_per_chunk*spacing:(chunk_idx+1)*self.knots_per_chunk*spacing], dim=-1)
                        print('diff mean', diff_init.mean().item(), 'max', diff_init.max().item(), 'median', diff_init.median().item(), 'for key', key, 'chunk', chunk_idx)
                        # if key == 'body_pose':
                        #     with torch.no_grad():
                        #         smplx_out1 = self.model(**{key: test_decoded})
                        #         smplx_out2 = self.model(**{key: self.smplx_params[key][chunk_idx*self.knots_per_chunk*spacing:(chunk_idx+1)*self.knots_per_chunk*spacing]})
                        #         plotMesh([smplx_out1['vertices'][-14]], self.model.faces, title='check decoded vs init for '+key+' chunk '+str(chunk_idx), show=False, elev=90, azim=-90)
                        #         plotMesh([smplx_out2['vertices'][-14]], self.model.faces, title='check init for '+key+' chunk '+str(chunk_idx), show=True, elev=90, azim=-90)
                        #         pdb.set_trace()

        else:
            self.num_chunks = (NUM_POSES + chunk_size - 1) // chunk_size
            for key in self.smplx_params.keys():
                if (not self.is_single(key)) and params_optim[key]:
                    self.smplx_params[key] = self.smplx_params[key].detach().requires_grad_(False)
                    if key in self.constraints and self.USE_SIGMOID_ACTIVATION:
                        self.smplx_params[key] = calculateInitialLogits(self.smplx_params[key].clone(), self.constraints[key])
                        assert not torch.isnan(self.smplx_params[key]).any()
                        assert not torch.isinf(self.smplx_params[key]).any()
                    self.optimized_params[key] = torch.chunk(self.smplx_params[key].clone(), self.num_chunks, dim=0)
                    self.optimized_params[key] = [p.requires_grad_(True) for p in self.optimized_params[key]]

        for key in self.optimized_params.keys():
            # print('POP parameter', key)
            self.smplx_params.pop(key, None)
        
        for key in self.smplx_params.keys():
            if not self.is_single(key):
                # split into chunks:
                if are_control_points:
                    raise NotImplementedError("Chunking not implemented for control points optimization yet!")
                else:
                    self.smplx_params[key] = torch.chunk(self.smplx_params[key], self.num_chunks, dim=0)

    def constraints_activation(self, x, param_key):
        if (param_key not in self.constraints) or (not self.USE_SIGMOID_ACTIVATION) or (not self.params_optim[param_key]):
            return x
        return torch.sigmoid(x) * self.constraints_diff[param_key] + self.constraints[param_key][:,0]

    def activateJRegressor(self):
        return torch.nn.functional.softmax(self.J_regressor_log * self.J_regressor_zero_mask + (1 - self.J_regressor_zero_mask) * (-1e9), dim=-1)

    def getOptimizableParamsList(self):
        params_list = []
        if self.OPTIM_J_REGRESSOR:
            print(f"Optimizing J_regressor with lr {self.LRs['J_regressor']} device {self.J_regressor_log.device}")
            params_list.append({'params': self.J_regressor_log, 'lr': self.LRs['J_regressor'], 'name': 'J_regressor_log'})
        for key in self.smplx_params.keys():
            if self.params_optim[key]:
                print(f"Optimizing SMPLX param: {key} with lr {self.LRs.get(key, self.LRs['lr'])} device {self.smplx_params[key].device}")
                params_list.append({'params': self.smplx_params[key], 'lr': self.LRs.get(key, self.LRs['lr']), 'name': key})
        for key in self.optimized_params.keys():
            if isinstance(self.optimized_params[key], NeuralSplineFilterDecoder):
                print(f"Optimizing SMPLX control points param: {key} with lr {self.LRs.get(key, self.LRs['lr'])} device {self.optimized_params[key].chunks[0].device}")
                params_list.append({'params': self.optimized_params[key].parameters(), 'lr': self.LRs.get(key, self.LRs['lr']), 'name': key})
            else:
                print(f"Optimizing SMPLX chunked param: {key} with lr {self.LRs.get(key, self.LRs['lr'])} device {self.optimized_params[key][0].device}")
                params_list.append({'params': self.optimized_params[key], 'lr': self.LRs.get(key, self.LRs['lr']), 'name': key})
        return params_list

    def apply_constraints(self, chunk_idx):
        if not self.USE_SIGMOID_ACTIVATION:
            with torch.no_grad():
                for key in self.constraints.keys():
                    if key in self.optimized_params.keys():
                        if self.are_control_points:
                            # chunks are (1, D, C) and constraints are (D), so unsqueeze accordingly to 1 x D x 1
                            self.optimized_params[key].chunks[chunk_idx].clamp_(min=self.constraints[key][:,0][None,:,None], max=self.constraints[key][:,1][None,:,None])
                            if chunk_idx-1 >= 0:
                                self.optimized_params[key].chunks[chunk_idx-1].clamp_(min=self.constraints[key][:,0][None,:,None], max=self.constraints[key][:,1][None,:,None])
                            elif chunk_idx+1 < self.num_chunks:
                                self.optimized_params[key].chunks[chunk_idx+1].clamp_(min=self.constraints[key][:,0][None,:,None], max=self.constraints[key][:,1][None,:,None])
                        else:
                            self.optimized_params[key][chunk_idx].clamp_(min=self.constraints[key][:,0][None], max=self.constraints[key][:,1][None])
                    else:
                        if key in self.smplx_params.keys() and self.params_optim[key]:
                            self.smplx_params[key].clamp_(min=self.constraints[key][:,0][None], max=self.constraints[key][:,1][None])
                if self.IS_BODY_POSE_6D and self.body_zero_angles_idxs is not None:
                    if self.are_control_points:
                        reshaped = self.optimized_params['body_pose'].chunks[chunk_idx].reshape(1, 21, 6, -1)
                        reshaped[:, self.body_zero_angles_idxs] = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], device=self.device)[None,None,:,None]
                        if chunk_idx-1 >= 0:
                            reshaped1 = self.optimized_params['body_pose'].chunks[chunk_idx-1].reshape(1, 21, 6, -1)
                            reshaped1[:, self.body_zero_angles_idxs] = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], device=self.device)[None,None,:,None]
                        elif chunk_idx+1 < self.num_chunks:
                            reshaped2 = self.optimized_params['body_pose'].chunks[chunk_idx+1].reshape(1, 21, 6, -1)
                            reshaped2[:, self.body_zero_angles_idxs] = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], device=self.device)[None,None,:,None]
                    else:
                        reshaped = self.optimized_params['body_pose'][chunk_idx].reshape(-1, 21, 6)
                        reshaped[:, self.body_zero_angles_idxs] = torch.tensor([1.0, 0.0, 0.0, 0.0, 1.0, 0.0], device=self.device)[None,None,:]

    def getParams(self):
        smplx_params = self.getSMPLXParams(chunk_idx=-1, get_loss=False)[0]
        for key in smplx_params.keys():
            smplx_params[key] = smplx_params[key].detach().clone()

        if self.are_control_points:
            optimized_params = {key: copy.deepcopy(self.optimized_params[key].state_dict()) for key in self.optimized_params.keys()}
            return smplx_params, optimized_params
        else:
            return smplx_params, {}

    def getSMPLXParams(self, chunk_idx, get_loss=True):
        smplx_params_chunk = {}
        for key in self.smplx_params.keys():
            if self.is_single(key):
                smplx_params_chunk[key] = self.smplx_params[key]
            else:
                if chunk_idx == -1:
                    smplx_params_chunk[key] = torch.cat(self.smplx_params[key], dim=0)
                else:
                    smplx_params_chunk[key] = self.smplx_params[key][chunk_idx]
        if self.J_regressor_log is not None:
            smplx_params_chunk['J_regressor'] = self.activateJRegressor()

        extra_losses = {}
        if chunk_idx == -1:
            assert not get_loss # assume loss is not needed
            params_list = {key: [] for key in self.optimized_params.keys()}
            for c_idx in range(self.num_chunks):
                for key in self.optimized_params.keys():
                    if self.are_control_points:
                        params_list[key].append(self.optimized_params[key](c_idx, get_loss=False)[0])
                    else:
                        params_list[key].append(self.constraints_activation(self.optimized_params[key][c_idx], key))
            for key in params_list.keys():
                smplx_params_chunk[key] = torch.cat(params_list[key])
        else:
            if self.are_control_points:
                for key in self.optimized_params.keys():
                    smplx_params_chunk[key], loss_key = self.optimized_params[key](chunk_idx, get_loss=get_loss)
                    if len(loss_key) > 0 and get_loss:
                        extra_losses[key] = loss_key
            else:
                for key in self.optimized_params.keys():
                    smplx_params_chunk[key] = self.constraints_activation(self.optimized_params[key][chunk_idx], key)

        # for key in smplx_params_chunk.keys():
        #     assert not torch.isnan(smplx_params_chunk[key]).any(), f'NaNs in SMPLX param {key} for chunk {chunk_idx}!'
        #     assert not torch.isinf(smplx_params_chunk[key]).any(), f'Infs in SMPLX param {key} for chunk {chunk_idx}!'

        return smplx_params_chunk, extra_losses

    def __call__(self, chunk_idx):
        smplx_params_chunk, extra_losses = self.getSMPLXParams(chunk_idx, get_loss=True)
        output = self.model(**smplx_params_chunk)
        return output, extra_losses, smplx_params_chunk


def chunkTensor(tensor, are_control_points, num_chunks, knots_per_chunk, spacing):
    if are_control_points:
        tensor_list = []
        for i in range(num_chunks):
            tensor_list.append(tensor[i*knots_per_chunk*spacing:(i + 1) * knots_per_chunk*spacing])
    else:
        tensor_list = torch.chunk(tensor, num_chunks, dim=0)
    return tensor_list

class SMPLXOptimizer():
    def __init__(self, smplx_kwargs, smplx_params_init, right_hand_constraints, left_hand_constraints, K, RT, mapper, poses_2d=None, poses_2d_weights=None, device='cuda', MAX_ITERS=10000, min_lr=5e-6, chunk_size=100000, loss_fn='huber', optimize_shape=True, optimize_hands=True, optimize_smplx_rt=True, OPTIMIZE_RT=True, OPTIMIZE_RT_FIX_CAM_0=True, MESH_INTERSECTION_LOSS_WEIGHT=0., other_params_init={}, lr=1e-3, SMOOTHNESS_TERMS=[], is_smplfx=True, keys=None, PRINT_EVERY=10, SCHEDULER_PATIENCE=20, GLOBAL_ORIENT_ZERO_X=False, GLOBAL_ORIENT_FIX_Z=True, body_zero_angles_idxs=None, LRs={}, are_control_points=False, REG_NORM_EXPRESSION_W=0.0, BETA_REG_W=1e-3, USE_SIGMOID_ACTIVATION=False, spacing=3, APPLY_SMOOTHNESS_ON_CP=True, apply_constraints_before_decode=False, ANGLES_REG_W=0.0, ANGLES_REG_IDXS=[2, 5, 8, 11, 14, 12, 13], OPTIMIZE_FLAME_SHAPE=True, body_constraints=None, is_single_smplx_rt=False, poses_3d=None, OPTIMIZE_FACES=True, IS_BODY_POSE_6D=False, OPTIMIZE_BODY_POSE=True, depths=None, depths_mask=None, DEPTH_LOSS_WEIGHT=0.0, z_3d=None, Z_3D_LOSS_WEIGHT=0.0, max_dist_contacts=0.075, segmentation_path=None, smplx_model_path=None):
        """
        K : (C, 3, 3) camera intrinsics
        RT : (C, 3, 4) camera extrinsics
        poses_2d: (N, C, J, 2) 2D keypoints
        poses_2d_weights: (N, C, J, 1) weights for 2D keypoints
        poses_3d: (N, J, 3) 3D keypoints for optional 3D supervision
        depths: (N, C, J) relative depths for 2D keypoints
        """
        if poses_2d is not None:
            NUM_POSES = poses_2d.shape[0]
        else:
            assert poses_3d is not None, "Either 2D or 3D poses must be provided!"
            NUM_POSES = poses_3d.shape[0]

        self.mapper = mapper
        for triplet in SMOOTHNESS_TERMS:
            assert triplet[1] in ['velocity', 'acceleration', 'jerk']
            assert isinstance(triplet[2], float)

        self.smoothness_losses = {'velocity': getVelocityStats, 'acceleration': getAccelerationStats, 'jerk': getJerkStats}
        if len(SMOOTHNESS_TERMS) == 0:
            print(PrintColors.WARNING + "Smoothness terms are important for SMPLX optimization!" + PrintColors.ENDC)

        if K is not None:
            assert K.ndim == 3 and K.shape[1:] == (3, 3)
            self.K = K
        if RT is not None:
            assert RT.ndim == 3 and RT.shape[1:] == (3, 4)
            self.P = K @ RT

        self.conv = 'XYZ'

        if OPTIMIZE_RT:
            self.r_extrinsics, self.t_extrinsics = matrix_to_euler_angles(RT[...,:3,:3].detach().clone(), self.conv).requires_grad_(True), RT[...,-1].detach().clone().requires_grad_(True)

        self.input_kwargs = {'right_hand_constraints': right_hand_constraints, 'left_hand_constraints': left_hand_constraints, 'device': device, 'MAX_ITERS': MAX_ITERS, 'min_lr': min_lr, 'chunk_size': chunk_size, 'loss_fn': loss_fn, 'optimize_shape': optimize_shape, 'optimize_hands': optimize_hands, 'K': K, 'RT': RT, 'OPTIMIZE_RT': OPTIMIZE_RT, 'MESH_INTERSECTION_LOSS_WEIGHT': MESH_INTERSECTION_LOSS_WEIGHT, 'lr': lr, 'OPTIMIZE_RT_FIX_CAM_0': OPTIMIZE_RT_FIX_CAM_0, 'SMOOTHNESS_TERMS': SMOOTHNESS_TERMS, 'PRINT_EVERY': PRINT_EVERY, 'SCHEDULER_PATIENCE': SCHEDULER_PATIENCE, 'are_control_points': are_control_points, 'NUM_POSES': NUM_POSES, 'is_smplfx': is_smplfx, 'GLOBAL_ORIENT_ZERO_X': GLOBAL_ORIENT_ZERO_X, 'GLOBAL_ORIENT_FIX_Z': GLOBAL_ORIENT_FIX_Z, 'body_zero_angles_idxs': body_zero_angles_idxs, 'REG_NORM_EXPRESSION_W': REG_NORM_EXPRESSION_W, 'BETA_REG_W': BETA_REG_W, 'USE_SIGMOID_ACTIVATION': USE_SIGMOID_ACTIVATION, 'spacing': spacing, 'APPLY_SMOOTHNESS_ON_CP': APPLY_SMOOTHNESS_ON_CP, 'apply_constraints_before_decode': apply_constraints_before_decode, 'ANGLES_REG_W': ANGLES_REG_W, 'ANGLES_REG_IDXS': ANGLES_REG_IDXS, 'smplx_kwargs': smplx_kwargs, 'OPTIMIZE_FLAME_SHAPE': OPTIMIZE_FLAME_SHAPE, 'body_constraints': body_constraints, 'is_single_smplx_rt': is_single_smplx_rt, 'OPTIMIZE_FACES': OPTIMIZE_FACES, 'optimize_smplx_rt': optimize_smplx_rt, 'IS_BODY_POSE_6D': IS_BODY_POSE_6D, 'OPTIMIZE_BODY_POSE': OPTIMIZE_BODY_POSE, 'DEPTH_LOSS_WEIGHT': DEPTH_LOSS_WEIGHT, 'Z_3D_LOSS_WEIGHT': Z_3D_LOSS_WEIGHT, 'max_dist_contacts': max_dist_contacts}
    
        self.smplx_kwargs = smplx_kwargs

        if LRs is None: LRs = {}
        LRs['lr'] = lr
        self.model = SMPLXWrapper(NUM_POSES, is_smplfx, smplx_kwargs, smplx_params_init, other_params_init, chunk_size, optimize_shape, optimize_hands, optimize_smplx_rt, are_control_points, device, GLOBAL_ORIENT_ZERO_X, GLOBAL_ORIENT_FIX_Z, right_hand_constraints, left_hand_constraints, body_zero_angles_idxs, LRs, USE_SIGMOID_ACTIVATION, spacing, apply_constraints_before_decode, OPTIMIZE_FLAME_SHAPE, body_constraints, is_single_smplx_rt, OPTIMIZE_FACES, IS_BODY_POSE_6D, OPTIMIZE_BODY_POSE)

        if MESH_INTERSECTION_LOSS_WEIGHT > 0:
            self.mesh_intersector = SegmentTriangleMeshIntersectionLossDualHands(torch.LongTensor(self.model.model.faces_original.copy().astype(np.int32)), num_triangles_test=6, target_remove_unlikely_parts=True, far_centroid_thr=0.03, is_intersect_prob=0.5, intersection_tolerance=2e-3, segmentation_path=segmentation_path, smplx_model_path=smplx_model_path).to(device)

        self.num_chunks = self.model.num_chunks
        if poses_2d is not None:
            assert poses_2d.ndim == 4 and poses_2d.shape[-1] == 2 # (N, C, J, 2)
            assert poses_2d_weights.ndim == 4 and poses_2d_weights.shape == poses_2d.shape[:-1] + (1,)
            if poses_2d_weights.max() > 1.0 or poses_2d_weights.min() < 0.0:
                print('WARNING: Weights are not normalized between 0 and 1, normalizing by max value', poses_2d_weights.max().item())

            self.poses_2d_chunks = chunkTensor(poses_2d, are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
            self.poses_2d_weights_chunks = chunkTensor(poses_2d_weights, are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
        else:
            self.poses_2d_chunks = None

        if poses_3d is not None:
            assert poses_3d.ndim == 3 and poses_3d.shape[-1] == 3 # (N, J, 3)
            self.poses_3d_chunks = chunkTensor(poses_3d, are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
        else:
            self.poses_3d_chunks = None

        from .utils_smplx import Contacts
        self.contacts = Contacts()
        if z_3d is not None:
            assert z_3d.ndim == 3 and z_3d.shape[-1] == 3 # (N, J, 3)
            # self.z_3d_chunks = chunkTensor(z_3d, are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
            z_3d_contacts = self.contacts.get_relative_contacts_joints(z_3d)
            z_3d_weights = getContactWeights(z_3d_contacts, max_dist_contact=self.input_kwargs['max_dist_contacts']).squeeze(-1)
            self.z_3d_chunks = chunkTensor(z_3d_contacts[...,2], are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
            self.z_3d_weights_chunks = chunkTensor(z_3d_weights, are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
        else:
            self.z_3d_chunks = None

        if depths is not None:
            assert depths.shape == poses_2d.shape[:-1]
            assert depths_mask.shape == poses_2d.shape[:-1]
            if depths_mask.max() > 1.0 or depths_mask.min() < 0.0:
                print('WARNING: Depth weights are not normalized between 0 and 1, normalizing by max value', depths_mask.max().item())
            
            depths_mask = depths_mask * poses_2d_weights.squeeze(-1).to(depths_mask.device) # apply keypoint weights to depth mask
            self.depth_chunks = chunkTensor(depths, are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
            self.depth_mask_chunks = chunkTensor(depths_mask, are_control_points, self.num_chunks, self.model.knots_per_chunk if are_control_points else None, spacing)
        else:
            self.depth_chunks, self.depth_mask_chunks = None, None

        self.is_temporal = False
        if keys is not None:
            self.is_temporal = True
            assert len(keys) <= NUM_POSES
            self.are_keys_consecutive = (np.max(keys[1:] - keys[:-1]) == 1).all()
            if are_control_points:
                assert self.are_keys_consecutive, "For control points based optimization, keys must be consecutive!"
                self.keys_chunked_vel, self.keys_chunked_acc, self.keys_chunked_jerk = [], [], []
                for i in range(self.num_chunks):
                    self.keys_chunked_vel.append(getTemporalIndicesBasedKeys(keys[i*self.model.knots_per_chunk*spacing:(i+1)*self.model.knots_per_chunk*spacing], 2))
                    self.keys_chunked_acc.append(getTemporalIndicesBasedKeys(keys[i*self.model.knots_per_chunk*spacing:(i+1)*self.model.knots_per_chunk*spacing], 3))
                    self.keys_chunked_jerk.append(getTemporalIndicesBasedKeys(keys[i*self.model.knots_per_chunk*spacing:(i+1)*self.model.knots_per_chunk*spacing], 4))
            else:
                self.keys_chunked_vel = createChunkedTemporalIndices(keys, NUM_POSES, self.num_chunks, diff=2)
                self.keys_chunked_acc = createChunkedTemporalIndices(keys, NUM_POSES, self.num_chunks, diff=3)
                self.keys_chunked_jerk = createChunkedTemporalIndices(keys, NUM_POSES, self.num_chunks, diff=4)

        self.angles_reg_idxs = []
        for idx in ANGLES_REG_IDXS:
            self.angles_reg_idxs += [idx*3, idx*3+1, idx*3+2]

    def get_optimizer(self):
        params_list = self.model.getOptimizableParamsList()
        if self.input_kwargs['OPTIMIZE_RT']:
            params_list.append({'params': [self.r_extrinsics, self.t_extrinsics], 'lr': self.input_kwargs['lr'] * 0.1, 'name': 'extrinsics'})
        return torch.optim.Adam(params_list)

    def getBestParams(self):
        smplx_params, other_params = self.model.getParams()
        if self.input_kwargs['OPTIMIZE_RT']:
            other_params['RTs'] = torch.cat([euler_angles_to_matrix(self.r_extrinsics, self.conv), self.t_extrinsics.unsqueeze(-1)], dim=-1).detach().clone()
        other_params['input_kwargs'] = self.input_kwargs
        return smplx_params, other_params

    def optimize(self):
        optimizer = self.get_optimizer()
        scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.8, patience=self.input_kwargs['SCHEDULER_PATIENCE'], min_lr=self.input_kwargs['min_lr'])
        best_loss, best_smplx_params, best_other_params = 1e10, None, None
        loss_fn = getLossFnc(self.input_kwargs['loss_fn'], reduction='none')
        trackers = {}
        if self.poses_2d_chunks is not None:
            trackers = {'body2d': 0.0, 'face2d': 0.0, 'hand2d': 0.0}
            if self.input_kwargs['DEPTH_LOSS_WEIGHT'] > 0.0:
                trackers['depth_loss'] = 0.0
        if self.input_kwargs['Z_3D_LOSS_WEIGHT'] > 0.0:
            trackers['z_3d_loss'] = 0.0
        if self.poses_3d_chunks is not None:
            trackers['pose3d'], trackers['face3d'], trackers['hand3d'] = 0.0, 0.0, 0.0
        if self.input_kwargs['REG_NORM_EXPRESSION_W'] > 0.0:
            trackers['exp_loss'] = 0.0
        if self.input_kwargs['MESH_INTERSECTION_LOSS_WEIGHT'] > 0:
            trackers['int_loss'], trackers['n_ints'] = 0.0, 0.0
        for triplet in self.input_kwargs['SMOOTHNESS_TERMS']:
            trackers[f'{triplet[0]}_{triplet[1]}_loss'] = 0.0
        if self.input_kwargs['ANGLES_REG_W'] > 0.0 and len(self.angles_reg_idxs) > 0:
            trackers['angles_reg_loss'] = 0.0
        time_start = time.time()
        saved_iter = 0
        optimizer_init_lrs = {param_group['name']: param_group['lr'] for param_group in optimizer.param_groups}
        for it in range(self.input_kwargs['MAX_ITERS']):
            for key in trackers.keys():
                trackers[key] = 0.0

            loss_it = 0.0
            for chunk_idx in range(self.num_chunks):
                optimizer.zero_grad()

                output, extra_losses, smplx_params_chunk = self.model(chunk_idx)

                joints = output.joints[:, self.mapper.smplx_idxs]
                loss = 0.0
                # plotMesh([output.vertices[50]], self.model.model.faces, title=f'Iteration {it} chunk {chunk_idx} mesh check', show=True, elev=-90, azim=-90)
                if self.poses_2d_chunks is not None:
                    if self.input_kwargs['OPTIMIZE_RT']:
                        proj_2d, depths = projectP(self.K.to(self.input_kwargs['device']) @ torch.cat([euler_angles_to_matrix(self.r_extrinsics, self.conv), self.t_extrinsics.unsqueeze(-1)], dim=-1).to(self.input_kwargs['device']), joints, with_depth=True)
                    else:
                        proj_2d, depths = projectP(self.P.to(self.input_kwargs['device']), joints, with_depth=True)

                    # if True:
                    if False:
                        from drawer import plotPoses2d
                        idx = 0
                        for cam_v in range(self.poses_2d_chunks[chunk_idx].shape[1]):
                            plotPoses2d([proj_2d[idx, cam_v], self.poses_2d_chunks[chunk_idx][idx, cam_v]], self.mapper.edges, title=f'2D projection check for cam {cam_v}', show=False, legend=['predicted', 'target'], mask=[np.ones(proj_2d.shape[2], dtype=bool), self.poses_2d_weights_chunks[chunk_idx][idx, cam_v].squeeze(-1).bool() > 0.0])
                            # plotPoses2d([torch.cat((proj_2d[idx, cam_v].cpu(), self.poses_2d_chunks[chunk_idx][idx, cam_v]))], getEdgesPairs(proj_2d.shape[-2]), title=f'2D projection check for cam {cam_v}', show=False, legend=['predicted', 'target'])
                            import matplotlib.pyplot as plt
                            plt.show()

                    loss2d = loss_fn(proj_2d, self.poses_2d_chunks[chunk_idx].to(self.input_kwargs['device']))
                    if self.poses_2d_weights_chunks is not None:
                        loss += (loss2d * self.poses_2d_weights_chunks[chunk_idx].to(self.input_kwargs['device'])).sum() / (self.poses_2d_weights_chunks[chunk_idx].to(self.input_kwargs['device']).sum() + 1e-9)
                    else:
                        loss += loss2d.mean()
                    
                    with torch.no_grad():
                        trackers['body2d'] += avgL2distance(proj_2d[:, :, self.mapper.body_subidxs], self.poses_2d_chunks[chunk_idx][:, :, self.mapper.body_subidxs].to(self.input_kwargs['device']), self.poses_2d_weights_chunks[chunk_idx][:, :, self.mapper.body_subidxs].to(self.input_kwargs['device'])).item()
                        trackers['hand2d'] += avgL2distance(proj_2d[:, :, self.mapper.hand_subidxs], self.poses_2d_chunks[chunk_idx][:, :, self.mapper.hand_subidxs].to(self.input_kwargs['device']), self.poses_2d_weights_chunks[chunk_idx][:, :, self.mapper.hand_subidxs].to(self.input_kwargs['device'])).item()
                        trackers['face2d'] += avgL2distance(proj_2d[:, :, self.mapper.face_subidxs], self.poses_2d_chunks[chunk_idx][:, :, self.mapper.face_subidxs].to(self.input_kwargs['device']), self.poses_2d_weights_chunks[chunk_idx][:, :, self.mapper.face_subidxs].to(self.input_kwargs['device'])).item()

                    if self.input_kwargs['DEPTH_LOSS_WEIGHT'] > 0:
                        # depth_loss = compute_ordinal_depth_loss(depths, self.depth_chunks[chunk_idx].to(depths.device), self.depth_mask_chunks[chunk_idx].to(depths.device), global_alignment=self.is_temporal)
                        depth_loss = compute_pearson_depth_loss(depths, self.depth_chunks[chunk_idx].to(depths.device), self.depth_mask_chunks[chunk_idx].to(depths.device))

                        loss += self.input_kwargs['DEPTH_LOSS_WEIGHT'] * depth_loss
                        trackers['depth_loss'] += depth_loss.detach().item()

                if self.poses_3d_chunks is not None:
                    loss3d = loss_fn(joints, self.poses_3d_chunks[chunk_idx].to(self.input_kwargs['device']))
                    if False:
                        from drawer import plotPoses
                        idx = min(100, joints.shape[0]-1)
                        plotPoses([joints[idx]], self.mapper.edges, title='3D check', show=True)
                        plotPoses([joints[idx], self.poses_3d_chunks[chunk_idx][idx]], self.mapper.edges, title='3D supervision check', show=False, legend=['predicted', 'target'])
                        plotPoses([torch.cat((joints[idx].cpu(), self.poses_3d_chunks[chunk_idx][idx]))], getEdgesPairs(len(joints[0])), title='3D supervision check', show=True, legend=['predicted', 'target'])
                    loss += loss3d.mean()
                    with torch.no_grad():
                        trackers['pose3d'] += avgL2distance(joints[:, self.mapper.body_subidxs], self.poses_3d_chunks[chunk_idx][:, self.mapper.body_subidxs].to(self.input_kwargs['device'])).item()
                        trackers['hand3d'] += avgL2distance(joints[:, self.mapper.hand_subidxs], self.poses_3d_chunks[chunk_idx][:, self.mapper.hand_subidxs].to(self.input_kwargs['device'])).item()
                        trackers['face3d'] += avgL2distance(joints[:, self.mapper.face_subidxs], self.poses_3d_chunks[chunk_idx][:, self.mapper.face_subidxs].to(self.input_kwargs['device'])).item()

                if self.input_kwargs['Z_3D_LOSS_WEIGHT'] > 0.0:
                    rel_contacts = self.contacts.get_relative_contacts_joints(output.joints)
                    rel_weights = getContactWeights(rel_contacts, max_dist_contact=self.input_kwargs['max_dist_contacts']).squeeze(-1)
                    rel_weights_combined = rel_weights + self.z_3d_weights_chunks[chunk_idx].to(rel_weights.device) # OR
                    loss_z_3d = ((rel_contacts[...,2] - self.z_3d_chunks[chunk_idx].to(joints.device)).abs() * rel_weights_combined).sum() / (1e-9 + rel_weights_combined.sum())
                    loss += self.input_kwargs['Z_3D_LOSS_WEIGHT'] * loss_z_3d
                    trackers['z_3d_loss'] += loss_z_3d.detach().item()

                for loss_k, loss_v in extra_losses.items():
                    for loss_k_n, loss_v_n in loss_v.items():
                        loss += loss_v_n
                        trackers[f"{loss_k}_{loss_k_n}"] += loss_v_n.detach().item()

                if self.input_kwargs['MESH_INTERSECTION_LOSS_WEIGHT'] > 0:
                    inpenetration_err_, num_total_penetrations = self.mesh_intersector(output, vis=False)#it > 100)
                    loss += self.input_kwargs['MESH_INTERSECTION_LOSS_WEIGHT'] * inpenetration_err_
                    trackers['int_loss'] += (1e3 * inpenetration_err_.detach().item()) # scale 1e3 times
                    trackers['n_ints'] += num_total_penetrations

                if len(self.input_kwargs['SMOOTHNESS_TERMS']) > 0:
                    for triplet in self.input_kwargs['SMOOTHNESS_TERMS']:
                        if self.are_keys_consecutive:
                            if triplet[0] == 'vertices':
                                smooth_loss = self.smoothness_losses[triplet[1]](output.vertices, with_stats=False)[0].mean()
                            else:
                                if self.input_kwargs['are_control_points'] and self.input_kwargs['APPLY_SMOOTHNESS_ON_CP']:
                                    smooth_loss = self.smoothness_losses[triplet[1]](self.model.optimized_params[triplet[0]].chunks[chunk_idx].squeeze(0).T, with_stats=False)[0].mean()
                                else:
                                    smooth_loss = self.smoothness_losses[triplet[1]](smplx_params_chunk[triplet[0]], with_stats=False)[0].mean()
                        else:
                            if triplet[0] == 'vertices':
                                if triplet[1] == 'velocity':
                                    smooth_loss = torch.linalg.norm(output.vertices[self.keys_chunked_vel[chunk_idx][:,1]] - output.vertices[self.keys_chunked_vel[chunk_idx][:,0]], dim=-1).mean()
                                elif triplet[1] == 'acceleration':
                                    smooth_loss = torch.linalg.norm(output.vertices[self.keys_chunked_acc[chunk_idx][:,2]] - 2 * output.vertices[self.keys_chunked_acc[chunk_idx][:,1]] + output.vertices[self.keys_chunked_acc[chunk_idx][:,0]], dim=-1).mean()
                                elif triplet[1] == 'jerk':
                                    smooth_loss = torch.linalg.norm(output.vertices[self.keys_chunked_jerk[chunk_idx][:,3]] - 3 * output.vertices[self.keys_chunked_jerk[chunk_idx][:,2]] + 3 * output.vertices[self.keys_chunked_jerk[chunk_idx][:,1]] - output.vertices[self.keys_chunked_jerk[chunk_idx][:,0]], dim=-1).mean()
                                else:
                                    assert False, "Invalid smoothness type in triplet: " + str(triplet)
                            else:
                                assert False, "THE CODE MUST BE REVISED:"
                                if self.input_kwargs['are_control_points'] and self.input_kwargs['APPLY_SMOOTHNESS_ON_CP']:
                                    smooth_loss = self.smoothness_losses[triplet[1]](self.model.optimized_params[triplet[0]].chunks[chunk_idx].squeeze(0).T, with_stats=False)[0].mean()
                                else:
                                    smooth_loss = self.smoothness_losses[triplet[1]](smplx_params_chunk[triplet[0]][self.keys_chunked[chunk_idx]], with_stats=False)[0].mean()
                        loss += triplet[2] * smooth_loss
                        trackers[f'{triplet[0]}_{triplet[1]}_loss'] += (1e3 * smooth_loss.detach().item()) # scale 1e3 times

                if self.input_kwargs['REG_NORM_EXPRESSION_W'] > 0.0:
                    loss_expression = torch.linalg.norm(smplx_params_chunk['expression'], dim=-1).mean()
                    loss += self.input_kwargs['REG_NORM_EXPRESSION_W'] * loss_expression
                    trackers['exp_loss'] += loss_expression.detach().item()

                if self.input_kwargs['BETA_REG_W'] > 0.0:
                    if self.input_kwargs['optimize_shape']:
                        loss += self.input_kwargs['BETA_REG_W'] * torch.linalg.norm(smplx_params_chunk['betas'], dim=-1).mean()
                    if self.input_kwargs['OPTIMIZE_FLAME_SHAPE'] and self.input_kwargs['is_smplfx']:
                        loss += self.input_kwargs['BETA_REG_W'] * torch.linalg.norm(smplx_params_chunk['shape'], dim=-1).mean()

                if self.input_kwargs['ANGLES_REG_W'] > 0.0 and len(self.angles_reg_idxs) > 0:
                    loss_angles_reg = torch.abs(smplx_params_chunk['body_pose'][:, self.angles_reg_idxs]).mean()
                    loss += self.input_kwargs['ANGLES_REG_W'] * loss_angles_reg
                    trackers['angles_reg_loss'] += loss_angles_reg.detach().item()

                loss.backward()

                if self.input_kwargs['OPTIMIZE_RT'] and self.input_kwargs.get('opt_rt_extrinsics_fix_first_view', False):
                    self.r_extrinsics.grad[0].zero_()
                    self.t_extrinsics.grad[0].zero_()

                optimizer.step()
                scheduler.step(loss)
                for optimizer_param_group in optimizer.param_groups:
                    optimizer_param_group['lr'] = min(optimizer_param_group['lr'], optimizer_init_lrs[optimizer_param_group['name']])

                loss_it += loss.detach().item()

                self.model.apply_constraints(chunk_idx)
    
            if it % self.input_kwargs['PRINT_EVERY'] == 0:
                log_str = f"Iter {it}, Loss {loss_it:.3f}, time {(time.time() - time_start)/60:.2f}m, lr {optimizer.param_groups[0]['lr']:.2e}"
                for key in trackers.keys():
                    log_str += f", {key} {trackers[key]/(chunk_idx+1):.3f}"
                if self.input_kwargs['optimize_shape']:
                    log_str += f' beta {torch.linalg.norm(self.model.smplx_params["betas"]).item():.3f}'
                if self.input_kwargs['is_smplfx'] and self.input_kwargs['OPTIMIZE_FLAME_SHAPE']:
                    log_str += f' shape {torch.linalg.norm(self.model.smplx_params["shape"]).item():.3f}'
                print(log_str)

            if best_loss > loss_it and it > 5: # never save on the first few iterations, as the loss can be very unstable at the beginning
                best_loss = loss_it
                saved_iter = it
                # print('SAVE BEST PARAMS at iter', it, 'with loss', best_loss)
                best_smplx_params, best_other_params = self.getBestParams()

            if optimizer.param_groups[0]['lr'] <= self.input_kwargs['min_lr']:
                print(f"Stopping optimization at iteration {it} due to learning rate reaching minimum.")
                break

        # torch.save({'smplx': best_smplx_params, 'others': best_other_params}, '')
        return best_smplx_params, best_other_params
