import sys, os
sys.path.append(os.path.join(os.path.abspath(os.path.dirname(__file__)), '..'))

import torch
from mmengine.registry import MODELS, METRICS
from mmengine.runner import Runner
from mmengine.config import Config
import torch.nn.functional as F
from src.utils.motion import recover_from_local_position_fast, recover_from_local_rotation
from tqdm import tqdm
import numpy as np

def vis_motion(motion, mean, std):
    from src.metrics.recons_metric import recover_from_local_position
    from src.utils.visualize import vis_3d_motion
    motion = (motion * std + mean).numpy()
    positions = recover_from_local_position(motion, 22) # 
    vis_3d_motion(positions[None], None, [f'video.mp4'], fps=30)

cfg = Config.fromfile('configs/eval_text.py')
ckp = torch.load('work_dirs/anymo_stage_3_3B_210e_48gpu/e210.pth')['state_dict']
lr = 2e-4
opt_iter = 400

# build model
model = MODELS.build(cfg.model)
model.load_state_dict(ckp)
model.eval().cuda()

# build dataloader
cfg.test_dataloader.batch_size = 2
dataloader = Runner.build_dataloader(cfg.test_dataloader)

# build evaluator
evaluator = METRICS.build(cfg.test_evaluator)

for data in tqdm(dataloader):
    # prepare data
    # import ipdb; ipdb.set_trace()
    data = {k: v.cuda() if isinstance(v, torch.Tensor) else v for k, v in data.items()}
    m_lens = data['motion_length']
    caps = data['text_data']

    _, _, ids = model(**data, mode='predict', return_ids=True)
    # import ipdb; ipdb.set_trace()
    ids_mask = (ids > -1).unsqueeze(1)
    ids[ids < 0] = 0 
    emb = model.vqvae.quantizer.dequantize(ids.permute(1, 2, 0)) # B, T, C
    emb = emb.detach().clone().permute(0, 2, 1) # B, C, T
    emb = emb * ids_mask[0][0].unsqueeze(1)

    emb.requires_grad = True
    optimizer = torch.optim.AdamW([emb], lr=lr, betas=(0.5, 0.9), weight_decay=1e-6)
    m_mask = torch.arange(max(m_lens)).cuda().expand(len(m_lens), max(m_lens)) < m_lens.unsqueeze(1)
    pbar = tqdm(total=opt_iter)
    for i in range(opt_iter):
        pred_motion = model.vqvae.decoder(emb, ids_mask[0].squeeze(0).sum(1)).permute(0, 2, 1)
        pred_traj = recover_from_local_position_fast(pred_motion * data['std'] + data['mean'], njoint=22)[:, :, 0]
        gt_traj = recover_from_local_position_fast(data['motion'] * data['std'] + data['mean'], njoint=22)[:, :, 0]

        tta_loss = F.mse_loss(pred_traj, gt_traj, reduction='none')
        tta_loss = (tta_loss * m_mask.unsqueeze(-1)).sum() / (m_mask.sum() * tta_loss.shape[-1])

        optimizer.zero_grad()
        tta_loss.backward()
        optimizer.step()
        pbar.update(1)
        pbar.set_description(f'tta_loss: {tta_loss.detach().cpu().item():.5f}')

    pred_motions = model.vqvae.decoder(emb, ids_mask[0].squeeze(0).sum(1)).permute(0, 2, 1)
    pred_motions = [pred_motions[i][:m_lens[i]] for i in range(pred_motions.shape[0])]

    import ipdb; ipdb.set_trace()
    # idx = 25;pred_data = recover_from_local_rotation((pred_motions[idx].cpu() * data['std'].cpu() + data['mean'].cpu()).detach().cpu().numpy()[:m_lens[idx]], njoint=22); np.savez('tmp.npz', **pred_data)
    # idx = 25;pred_data = recover_from_local_rotation((data['motion'].cpu() * data['std'].cpu() + data['mean'].cpu()).cpu().numpy()[idx][:m_lens[idx]], njoint=22); np.savez('tmp.npz', **pred_data)
    # idx = 25;vis_motion(pred_motions[idx].detach().cpu(), data['mean'].cpu(), data['std'].cpu())
    # idx = 25;vis_motion(data['motion'][idx][:m_lens[idx]].cpu(), data['mean'].cpu(), data['std'].cpu())

    # evaluator.process(data, (pred_motions, m_lens))
    # print(evaluator.compute_metrics(evaluator.results))