# 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.embodied.kinematics import IdentityModel
from theseus.geometry import SE2, Point2

from .signed_distance_field import SignedDistanceField2D


class EffectorObjectContactPlanar(CostFunction):
    def __init__(
        self,
        obj: SE2,
        eff: SE2,
        sdf_origin: Variable,
        sdf_data: Variable,
        sdf_cell_size: Variable,
        eff_radius: Union[float, Variable, torch.Tensor],
        cost_weight: CostWeight,
        name: Optional[str] = None,
        use_huber_loss: bool = False,
    ):
        super().__init__(cost_weight, name=name)
        self.obj = obj
        self.eff = eff
        self.sdf_origin = sdf_origin
        self.sdf_data = sdf_data
        self.sdf_cell_size = sdf_cell_size
        if not isinstance(eff_radius, Variable):
            if not isinstance(eff_radius, torch.Tensor):
                eff_radius = torch.tensor(eff_radius)
            self.eff_radius = Variable(eff_radius)
        else:
            self.eff_radius = eff_radius
        if self.eff_radius.tensor.squeeze().ndim > 1:
            raise ValueError("eff_radius must be a 0-D or 1-D tensor.")
        self.eff_radius.tensor = self.eff_radius.tensor.view(-1, 1)
        self.register_optim_vars(["obj", "eff"])
        self.register_aux_vars(
            ["sdf_origin", "sdf_data", "sdf_cell_size", "eff_radius"]
        )
        self.robot = IdentityModel()
        self.sdf = SignedDistanceField2D(sdf_origin, sdf_cell_size, sdf_data)
        self._use_huber = use_huber_loss

        if use_huber_loss:
            raise NotImplementedError(
                "Jacobians for huber loss are not yet implemented."
            )

    def _compute_distances_and_jacobians(
        self,
    ) -> Tuple[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
        J_transf: List[torch.Tensor] = []
        J_xy: List[torch.Tensor] = []
        eff__obj = self.obj.transform_to(
            self.eff.xy(jacobians=J_xy), jacobians=J_transf
        )
        J_transf_obj = J_transf[0]
        J_transf_eff = J_transf[1].matmul(J_xy[0])
        robot_state = cast(Point2, self.robot.forward_kinematics(eff__obj)["state"])
        dist, J_dist = self.sdf.signed_distance(robot_state.tensor.view(-1, 2, 1))
        J_out = (J_dist.matmul(J_transf_obj), J_dist.matmul(J_transf_eff))
        return dist, J_out

    def _error_from_distances(self, distances: torch.Tensor):
        if self._use_huber:
            eff_rad = self.eff_radius.tensor
            err = distances.clone()
            # linear (two-sided, otherwise this would be 0)
            gt_r_idx = distances >= 2 * eff_rad
            err[gt_r_idx] = distances[gt_r_idx] - 0.5 * eff_rad
            # quadratic
            lt_r_pos_idx = (distances < 2 * eff_rad).logical_or(distances > 0)
            err[lt_r_pos_idx] = (
                0.5 * (distances[lt_r_pos_idx] - eff_rad).square() / eff_rad
            )
            # linear
            neg_idx = distances < 0
            err[neg_idx] = -distances[neg_idx] + 0.5 * eff_rad
            return err
        else:
            return (distances - self.eff_radius.tensor).abs()

    def error(self) -> torch.Tensor:
        distances, _ = self._compute_distances_and_jacobians()
        return self._error_from_distances(distances)

    def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
        distances, jacobians = self._compute_distances_and_jacobians()
        error = self._error_from_distances(distances)
        if self._use_huber:
            raise NotImplementedError(
                "Jacobians for huber loss are not yet implemented."
            )
        else:
            lt_idx = distances < self.eff_radius.tensor
            jacobians[0][lt_idx] *= -1.0
            jacobians[1][lt_idx] *= -1.0
        return [jacobians[0], jacobians[1]], error

    def _copy_impl(
        self, new_name: Optional[str] = None
    ) -> "EffectorObjectContactPlanar":
        return EffectorObjectContactPlanar(
            self.obj.copy(),
            self.eff.copy(),
            self.sdf_origin.copy(),
            self.sdf_data.copy(),
            self.sdf_cell_size.copy(),
            self.eff_radius.copy(),
            self.weight.copy(),
            name=new_name,
        )

    def dim(self) -> int:
        return 1

    # This is needed so that the SDF container also updates with the new aux var
    def set_aux_var_at(self, index: int, variable: Variable):
        super().set_aux_var_at(index, variable)
        self.sdf.update_data(self.sdf_origin, self.sdf_data, self.sdf_cell_size)
