import argparse
import os
from os.path import join as pjoin

import numpy as np
import torch
import utils.paramUtil as paramUtil
from models import MotionTransformer
from torch.utils.data import DataLoader
from trainers import DDPMTrainer
from utils.get_opt import get_opt
from utils.motion_process import recover_from_ric
from utils.plot_script import *
from utils.utils import *
from utils.word_vectorizer import POS_enumerator

# def plot_t2m(opt, data, result_path, caption):
#     joint = recover_from_ric(torch.from_numpy(data).float(), opt.joints_num).numpy()
#     # joint = motion_temporal_filter(joint, sigma=1)
#     plot_3d_motion(result_path, paramUtil.t2m_kinematic_chain, joint, title=caption, fps=20)


# def process(trainer, opt, device, mean, std, text, motion_length, result_path):

#     result_dict = {}
#     with torch.no_grad():
#         if motion_length != -1:
#             caption = [text]
#             m_lens = torch.LongTensor([motion_length]).to(device)
#             pred_motions = trainer.generate(caption, m_lens, opt.dim_pose)
#             motion = pred_motions[0].cpu().numpy()
#             motion = motion * std + mean
#             title = text + " #%d" % motion.shape[0]
#             plot_t2m(opt, motion, result_path, title

def plot_t2m(data, result_path, npy_path, caption, joints_n):
    joint = recover_from_ric(torch.from_numpy(data).float(), joints_n).numpy()

def get_numpy_file_path(prompt, epoch, n_frames):
    # e.g. "airplane_fly_1_1000_60f.npy"
    prompt_no_spaces = prompt.replace(' ', '_')
    return f"{prompt_no_spaces}_{epoch}_{n_frames}f"

def get_wordvec_model(opt):
    encoder = MotionTransformer(
        input_feats=opt.dim_pose,
        num_frames=opt.max_motion_length,
        num_layers=opt.num_layers,
        latent_dim=opt.latent_dim,
        no_clip=opt.no_clip,
        no_eff=opt.no_eff)
    return encoder


if __name__ == '__main__':
    print("inferencemake started")
    parser = argparse.ArgumentParser()
    parser.add_argument('--opt_path', type=str, help='Opt path')
    parser.add_argument('--text', type=str, default="", help='Text description for motion generation')
    parser.add_argument('--motion_length', type=int, default=60, help='Number of frames for motion generation')
    parser.add_argument('--result_path', type=str, default="test_sample.gif", help='Path to save generation result')
    parser.add_argument('--npy_path', type=str, default="", help='Path to save 3D keypoints sequence')
    parser.add_argument('--gpu_id', type=int, default=-1, help="which gpu to use")
    parser.add_argument('--seed', type=int, default=0, help="random seed")
    # add which_epoch
    parser.add_argument('--which_epoch', type=str, default="latest", help="which epoch to load")
    args = parser.parse_args()
    
    set_random_seed(args.seed)
    print(f"set random seed to {args.seed}")
    device = torch.device('cuda:%d' % args.gpu_id if args.gpu_id != -1 else 'cpu')
    opt = get_opt(args.opt_path, device)
    opt.do_denoise = True
    opt.which_epoch = args.which_epoch

    # TODO (elmc): re-enable this
    # assert opt.dataset_name == "t2m"
    # assert args.motion_length <= 196
    # opt.data_root = './dataset/HumanML3D'
    opt.data_root = './data/GRAB'
    # opt.motion_dir = pjoin(opt.data_root, 'new_joint_vecs')
    opt.text_dir = pjoin(opt.data_root, 'texts')
    # TODO (elmc): re-enable this
    # opt.joints_num = 22
    # opt.dim_pose = 263
    opt.dim_pose = 212
    dim_word = 300
    dim_pos_ohot = len(POS_enumerator)
    # TODO (elmc): re-enable this
    # num_classes = 200 // opt.unit_length

    mean = np.load(pjoin(opt.meta_dir, 'mean.npy'))
    std = np.load(pjoin(opt.meta_dir, 'std.npy'))

    print("Loading word vectorizer...")
    encoder = get_wordvec_model(opt).to(device)
    print("Loading model...")
    trainer = DDPMTrainer(opt, encoder)
    trainer.load(pjoin(opt.model_dir, opt.which_epoch + '.tar'))

    trainer.eval_mode()
    trainer.to(opt.device)

    result_dict = {}
    with torch.no_grad():
        if args.motion_length != -1:
            caption = [args.text]
            file_name = f"{opt.which_epoch}_{args.motion_length}f.npy"
            m_lens = torch.LongTensor([args.motion_length]).to(device)
            pred_motions = trainer.generate(caption, m_lens, opt.dim_pose)
            motion = pred_motions[0].cpu().numpy()
            motion = motion * std + mean
            title = args.text + " #%d" % motion.shape[0]
            print(f"trying to plot {title}")
            # write motion to numpy file
            if not os.path.exists(args.npy_path):
                os.makedirs(args.npy_path) 
            full_npy_path = f"{args.npy_path}/{get_numpy_file_path(caption[0], opt.which_epoch, args.motion_length)}.npy"
            with open(full_npy_path, 'wb') as f:
                print(f"saving output to {full_npy_path}")
                np.save(f, motion)

            # plot_t2m(motion, args.result_path, args.npy_path, title)