# 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, Tuple, Union, cast

import torch

from theseus.core import CostFunction, CostWeight, Variable
from theseus.geometry import SE2, OptionalJacobians, Point2
from theseus.geometry.so2 import SO2


# Based on the model described in
# Zhou et al. A Fast Stochastic Contact Model for Planar Pushing and Grasping:
# Theory and Experimental Validation, 2017
# https://arxiv.org/abs/1705.10664
class QuasiStaticPushingPlanar(CostFunction):
    def __init__(
        self,
        obj1: SE2,
        obj2: SE2,
        eff1: SE2,
        eff2: SE2,
        c_square: Union[Variable, torch.Tensor, float],
        cost_weight: CostWeight,
        name: Optional[str] = None,
    ):
        super().__init__(cost_weight, name=name)
        self.obj1 = obj1
        self.obj2 = obj2
        self.eff1 = eff1
        self.eff2 = eff2
        self.register_optim_vars(["obj1", "obj2", "eff1", "eff2"])

        if not isinstance(c_square, Variable):
            if isinstance(c_square, float):
                c_square = torch.tensor(c_square, dtype=self.obj1.dtype).view(1, 1)
            c_square = Variable(c_square, name=f"csquare_{name}")
        if c_square.tensor.squeeze().ndim > 1:
            raise ValueError("dt must be a 0-D or 1-D tensor.")
        self.c_square = c_square
        self.register_aux_vars(["c_square"])

    def _get_tensor_info(self) -> Tuple[int, torch.device, torch.dtype]:
        return self.obj1.shape[0], self.obj1.device, self.obj1.dtype

    # See Eqs. 3-7 in Zhou et al.
    def _compute_D(
        self,
        contact_point_2: torch.Tensor,
        J_cp2_e2: OptionalJacobians,
        get_jacobians: bool = False,
    ) -> Tuple[torch.Tensor, OptionalJacobians]:
        batch_size, device, dtype = self._get_tensor_info()

        # Setting up intermediate jacobians, if needed
        J_cp2o: OptionalJacobians = [] if get_jacobians else None

        # Current contact point in object frame
        contact_point_2__obj = self.obj2.transform_to(contact_point_2, jacobians=J_cp2o)
        px = contact_point_2__obj.tensor[:, 0]
        py = contact_point_2__obj.tensor[:, 1]

        # Putting D together
        D = (
            torch.eye(3, 3, device=device, dtype=dtype)
            .unsqueeze(0)
            .repeat(batch_size, 1, 1)
        )
        D[:, 0, 2] = D[:, 2, 0] = -py
        D[:, 1, 2] = D[:, 2, 1] = px
        D[:, 2, 2] = -self.c_square.tensor.view(-1)

        if not get_jacobians:
            return D, None

        # Computing jacobians
        J_cp2o_o2, J_cp2o_cp2 = J_cp2o
        J_cp2o_e2 = J_cp2o_cp2.matmul(J_cp2_e2[0])

        def _get_J_D_var(J_cp2o_var_):
            # For matmul batching, this is set up with dimensions:
            #   batch_size x se2_dim x 3 x 3,
            # so that J_D_var_[b, d] is the (3, 3) matrix derivative
            # of the b-th batch D wrt to the d-th dimention of var
            J_D_var_ = torch.zeros(batch_size, 3, 3, 3, device=device, dtype=dtype)
            J_D_var_[..., 0, 2] = J_D_var_[..., 2, 0] = -J_cp2o_var_[:, 1]
            J_D_var_[..., 1, 2] = J_D_var_[..., 2, 1] = J_cp2o_var_[:, 0]
            return J_D_var_

        J_D_o2 = _get_J_D_var(J_cp2o_o2)
        J_D_e2 = _get_J_D_var(J_cp2o_e2)

        return D, [J_D_o2, J_D_e2]

    # See Eqs. 3-7 in Zhou et al.
    def _compute_V(
        self,
        obj2_angle: SO2,
        J_o2angle_o2: OptionalJacobians,
        get_jacobians: bool = False,
    ) -> Tuple[torch.Tensor, OptionalJacobians]:
        batch_size, device, dtype = self._get_tensor_info()

        # Setting up intermediate jacobians, if needed
        J_o2xy_o2: OptionalJacobians = [] if get_jacobians else None
        J_o1xy_o1: OptionalJacobians = [] if get_jacobians else None
        J_vxyoo: OptionalJacobians = [] if get_jacobians else None
        J_odw: OptionalJacobians = [] if get_jacobians else None
        J_omega_odw: OptionalJacobians = [] if get_jacobians else None

        # Compute object velocities using consecutive object poses
        vel_xy_obj__world = self.obj2.xy(jacobians=J_o2xy_o2) - self.obj1.xy(
            jacobians=J_o1xy_o1
        )
        vel_xy_obj__obj = obj2_angle.unrotate(vel_xy_obj__world, jacobians=J_vxyoo)

        # Putting V together
        obj_diff__world = cast(SE2, self.obj1.between(self.obj2, jacobians=J_odw))
        vx = vel_xy_obj__obj.tensor[:, 0]
        vy = vel_xy_obj__obj.tensor[:, 1]
        omega = obj_diff__world.theta(jacobians=J_omega_odw)
        V = torch.stack([vx, vy, omega.squeeze(1)], dim=1)

        if not get_jacobians:
            return V, None

        # Computing jacobians
        def _get_J_V_var(J_vxyoo_var_, J_omega_var_):
            # For matmul batching, this is set up with dimensions:
            #   batch_size x se2_dim x 3 x 1,
            # so that J_V_var_[b, d] is the (3, 1) vector derivative
            # of the b-th batch V wrt to the d-th dimention of var
            J_V_var_ = torch.zeros(batch_size, 3, 3, device=device, dtype=dtype)
            J_V_var_[..., 0] = J_vxyoo_var_[:, 0]
            J_V_var_[..., 1] = J_vxyoo_var_[:, 1]
            J_V_var_[..., 2] = J_omega_var_[:, 0]
            return J_V_var_.unsqueeze(3)

        # Computing velocity derivates wrt to objects 1, 2 variables
        J_vxyow_o2 = J_o2xy_o2[0]
        J_vxyow_o1 = -J_o1xy_o1[0]
        J_vxyoo_o2angle, J_vxyoo_vxyow = J_vxyoo
        J_vxyoo_o1 = J_vxyoo_vxyow.matmul(J_vxyow_o1)
        J_vxyoo_o2 = J_vxyoo_o2angle.matmul(J_o2angle_o2[0]) + J_vxyoo_vxyow.matmul(
            J_vxyow_o2
        )

        J_odw_o1, J_odw_o2 = J_odw
        J_omega_o1 = J_omega_odw[0].matmul(J_odw_o1)
        J_omega_o2 = J_omega_odw[0].matmul(J_odw_o2)

        J_V_o1 = _get_J_V_var(J_vxyoo_o1, J_omega_o1)
        J_V_o2 = _get_J_V_var(J_vxyoo_o2, J_omega_o2)

        return V, [J_V_o1, J_V_o2]

    # See Eqs. 3-7 in Zhou et al.
    def _compute_Vp(
        self,
        obj2_angle: SO2,
        J_o2angle_o2: OptionalJacobians,
        contact_point_2: torch.Tensor,
        J_cp2_e2: OptionalJacobians,
        get_jacobians: bool = False,
    ) -> Tuple[torch.Tensor, List[torch.Tensor]]:
        batch_size, device, dtype = self._get_tensor_info()

        # Setting up intermediate jacobians, if needed
        J_cp1_e1: OptionalJacobians = [] if get_jacobians else None
        J_vxyco: OptionalJacobians = [] if get_jacobians else None

        # Compute contact point velocities using consecutive end effector poses
        contact_point_1 = self.eff1.xy(jacobians=J_cp1_e1)
        vel_xy_contact__world = cast(Point2, contact_point_2 - contact_point_1)

        # Transform velocity to object's 2 frame
        vel_xy_contact__obj = obj2_angle.unrotate(
            vel_xy_contact__world, jacobians=J_vxyco
        )

        # Putting Vp together
        v_px = vel_xy_contact__obj.tensor[:, 0]
        v_py = vel_xy_contact__obj.tensor[:, 1]
        Vp = torch.stack([v_px, v_py, torch.zeros_like(v_px)], dim=1)

        if not get_jacobians:
            return Vp, None

        # Computing jacobians
        def _get_J_Vp_var(J_vxyco_var_):
            # For matmul batching, this is set up with dimensions:
            #   batch_size x se2_dim x 3 x 1,
            # so that J_Vp_var_[b, d] is the (3, 1) vector derivative
            # of the b-th batch Vp wrt to the d-th dimention of var
            J_Vp_var_ = torch.zeros(batch_size, 3, 3, device=device, dtype=dtype)
            J_Vp_var_[..., 0] = J_vxyco_var_[:, 0]  # batch_size x se2_dim
            J_Vp_var_[..., 1] = J_vxyco_var_[:, 1]  # batch_size x se2_dim
            return J_Vp_var_.unsqueeze(3)

        J_vxyco_o2angle, J_vxyco_vxycw = J_vxyco
        J_vxyco_o2 = J_vxyco_o2angle.matmul(J_o2angle_o2[0])
        J_vxyco_e1 = J_vxyco_vxycw.matmul(-J_cp1_e1[0])
        J_vxyco_e2 = J_vxyco_vxycw.matmul(J_cp2_e2[0])

        J_Vp_o2 = _get_J_Vp_var(J_vxyco_o2)
        J_Vp_e1 = _get_J_Vp_var(J_vxyco_e1)
        J_Vp_e2 = _get_J_Vp_var(J_vxyco_e2)

        return Vp, [J_Vp_o2, J_Vp_e1, J_Vp_e2]

    def _error_and_jacobians_impl(
        self, get_jacobians: bool
    ) -> Tuple[torch.Tensor, OptionalJacobians]:
        J_o2angle_o2: OptionalJacobians = [] if get_jacobians else None
        J_cp2_e2: OptionalJacobians = [] if get_jacobians else None

        # These quantities are needed by two or more of D, V, Vp, so we compute once.
        obj2_angle = self.obj2.rotation
        self.obj2.theta(jacobians=J_o2angle_o2)
        contact_point_2 = self.eff2.xy(jacobians=J_cp2_e2)

        # D * V = Vp, See Zhou et al.
        # A Fast Stochastic Contact Model for Planar Pushing and Grasping:
        # Theory and Experimental Validation, 2017
        # https://arxiv.org/abs/1705.10664
        D, J_D = self._compute_D(contact_point_2, J_cp2_e2, get_jacobians=get_jacobians)
        V, J_V = self._compute_V(obj2_angle, J_o2angle_o2, get_jacobians=get_jacobians)
        Vp, J_Vp = self._compute_Vp(
            obj2_angle,
            J_o2angle_o2,
            contact_point_2,
            J_cp2_e2,
            get_jacobians=get_jacobians,
        )

        error = torch.bmm(D, V.unsqueeze(2)).squeeze(2) - Vp

        if not get_jacobians:
            return error, None

        # Computing jacobians
        V_exp = V.view(-1, 1, 3, 1)
        D_exp = D.view(-1, 1, 3, 3)
        sum_str = "bdij,bdjk->bdik"

        J_D_o2, J_D_e2 = J_D
        J_V_o1, J_V_o2 = J_V
        J_Vp_o2, J_Vp_e1, J_Vp_e2 = J_Vp

        def _get_J_err_var_(J_D_var_, J_V_var_, J_Vp_var_):
            # After einsum, each result is a tensor with size:
            #   batch_size x se2_dim x 3 x 1
            # so that after all terms are summed,
            # J_err_var_[b, d] is a (3, 1) vector derivative of
            # the b-th batch element of the 3-D error wrt to the
            # d-th dimension of the optimization variable.
            # At return, we squeeze the trailing dimension, and permute, so that the
            # output has shape (batch_size, 3, se2_dim), which is the expected jacobian
            # shape.
            J_err_var_ = torch.zeros_like(J_V_o1)
            if J_D_var_ is not None:
                J_err_var_ += torch.einsum(sum_str, J_D_var_, V_exp)
            if J_V_var_ is not None:
                J_err_var_ += torch.einsum(sum_str, D_exp, J_V_var_)
            if J_Vp_var_ is not None:
                J_err_var_ += -J_Vp_var_
            return J_err_var_.squeeze(3).permute(0, 2, 1)

        J_err_o1 = _get_J_err_var_(None, J_V_o1, None)
        J_err_o2 = _get_J_err_var_(J_D_o2, J_V_o2, J_Vp_o2)
        J_err_e1 = _get_J_err_var_(None, None, J_Vp_e1)
        J_err_e2 = _get_J_err_var_(J_D_e2, None, J_Vp_e2)

        return error, [J_err_o1, J_err_o2, J_err_e1, J_err_e2]

    def error(self) -> torch.Tensor:
        return self._error_and_jacobians_impl(get_jacobians=False)[0]

    def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
        error, jacobians = self._error_and_jacobians_impl(get_jacobians=True)
        return jacobians, error

    def dim(self) -> int:
        return 3

    def _copy_impl(self, new_name: Optional[str] = None) -> "QuasiStaticPushingPlanar":
        return QuasiStaticPushingPlanar(
            self.obj1.copy(),
            self.obj2.copy(),
            self.eff1.copy(),
            self.eff2.copy(),
            self.c_square.copy(),
            self.weight.copy(),
            name=new_name,
        )
