import torch
import numpy as np
import os 

from .mesh.io import save_obj, to_mesh
from .mesh.smpl2mesh import SMPL2Mesh
from .skeleton import SkeletonAMASS, convert2humanml
from .skeleton2smpl.skeleton2smpl import Skeleton2Obj
import json


def save_mesh(vertices, faces, npy_file):
    def npy_path_to_obj_path(npy_path: str) -> str:
        return os.path.join(os.path.dirname(npy_path) ,  f"{npy_path}_obj")
    results_dir = npy_path_to_obj_path(npy_file)
    os.makedirs(results_dir, exist_ok=True)
    for frame_i in range(vertices.shape[-1]):
        file_path = os.path.join(results_dir, f"frame{frame_i:03d}.obj")
        mesh = to_mesh(vertices[..., frame_i], faces)
        save_obj(mesh, file_path)
    print(f"Saved obj files to [{results_dir}]")

def main():
    num_smplify_iters = 20 # This is what requires most time. It can be decreased or increasd depending on the output quality we want (or how quick we canr each it)
    device = "cuda"

    # get observation smpl params
    json_file_path = "./smpl_params.json"
    with open(json_file_path, "r") as json_file:
        loaded_data = json.load(json_file)
    person_idx = 0
    smpl_dict_last_obs = loaded_data[-1]        
    smpl_dict_last_obs = {k: torch.from_numpy(np.array(v)).float().to(device) for k,v in smpl_dict_last_obs.items()}

    input_kpts = smpl_dict_last_obs['joints3d'] 
    input_kpts = torch.stack([input_kpts[..., 0], input_kpts[..., 2], -input_kpts[..., 1]], dim=-1)
    input_kpts = input_kpts/1000
    input_kpts = input_kpts - input_kpts[..., 0:1, :]
    
    # get predictions
    pred_motions = torch.from_numpy(np.load( "./joints3d.npy", allow_pickle=True)).to(device) 


    pred_motions = torch.stack([pred_motions[..., 0], pred_motions[..., 2], -pred_motions[..., 1]], dim=-1)

    
    # remove bacth dimension, add a zero hip joint
    pred_motions = pred_motions.squeeze(0)
    pred_motions = torch.cat([torch.zeros(*pred_motions.shape[:2], 1, 3).to(device), pred_motions], dim=-2)

    # select just some of the motions
    # TO DO use the previous code with the limb length variance error to choose the sample
    # Or pick the most diverse
    pred_motions = pred_motions[:1]
    pred_motions = pred_motions.view(-1, 22, 3)

    skeleton = SkeletonAMASS
    pred_motions = convert2humanml(pred_motions, skeleton.LANDMARKS, skeleton.TO_HUMANML_NAMES)

    init_params = {}
    init_params["betas"] = smpl_dict_last_obs["betas"][person_idx].unsqueeze(0).expand(pred_motions.shape[0], -1)
    init_params["pose"] = smpl_dict_last_obs["body_pose"][person_idx].view(-1, 3)
    assert init_params["pose"].shape[0] == 24, "the body pose should have 24 joints, it is the output of NLF"

    init_params["pose"] = init_params["pose"].unsqueeze(0).expand(pred_motions.shape[0], -1, -1).view(pred_motions.shape[0], -1).to(device)
    init_params["cam"] = smpl_dict_last_obs["transl"][person_idx].unsqueeze(0).unsqueeze(-2).expand(pred_motions.shape[0], -1, -1).to(device)

    skeleton2obj = Skeleton2Obj(
        device=device, num_smplify_iters=num_smplify_iters, 
        smpl_model_dir="./body_models/", #path to smpl body models
        gmm_model_dir="./joint2smpl_models/", #path to gmm model
    )
    rot_motions, smpl_dict = skeleton2obj.convert_motion_2smpl(pred_motions, hmp=True, init_params=init_params, fix_betas=True)

    smpl2mesh = SMPL2Mesh(device)
    vertices, faces = smpl2mesh.convert_smpl_to_mesh(rot_motions, pred_motions)

    pred_files = [('./hanyu')]
    vertices = vertices.reshape(*vertices.shape[:2], len(pred_files), -1)
    for v, npy_file in zip(np.moveaxis(vertices, 2, 0), pred_files):
        save_mesh(v, faces, npy_file)


if __name__ == "__main__":
    main()