import sys
sys.path.append('/mnt/shenzhen2cephfs/capybarali/codes/humanoid')

import torch, yaml, os
from tqdm import tqdm
from smplx import SMPLX
from src.utils.rotation_conversions import quaternion_to_matrix, matrix_to_rotation_6d, axis_angle_to_6d
import numpy as np
from argparse import ArgumentParser
import joblib
from data.vis import vis_3d_motion
from data.vis_g1 import vis_3d_g1
from copy import deepcopy

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

parser = ArgumentParser(description="Launch MoCap processing")
parser.add_argument('--save_root', type=str, default="data/smplx_data")
parser.add_argument('--start_idx', type=int, default=0)
parser.add_argument('--interval', type=int, default=1)
args = parser.parse_args()

os.makedirs(args.save_root, exist_ok=True)   
os.makedirs(os.path.join(args.save_root, 'motions'), exist_ok=True)   

smplx_model = SMPLX(
    model_path='checkpoints/human_model/SMPLX_NEUTRAL.npz', 
    use_pca=False, num_expression_coeffs=100, num_betas=10, ext='npz'
).to(device)

def get_smplx_motion(data_path):
    data = dict(np.load(data_path))
    import ipdb; ipdb.set_trace()
    smplx_data = dict(
        smplx_transl=data["transl"],
        smplx_global_orient=data["global_orient"],
        smplx_body_pose=data["body_pose"],
    )

    transl = torch.from_numpy(smplx_data['smplx_transl'])
    betas = torch.from_numpy(np.load('checkpoints/humanoid_model/g1/betas.npy'))
    global_orient = torch.from_numpy(smplx_data['smplx_global_orient'])
    body_pose = torch.from_numpy(smplx_data['smplx_body_pose'])

    N = transl.shape[0]
    motion_params = dict(
        transl=transl, 
        global_orient=global_orient, 
        body_pose=body_pose, 
        betas=betas.unsqueeze(0).repeat(N, 1).float()
    )
    # 1. process positions
    frame_params = {k: v.to(device) for k, v in motion_params.items()}
    frame_params['leye_pose'] = torch.zeros((N, 3)).to(device)
    frame_params['reye_pose'] = torch.zeros((N, 3)).to(device)
    frame_params['left_hand_pose'] = torch.zeros((N, 45)).to(device)
    frame_params['right_hand_pose'] = torch.zeros((N, 45)).to(device)
    frame_params['jaw_pose'] = torch.zeros((N, 3)).to(device)
    frame_params['expression'] = torch.zeros((N, 100)).to(device)
    output = smplx_model(**frame_params)
    position_data = output.joints.detach().cpu()[:, :22] # T, 22 ,3
    position_val_data = position_data[1:] - position_data[:-1]

    root_idx = 0
    # put on floor and put root on origin for the first frame
    # import ipdb; ipdb.set_trace()
    ori = deepcopy(position_data[0, root_idx]) # first frame root position
    y_min = torch.min(position_data[:, :, 1])
    ori[1] = y_min
    position_data = position_data - ori
    velocities_root = position_data[1:, root_idx, :] - position_data[:-1, root_idx, :]
    # import ipdb; ipdb.set_trace()
    position_data_cp = deepcopy(position_data)
    position_data[:,:,0] -= position_data_cp[:,0:1,0]
    position_data[:,:,2] -= position_data_cp[:,0:1,2]

    T, njoint, _ = position_data.shape
    final_x = torch.zeros((T, 2 + 6 + njoint * 3 + njoint * 3))
    final_x[1:, 0] = velocities_root[:, 0]
    final_x[1:, 1] = velocities_root[:, 1]
    final_x[:, 2:2+6] = axis_angle_to_6d(global_orient)
    final_x[:, 8:8+njoint*3] = position_data.flatten(1, 2)
    final_x[1:, 8+njoint*3:8+njoint*6] = position_val_data.flatten(1, 2) # T, 140

    # import ipdb; ipdb.set_trace()
    # vis_3d_motion([position_data.numpy()], None, ['video.mp4'], fps=30)
    return final_x

# python -m data.motionmillion.tools.process --save_root "data/motionmillion/final_data"
if __name__ == '__main__':
    with open('data/smplx_path.txt', 'r') as f:
        paths = f.readlines()

    for line in tqdm(paths[args.start_idx::args.interval]):
        data_path = line.strip()
        save_path = data_path.replace('.npz', '.npy').replace('/mnt/shenzhen2cephfs/capybarali/codes/neobot/data/motionmillion/1_filtered_smplx/', '')
        save_path = args.save_root + '/' + save_path
        if os.path.exists(save_path):
            continue
        smplx_motion = get_smplx_motion(data_path)
        os.makedirs(os.path.dirname(save_path), exist_ok=True)
        np.save(save_path, smplx_motion)
        # joblib.dump(data, save_path)