import os

os.environ["CUDA_VISIBLE_DEVICES"] = '7'

import numpy as np
from smpl_torch import SMPL as SMPL_torch
from smpl_tf import SMPL as SMPL_tf
from utils import plot_mesh, plot_mesh_2
from time import time

import torch
import tensorflow as tf

if __name__ == '__main__':
    device = torch.device('cuda')

    smpl_tf : SMPL_tf = SMPL_tf('model_files')
    smpl_torch : SMPL_torch = SMPL_torch(device=device, path='model_files')

    np.random.seed(4444)
    pose = ((np.random.rand(4, 72) - 0.5) * 0.4).astype(np.float32)
    betas = ((np.random.rand(4, 10) - 0.5) * 0.06).astype(np.float32)
    betas = np.concatenate((betas, np.array([0,1,0,1], dtype=np.float32)[:,None]), axis=1)

    # pose = np.zeros((4, 72), dtype=np.float32)
    # betas = np.zeros((4, 11), dtype=np.float32)
    
    pose_torch = torch.tensor(pose, device=device)
    betas_torch = torch.tensor(betas, device=device)

    result_torch, joints_torch, _, _ = smpl_torch(betas_torch, pose_torch)
    s = time()
    result_torch, joints_torch, v_shaped_1, v_posed_1 = smpl_torch(betas_torch, pose_torch)
    print(time() - s)
    R_1 = smpl_torch.rodrigues(pose_torch.view(-1, 1, 3)).reshape(4, -1, 3, 3)

    # plot_mesh(result[0].cpu().numpy(), model.f).show()

    R = smpl_tf.run_rodrigues(pose)
    result_tf, joints_tf, _, _ = smpl_tf.run(R, betas)
    s = time()
    R = smpl_tf.run_rodrigues(pose)
    result_tf, joints_tf, v_shaped_2, v_posed_2 = smpl_tf.run(R, betas)
    print(time() - s)

    result_torch = result_torch.cpu().numpy()
    result_tf = result_tf.numpy()
    diff = np.abs(result_torch - result_tf).sum(-1).sum(-1)

    plot_mesh_2(result_torch[0], result_tf[0], smpl_tf.f).show()
    plot_mesh_2(result_torch[1], result_tf[1], smpl_tf.f).show()
    pass