# 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, as_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."
            )
        self.dt = as_variable(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)
        dt = as_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)

        Qc_inv = as_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."
            )
        self.dt = as_variable(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))
