import torch
from utils.transforms import quat2mat, repr6d2mat, euler2mat


class ForwardKinematics:
    def __init__(self, parents, offsets=None):
        self.parents = parents
        if offsets is not None and len(offsets.shape) == 2:
            offsets = offsets.unsqueeze(0)
        self.offsets = offsets

    def forward(self, rots, offsets=None, global_pos=None):
        """
        Forward Kinematics: returns a per-bone transformation
        @param rots: local joint rotations (batch_size, bone_num, 3, 3)
        @param offsets: (batch_size, bone_num, 3) or None
        @param global_pos: global_position: (batch_size, 3) or keep it as in offsets (default)
        @return: (batch_szie, bone_num, 3, 4)
        """
        rots = rots.clone()
        if offsets is None:
            offsets = self.offsets.to(rots.device)
        if global_pos is None:
            global_pos = offsets[:, 0]

        pos = torch.zeros((rots.shape[0], rots.shape[1], 3), device=rots.device)
        rest_pos = torch.zeros_like(pos)
        res = torch.zeros((rots.shape[0], rots.shape[1], 3, 4), device=rots.device)

        pos[:, 0] = global_pos
        rest_pos[:, 0] = offsets[:, 0]

        for i, p in enumerate(self.parents):
            if i != 0:
                rots[:, i] = torch.matmul(rots[:, p], rots[:, i])
                pos[:, i] = torch.matmul(rots[:, p], offsets[:, i].unsqueeze(-1)).squeeze(-1) + pos[:, p]
                rest_pos[:, i] = rest_pos[:, p] + offsets[:, i]

            res[:, i, :3, :3] = rots[:, i]
            res[:, i, :, 3] = torch.matmul(rots[:, i], -rest_pos[:, i].unsqueeze(-1)).squeeze(-1) + pos[:, i]

        return res

    def accumulate(self, local_rots):
        """
        Get global joint rotation from local rotations
        @param local_rots: (batch_size, n_bone, 3, 3)
        @return: global_rotations
        """
        res = torch.empty_like(local_rots)
        for i, p in enumerate(self.parents):
            if i == 0:
                res[:, i] = local_rots[:, i]
            else:
                res[:, i] = torch.matmul(res[:, p], local_rots[:, i])
        return res

    def unaccumulate(self, global_rots):
        """
        Get local joint rotation from global rotations
        @param global_rots: (batch_size, n_bone, 3, 3)
        @return: local_rotations
        """
        res = torch.empty_like(global_rots)
        inv = torch.empty_like(global_rots)

        for i, p in enumerate(self.parents):
            if i == 0:
                inv[:, i] = global_rots[:, i].transpose(-2, -1)
                res[:, i] = global_rots[:, i]
                continue
            res[:, i] = torch.matmul(inv[:, p], global_rots[:, i])
            inv[:, i] = torch.matmul(res[:, i].transpose(-2, -1), inv[:, p])

        return res


class ForwardKinematicsJoint:
    def __init__(self, parents, offset):
        self.parents = parents
        self.offset = offset

    '''
        rotation should have shape batch_size * Joint_num * (3/4) * Time
        position should have shape batch_size * 3 * Time
        offset should have shape batch_size * Joint_num * 3
        output have shape batch_size * Time * Joint_num * 3
    '''

    def forward(self, rotation: torch.Tensor, position: torch.Tensor, offset=None,
                world=True):
        '''
        if not quater and rotation.shape[-2] != 3: raise Exception('Unexpected shape of rotation')
        if quater and rotation.shape[-2] != 4: raise Exception('Unexpected shape of rotation')
        rotation = rotation.permute(0, 3, 1, 2)
        position = position.permute(0, 2, 1)
        '''
        if rotation.shape[-1] == 6:
            transform = repr6d2mat(rotation)
        elif rotation.shape[-1] == 4:
            norm = torch.norm(rotation, dim=-1, keepdim=True)
            rotation = rotation / norm
            transform = quat2mat(rotation)
        elif rotation.shape[-1] == 3:
            transform = euler2mat(rotation)
        else:
            raise Exception('Only accept quaternion rotation input')
        result = torch.empty(transform.shape[:-2] + (3,), device=position.device)

        if offset is None:
            offset = self.offset
        offset = offset.reshape((-1, 1, offset.shape[-2], offset.shape[-1], 1))

        result[..., 0, :] = position
        for i, pi in enumerate(self.parents):
            if pi == -1:
                assert i == 0
                continue

            result[..., i, :] = torch.matmul(transform[..., pi, :, :], offset[..., i, :, :]).squeeze()
            transform[..., i, :, :] = torch.matmul(transform[..., pi, :, :].clone(), transform[..., i, :, :].clone())
            if world: result[..., i, :] += result[..., pi, :]
        return result


class InverseKinematicsJoint:
    def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains):
        self.rotations = rotations.detach().clone()
        self.rotations.requires_grad_(True)
        self.position = positions.detach().clone()
        self.position.requires_grad_(True)

        self.parents = parents
        self.offset = offset
        self.constrains = constrains

        self.optimizer = torch.optim.Adam([self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999))
        self.criteria = torch.nn.MSELoss()

        self.fk = ForwardKinematicsJoint(parents, offset)

        self.glb = None

    def step(self):
        self.optimizer.zero_grad()
        glb = self.fk.forward(self.rotations, self.position)
        loss = self.criteria(glb, self.constrains)
        loss.backward()
        self.optimizer.step()
        self.glb = glb
        return loss.item()


class InverseKinematicsJoint2:
    def __init__(self, rotations: torch.Tensor, positions: torch.Tensor, offset, parents, constrains, cid,
                 lambda_rec_rot=1., lambda_rec_pos=1., use_velo=False):
        self.use_velo = use_velo
        self.rotations_ori = rotations.detach().clone()
        self.rotations = rotations.detach().clone()
        self.rotations.requires_grad_(True)
        self.position_ori = positions.detach().clone()
        self.position = positions.detach().clone()
        if self.use_velo:
            self.position[1:] = self.position[1:] - self.position[:-1]
        self.position.requires_grad_(True)

        self.parents = parents
        self.offset = offset
        self.constrains = constrains.detach().clone()
        self.cid = cid

        self.lambda_rec_rot = lambda_rec_rot
        self.lambda_rec_pos = lambda_rec_pos

        self.optimizer = torch.optim.Adam([self.position, self.rotations], lr=1e-3, betas=(0.9, 0.999))
        self.criteria = torch.nn.MSELoss()

        self.fk = ForwardKinematicsJoint(parents, offset)

        self.glb = None

    def step(self):
        self.optimizer.zero_grad()
        if self.use_velo:
            position = torch.cumsum(self.position, dim=0)
        else:
            position = self.position
        glb = self.fk.forward(self.rotations, position)
        self.constrain_loss = self.criteria(glb[:, self.cid], self.constrains)
        self.rec_loss_rot = self.criteria(self.rotations, self.rotations_ori)
        self.rec_loss_pos = self.criteria(self.position, self.position_ori)
        loss = self.constrain_loss + self.rec_loss_rot * self.lambda_rec_rot + self.rec_loss_pos * self.lambda_rec_pos
        loss.backward()
        self.optimizer.step()
        self.glb = glb
        return loss.item()

    def get_position(self):
        if self.use_velo:
            position = torch.cumsum(self.position.detach(), dim=0)
        else:
            position = self.position.detach()
        return position