# This script is borrowed and extended from https://github.com/nkolot/SPIN/blob/master/models/hmr.py
# Adhere to their licence to use this script

import torch
import numpy as np
import os.path as osp
from smplx import SMPL as _SMPL
from smplx.utils import ModelOutput, SMPLOutput
import warnings
warnings.filterwarnings("ignore", category=DeprecationWarning)

# Map joints to SMPL joints
JOINT_MAP = {
    'OP Nose': 24, 'OP Neck': 12, 'OP RShoulder': 17,
    'OP RElbow': 19, 'OP RWrist': 21, 'OP LShoulder': 16,
    'OP LElbow': 18, 'OP LWrist': 20, 'OP MidHip': 0,
    'OP RHip': 2, 'OP RKnee': 5, 'OP RAnkle': 8,
    'OP LHip': 1, 'OP LKnee': 4, 'OP LAnkle': 7,
    'OP REye': 25, 'OP LEye': 26, 'OP REar': 27,
    'OP LEar': 28, 'OP LBigToe': 29, 'OP LSmallToe': 30,
    'OP LHeel': 31, 'OP RBigToe': 32, 'OP RSmallToe': 33, 'OP RHeel': 34,
    'Right Ankle': 8, 'Right Knee': 5, 'Right Hip': 45,
    'Left Hip': 46, 'Left Knee': 4, 'Left Ankle': 7,
    'Right Wrist': 21, 'Right Elbow': 19, 'Right Shoulder': 17,
    'Left Shoulder': 16, 'Left Elbow': 18, 'Left Wrist': 20,
    'Neck (LSP)': 47, 'Top of Head (LSP)': 48,
    'Pelvis (MPII)': 49, 'Thorax (MPII)': 50,
    'Spine (H36M)': 51, 'Jaw (H36M)': 52,
    'Head (H36M)': 53, 'Nose': 24, 'Left Eye': 26,
    'Right Eye': 25, 'Left Ear': 28, 'Right Ear': 27
}

JOINT_PAIRS = [('OP Neck', 'OP MidHip'), ('OP RHip', 'OP RKnee'), ('OP RKnee', 'OP RAnkle'), ('OP MidHip', 'OP RHip'), ('OP MidHip', 'OP LHip'), ('OP LHip', 'OP LKnee'),
              ('OP LKnee', 'OP LAnkle'), ('OP Neck', 'OP RShoulder'), ('OP RShoulder', 'OP RElbow'), ('OP RElbow', 'OP RWrist'), ('OP Neck', 'OP LShoulder'),
              ('OP LShoulder', 'OP LElbow'), ('OP LElbow', 'OP LWrist'), ('OP Neck', 'OP Nose'), ('OP Nose', 'OP REye'), ('OP Nose', 'OP LEye'),
              ('OP REye', 'OP REar'), ('OP LEye', 'OP LEar'), ('OP LAnkle', 'OP LBigToe'), ('OP RAnkle', 'OP RBigToe')]


OP_21_JOINTS = [
    'OP LWrist', 'OP RWrist', 'OP LAnkle', 'OP RAnkle',
    'OP LElbow', 'OP RElbow', 'OP LKnee', 'OP RKnee',
    'OP LShoulder','OP RShoulder',
    'OP LHip',  'OP MidHip','OP RHip',
    'OP Nose', 'OP Neck',
    'OP LEye', 'OP REye',  'OP LEar', 'OP REar',
    'OP LBigToe', 'OP RBigToe',

]


JOINT_IDS = {OP_21_JOINTS[i]: i for i in range(len(OP_21_JOINTS))}

JOINT_ID_PAIRS = [(JOINT_IDS[first], JOINT_IDS[second]) for first, second in JOINT_PAIRS]

# JOINT_REGRESSOR_TRAIN_EXTRA = osp.join(VIBE_DATA_DIR, 'J_regressor_extra.npy')
# SMPL_MEAN_PARAMS = osp.join(VIBE_DATA_DIR, 'smpl_mean_params.npz')
SMPL_MODEL_DIR = 'smpl/data/smpl_data'
H36M_TO_J17 = [6, 5, 4, 1, 2, 3, 16, 15, 14, 11, 12, 13, 8, 10, 0, 7, 9]
H36M_TO_J14 = H36M_TO_J17[:14]


class SMPL(_SMPL):
    """ Extension of the official SMPL implementation to support more joints """

    def __init__(self, *args, **kwargs):
        super(SMPL, self).__init__(*args, **kwargs)
        joints = [JOINT_MAP[i] for i in OP_21_JOINTS]
        # J_regressor_extra = np.load(JOINT_REGRESSOR_TRAIN_EXTRA)
        # self.register_buffer('J_regressor_extra', torch.tensor(J_regressor_extra, dtype=torch.float32))
        self.joint_map = torch.tensor(joints, dtype=torch.long)

    def forward(self, *args, **kwargs):
        kwargs['get_skin'] = True
        smpl_output = super(SMPL, self).forward(*args, **kwargs)
        # extra_joints = vertices2joints(self.J_regressor_extra, smpl_output.vertices)
        # joints = torch.cat([smpl_output.joints, extra_joints], dim=1)
        # joints = joints[:, self.joint_map, :]
        joints = smpl_output.joints
        output = SMPLOutput(vertices=smpl_output.vertices,
                            global_orient=smpl_output.global_orient,
                            body_pose=smpl_output.body_pose,
                            joints=joints,
                            betas=smpl_output.betas,
                            full_pose=smpl_output.full_pose)
        return output


def get_smpl_faces():
    smpl = SMPL(SMPL_MODEL_DIR, batch_size=1, create_transl=False)
    return smpl.faces