from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import torch

from ..utils.normalizer import Normalizer
from ...models import build_body_model
from ...utils import transforms
from ...utils.kp_utils import root_centering
from ...utils.imutils import compute_cam_intrinsics

KEYPOINTS_THR = 0.3

def convert_dpvo_to_cam_angvel(traj, fps):
    """Function to convert DPVO trajectory output to camera angular velocity"""
    
    # 0 ~ 3: translation, 3 ~ 7: Quaternion
    quat = traj[:, 3:]
    
    # Convert (x,y,z,q) to (q,x,y,z)
    quat = quat[:, [3, 0, 1, 2]]
    
    # Quat is camera to world transformation. Convert it to world to camera
    world2cam = transforms.quaternion_to_matrix(torch.from_numpy(quat)).float()
    R = world2cam.mT
    
    # Compute the rotational changes over time.
    cam_angvel = transforms.matrix_to_axis_angle(R[:-1] @ R[1:].transpose(-1, -2))
    
    # Convert matrix to 6D representation
    cam_angvel = transforms.matrix_to_rotation_6d(transforms.axis_angle_to_matrix(cam_angvel))
    
    # Normalize 6D angular velocity
    cam_angvel = cam_angvel - torch.tensor([[1, 0, 0, 0, 1, 0]]).to(cam_angvel) # Normalize
    cam_angvel = cam_angvel * fps
    cam_angvel = torch.cat((cam_angvel, cam_angvel[:1]), dim=0)
    return cam_angvel


class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, cfg, tracking_results, slam_results, width, height, fps):
        
        self.tracking_results = tracking_results
        self.slam_results = slam_results
        self.width = width
        self.height = height
        self.fps = fps
        self.res = torch.tensor([width, height]).float()
        self.intrinsics = compute_cam_intrinsics(self.res)
        
        self.device = cfg.DEVICE.lower()
        
        self.smpl = build_body_model('cpu')
        self.keypoints_normalizer = Normalizer(cfg)
        
        self._to = lambda x: x.unsqueeze(0).to(self.device)
        
    def __len__(self):
        return len(self.tracking_results.keys())

    def load_data(self, index, flip=False):
        if flip:
            self.prefix = 'flipped_'
        else:
            self.prefix = ''
        
        return self.__getitem__(index)
    
    def __getitem__(self, _index):
        if _index >= len(self): return
        
        index = sorted(list(self.tracking_results.keys()))[_index]
            
        # Process 2D keypoints
        kp2d = torch.from_numpy(self.tracking_results[index][self.prefix + 'keypoints']).float()
        mask = kp2d[..., -1] < KEYPOINTS_THR
        bbox = torch.from_numpy(self.tracking_results[index][self.prefix + 'bbox']).float()
        
        norm_kp2d, _ = self.keypoints_normalizer(
            kp2d[..., :-1].clone(), self.res, self.intrinsics, 224, 224, bbox
        )
        
        # Process image features
        features = self.tracking_results[index][self.prefix + 'features']
        
        # Process initial pose
        init_output = self.smpl.get_output(
            global_orient=self.tracking_results[index][self.prefix + 'init_global_orient'],
            body_pose=self.tracking_results[index][self.prefix + 'init_body_pose'],
            betas=self.tracking_results[index][self.prefix + 'init_betas'],
            pose2rot=False,
            return_full_pose=True
        )
        init_kp3d = root_centering(init_output.joints[:, :17], 'coco')
        init_kp = torch.cat((init_kp3d.reshape(1, -1), norm_kp2d[0].clone().reshape(1, -1)), dim=-1)
        init_smpl = transforms.matrix_to_rotation_6d(init_output.full_pose)
        init_root = transforms.matrix_to_rotation_6d(init_output.global_orient)
        
        # Process SLAM results
        cam_angvel = convert_dpvo_to_cam_angvel(self.slam_results, self.fps)
        
        return (
            index,                                          # subject id
            self._to(norm_kp2d),                            # 2d keypoints
            (self._to(init_kp), self._to(init_smpl)),       # initial pose
            self._to(features),                             # image features
            self._to(mask),                                 # keypoints mask
            init_root.to(self.device),                      # initial root orientation
            self._to(cam_angvel),                           # camera angular velocity
            self.tracking_results[index]['frame_id'],       # frame indices
            {'cam_intrinsics': self._to(self.intrinsics),   # other keyword arguments
             'bbox': self._to(bbox),
             'res': self._to(self.res)},
            )