# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import List, Optional

import torch

from torchlie.functional import SE3

from .robot import Joint, Link, Robot

# TODO: Add support for joints with DOF>1


def ForwardKinematicsFactory(robot: Robot, link_names: Optional[List[str]] = None):
    links: List[Link] = robot.get_links(link_names)
    link_ids: List[int] = [link.id for link in links]

    ancestor_links: List[Link] = []
    joint_ids: List[int] = []
    for link in links:
        ancestor_links += [anc for anc in link.ancestor_links]
        joint_ids += link.ancestor_non_fixed_joint_ids
    related_link_ids = sorted(list(set([anc.id for anc in ancestor_links] + link_ids)))
    joint_ids = sorted(list(set(joint_ids)))

    def _forward_kinematics_helper(angles: torch.Tensor):
        if angles.ndim != 2 or angles.shape[1] != robot.dof:
            raise ValueError(
                f"Joint angles for {robot.name} should be {robot.dof}-D vectors"
            )

        poses: List[Optional[torch.Tensor]] = [None] * robot.num_links
        poses[0] = angles.new_zeros(angles.shape[0], 3, 4)
        poses[0][:, 0, 0] = 1
        poses[0][:, 1, 1] = 1
        poses[0][:, 2, 2] = 1

        for id in related_link_ids[1:]:
            curr: Link = robot.links[id]
            joint: Joint = robot.links[id].parent_joint
            prev: Link = joint.parent_link
            relative_pose = (
                joint.relative_pose(angles[:, joint.id])
                if joint.id < robot.dof
                else joint.relative_pose()
            )
            poses[curr.id] = SE3.compose(poses[prev.id], relative_pose)

        return poses

    def _forward_kinematics_impl(angles: torch.Tensor):
        poses = _forward_kinematics_helper(angles)
        return tuple(poses[id] for id in link_ids)

    def _forward_kinematics_backward_helper(
        poses: List[Optional[torch.Tensor]],
    ) -> torch.Tensor:
        jposes = poses[0].new_zeros(poses[0].shape[0], 6, robot.dof)

        for id in related_link_ids[1:]:
            link: Link = robot.links[id]
            joint: Joint = link.parent_joint
            if joint.id >= robot.dof:
                break
            jposes[:, :, joint.id : joint.id + 1] = SE3.adj(poses[link.id]) @ joint.axis

        return jposes

    class ForwardKinematics(torch.autograd.Function):
        generate_vmap_rule = True

        @staticmethod
        def forward(angles):
            poses = _forward_kinematics_helper(angles)
            rets = tuple(poses[id] for id in link_ids)
            return *rets, poses

        @staticmethod
        def setup_context(ctx, inputs, output):
            ctx.save_for_backward(*output[:-1])
            ctx.poses = output[-1]

        @staticmethod
        def backward(ctx, *grad_outputs):
            rets: tuple(torch.Tensor) = ctx.saved_tensors
            if not hasattr(ctx, "jposes"):
                ctx.jposes: torch.Tensor = _forward_kinematics_backward_helper(
                    ctx.poses
                )
            grad_pose = grad_outputs[0].new_zeros(
                grad_outputs[0].shape[0], 6, robot.dof
            )
            grad_input = grad_outputs[0].new_zeros(grad_outputs[0].shape[0], robot.dof)

            for link_id, ret, grad_output in zip(link_ids, rets, grad_outputs):
                ancestor_non_fixed_joint_ids = robot.links[
                    link_id
                ].ancestor_non_fixed_joint_ids
                temp = SE3.project(
                    torch.cat(
                        (grad_output @ ret.transpose(1, 2), grad_output[:, :, 3:]),
                        dim=-1,
                    )
                ).unsqueeze(-1)
                grad_pose[:, :, ancestor_non_fixed_joint_ids] += temp
            grad_input[:, joint_ids] = (
                ctx.jposes[:, :, joint_ids] * grad_pose[:, :, joint_ids]
            ).sum(-2)

            return grad_input

    return (
        ForwardKinematics,
        _forward_kinematics_impl,
        _forward_kinematics_helper,
        _forward_kinematics_backward_helper,
        link_ids,
        joint_ids,
        related_link_ids,
    )


def get_forward_kinematics_fns(robot: Robot, link_names: Optional[List[str]] = None):
    (
        ForwardKinematics,
        _,
        _,
        backward_helper,
        link_ids,
        _,
        related_link_ids,
    ) = ForwardKinematicsFactory(robot, link_names)

    links = robot.get_links()
    selected_link_names = [links[id].name for id in related_link_ids]
    SelectedForwardKinematics, *_ = ForwardKinematicsFactory(robot, selected_link_names)

    def forward_kinematics(angles: torch.Tensor):
        output = ForwardKinematics.apply(angles)
        return output[:-1]

    def forward_kinematics_spatial_jacobian(angles: torch.Tensor):
        selected_poses: torch.Tensor = SelectedForwardKinematics.apply(angles)
        poses = [None] * robot.num_links
        for id, pose in zip(related_link_ids, selected_poses):
            poses[id] = pose
        jposes: torch.Tensor = backward_helper(poses)

        rets = tuple(poses[id] for id in link_ids)
        jacs_s: List[torch.Tensor] = []

        for link_id in link_ids:
            pose = poses[link_id]
            jac_s = jposes.new_zeros(angles.shape[0], 6, robot.dof)
            sel = robot.links[link_id].ancestor_non_fixed_joint_ids
            jac_s[:, :, sel] = jposes[:, :, sel]
            jacs_s.append(jac_s)

        return jacs_s, rets

    def forward_kinematics_body_jacobian(angles: torch.Tensor):
        jacs_s, rets = forward_kinematics_spatial_jacobian(angles)
        jacs_b: List[torch.Tensor] = []

        for jac_s, pose in zip(jacs_s, rets):
            jacs_b.append(SE3.adj(SE3.inv(pose)) @ jac_s)

        return jacs_b, rets

    return (
        forward_kinematics,
        forward_kinematics_body_jacobian,
        forward_kinematics_spatial_jacobian,
    )
