# 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 LieGroup, Vector


class DoubleIntegrator(CostFunction):
    def __init__(
        self,
        pose1: LieGroup,
        vel1: Vector,
        pose2: LieGroup,
        vel2: Vector,
        dt: Union[float, torch.Tensor, Variable],
        cost_weight: CostWeight,
        name: Optional[str] = None,
    ):
        super().__init__(cost_weight, name=name)
        dof = pose1.dof()
        if not (vel1.dof() == pose2.dof() == vel2.dof() == dof):
            raise ValueError(
                "All variables for a DoubleIntegrator must have the same dimension."
            )
        if not isinstance(dt, Variable):
            if not isinstance(dt, torch.Tensor):
                dt = torch.tensor(dt)
            self.dt = Variable(dt)
        else:
            self.dt = dt
        if self.dt.tensor.squeeze().ndim > 1:
            raise ValueError(
                "dt data must be a 0-D or 1-D tensor with numel in {1, batch_size}."
            )
        self.dt.tensor = self.dt.tensor.view(-1, 1)
        self.pose1 = pose1
        self.vel1 = vel1
        self.pose2 = pose2
        self.vel2 = vel2
        self.register_optim_vars(["pose1", "vel1", "pose2", "vel2"])
        self.register_aux_vars(["dt"])
        self.weight = cost_weight

    def dim(self):
        return 2 * self.pose1.dof()

    def _new_pose_diff(
        self, jacobians: Optional[List[torch.Tensor]] = None
    ) -> torch.Tensor:
        return self.pose1.local(self.pose2, jacobians=jacobians)

    def _error_from_pose_diff(self, pose_diff: torch.Tensor) -> torch.Tensor:
        pose_diff_err = pose_diff - self.dt.tensor.view(-1, 1) * self.vel1.tensor
        vel_diff = self.vel2.tensor - self.vel1.tensor
        return torch.cat([pose_diff_err, vel_diff], dim=1)

    def error(self) -> torch.Tensor:
        return self._error_from_pose_diff(self._new_pose_diff())

    def jacobians(self) -> Tuple[List[torch.Tensor], torch.Tensor]:
        # Pre-allocate jacobian tensors
        batch_size = self.pose1.shape[0]
        dof = self.pose1.dof()
        dtype = self.pose1.dtype
        device = self.pose1.device
        Jerr_pose1 = torch.zeros(batch_size, 2 * dof, dof, dtype=dtype, device=device)
        Jerr_vel1 = torch.zeros_like(Jerr_pose1)
        Jerr_pose2 = torch.zeros_like(Jerr_pose1)
        Jerr_vel2 = torch.zeros_like(Jerr_pose1)

        Jlocal: List[torch.Tensor] = []
        error = self._error_from_pose_diff(self._new_pose_diff(Jlocal))
        Jerr_pose1[:, :dof, :] = Jlocal[0]
        identity = torch.eye(dof, dtype=dtype, device=device).repeat(batch_size, 1, 1)
        Jerr_vel1[:, :dof, :] = -self.dt.tensor.view(-1, 1, 1) * identity
        Jerr_vel1[:, dof:, :] = -identity
        Jerr_pose2[:, :dof, :] = Jlocal[1]
        Jerr_vel2[:, dof:, :] = identity
        return [Jerr_pose1, Jerr_vel1, Jerr_pose2, Jerr_vel2], error

    def _copy_impl(self, new_name: Optional[str] = None) -> "DoubleIntegrator":
        return DoubleIntegrator(
            self.pose1.copy(),
            self.vel1.copy(),
            self.pose2.copy(),
            self.vel2.copy(),
            self.dt.copy(),
            self.weight.copy(),
            name=new_name,
        )


class GPCostWeight(CostWeight):
    # Qc_inv is either a single square matrix or a batch of square matrices
    # dt is either a single element or a 1-D batch
    def __init__(
        self,
        Qc_inv: Union[Variable, torch.Tensor],
        dt: Union[float, Variable, torch.Tensor],
        name: Optional[str] = None,
    ):
        super().__init__(name=name)
        if not isinstance(dt, Variable):
            if not isinstance(dt, torch.Tensor):
                dt = torch.tensor(dt)
            dt = Variable(dt)
        if dt.tensor.squeeze().ndim > 1:
            raise ValueError("dt must be a 0-D or 1-D tensor.")
        self.dt = dt
        self.dt.tensor = self.dt.tensor.view(-1, 1)

        if not isinstance(Qc_inv, Variable):
            Qc_inv = Variable(Qc_inv)
        if Qc_inv.ndim not in [2, 3]:
            raise ValueError("Qc_inv must be a single matrix or a batch of matrices.")
        if not Qc_inv.shape[-2] == Qc_inv.shape[-1]:
            raise ValueError("Qc_inv must contain square matrices.")
        self.Qc_inv = Qc_inv
        self.Qc_inv.tensor = (
            Qc_inv.tensor if Qc_inv.ndim == 3 else Qc_inv.tensor.unsqueeze(0)
        )
        self.register_aux_vars(["Qc_inv", "dt"])

    def _compute_cost_weight(self) -> torch.Tensor:
        batch_size, dof, _ = self.Qc_inv.shape
        cost_weight = torch.empty(
            batch_size,
            2 * dof,
            2 * dof,
            dtype=self.Qc_inv.dtype,
            device=self.Qc_inv.device,
        )
        dt_data = self.dt.tensor.view(-1, 1, 1)
        Q11 = 12.0 * dt_data.pow(-3.0) * self.Qc_inv.tensor
        Q12 = -6.0 * dt_data.pow(-2.0) * self.Qc_inv.tensor
        Q22 = 4.0 * dt_data.reciprocal() * self.Qc_inv.tensor
        cost_weight[:, :dof, :dof] = Q11
        cost_weight[:, :dof:, dof:] = Q12
        cost_weight[:, dof:, :dof] = Q12
        cost_weight[:, dof:, dof:] = Q22
        return (
            torch.linalg.cholesky(cost_weight.transpose(-2, -1).conj())
            .transpose(-2, -1)
            .conj()
        )

    def weight_error(self, error: torch.Tensor) -> torch.Tensor:
        weights = self._compute_cost_weight()
        return torch.matmul(weights, error.unsqueeze(2)).squeeze(2)

    def weight_jacobians_and_error(
        self,
        jacobians: List[torch.Tensor],
        error: torch.Tensor,
    ) -> Tuple[List[torch.Tensor], torch.Tensor]:
        cost_weight = self._compute_cost_weight()
        error = torch.matmul(cost_weight, error.unsqueeze(2)).squeeze(2)
        new_jacobians = []
        for jac in jacobians:
            # Jacobian is batch_size x cost_function_dim x var_dim
            # This left multiplies the weights (inv cov.) to jacobian
            new_jacobians.append(torch.matmul(cost_weight, jac))
        return new_jacobians, error

    def _copy_impl(self, new_name: Optional[str] = None) -> "GPCostWeight":
        # need to pass data, since it could be Parameter(self.Qc_inv) in which
        # case Qc_inv won't be up to date.
        # This will change with the "learning happens outside API"
        return GPCostWeight(self.Qc_inv.copy(), self.dt.copy(), name=new_name)


class GPMotionModel(DoubleIntegrator):
    def __init__(
        self,
        pose1: LieGroup,
        vel1: Vector,
        pose2: LieGroup,
        vel2: Vector,
        dt: Union[float, Variable, torch.Tensor],
        cost_weight: GPCostWeight,
        name: Optional[str] = None,
    ):
        if not isinstance(cost_weight, GPCostWeight):
            raise ValueError(
                "GPMotionModel only accepts cost weights of type GPCostWeight. "
                "For other weight types, consider using DoubleIntegrator instead."
            )
        if not isinstance(dt, Variable):
            if not isinstance(dt, torch.Tensor):
                dt = torch.tensor(dt)
            self.dt = Variable(dt)
        else:
            self.dt = dt
        if self.dt.tensor.squeeze().ndim > 1:
            raise ValueError("dt must be a 0-D or 1-D tensor.")
        self.dt.tensor = self.dt.tensor.view(-1, 1)
        super().__init__(pose1, vel1, pose2, vel2, dt, cost_weight, name=name)

    def _copy_impl(self, new_name: Optional[str] = None) -> "GPMotionModel":
        return cast(GPMotionModel, super()._copy_impl(new_name=new_name))
