# 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.geometry.manifold import Manifold

from .lie_group import LieGroup


class Vector(LieGroup):
    def __init__(
        self,
        dof: Optional[int] = None,
        tensor: Optional[torch.Tensor] = None,
        name: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        if dof is not None and tensor is not None:
            raise ValueError("Please provide only one of dof or tensor.")
        if dof is None:
            if tensor is not None and tensor.ndim == 1:
                tensor = tensor.view(1, -1)
            if tensor is None or tensor.ndim != 2:
                raise ValueError(
                    "If dof not provided, tensor must " "have shape (batch_size, dof)"
                )
            dof = tensor.shape[1]

        super().__init__(dof, tensor=tensor, name=name, dtype=dtype)

    # Vector variables are of shape [batch_size, dof]
    @staticmethod
    def _init_tensor(dof: int) -> torch.Tensor:  # type: ignore
        return torch.zeros(1, dof)

    def dof(self) -> int:
        return self.tensor.shape[1]

    @staticmethod
    def rand(
        *size: int,
        generator: Optional[torch.Generator] = None,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
        requires_grad: bool = False,
    ) -> "Vector":
        if len(size) != 2:
            raise ValueError("The size should be 2D.")
        return Vector(
            tensor=torch.rand(
                size,
                generator=generator,
                dtype=dtype,
                device=device,
                requires_grad=requires_grad,
            )
        )

    @staticmethod
    def randn(
        *size: int,
        generator: Optional[torch.Generator] = None,
        dtype: Optional[torch.dtype] = None,
        device: Optional[torch.device] = None,
        requires_grad: bool = False,
    ) -> "Vector":
        if len(size) != 2:
            raise ValueError("The size should be 2D.")
        return Vector(
            tensor=torch.randn(
                size,
                generator=generator,
                dtype=dtype,
                device=device,
                requires_grad=requires_grad,
            )
        )

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}"
            f"(dof={self.tensor.shape[1]}, tensor={self.tensor}, name={self.name})"
        )

    def allclose(self, other: "Vector", *args, **kwargs) -> bool:
        return torch.allclose(self.tensor, other.tensor, *args, **kwargs)

    def __add__(self, other: "Vector") -> "Vector":
        return self.__class__._compose_impl(self, other)

    def __sub__(self, other: "Vector") -> "Vector":
        return self.__class__(tensor=torch.sub(self.tensor, other.tensor))

    def __neg__(self) -> "Vector":
        return self.__class__._inverse_impl(self)

    def __mul__(self, other: Union["Vector", torch.Tensor]) -> "Vector":
        if isinstance(other, Vector):
            return self.__class__(tensor=torch.mul(self.tensor, other.tensor))
        elif isinstance(other, torch.Tensor):
            return self.__class__(tensor=torch.mul(self.tensor, other))
        else:
            raise ValueError(
                f"expected type 'Vector' or 'Tensor', got {type(other).__name__}"
            )

    __rmul__ = __mul__

    def __matmul__(self, other: torch.Tensor) -> torch.Tensor:
        if not isinstance(other, torch.Tensor):
            raise ValueError("Vector matmul only accepts torch tensors.")
        if other.ndim != 3:
            raise ValueError(
                f"Vector matmul only accepts tensors with ndim=3 "
                f"but tensor has ndim={other.ndim}."
            )
        return torch.bmm(self.tensor.unsqueeze(1), other).squeeze(1)

    def __rmatmul__(self, other: torch.Tensor) -> torch.Tensor:
        if isinstance(other, Vector):
            raise ValueError("Vector matmul only accepts torch tensors.")
        if other.ndim != 3:
            raise ValueError(
                f"Vector matmul only accepts tensors with ndim=3 but "
                f"tensor has ndim={other.ndim}."
            )
        return torch.bmm(other, self.tensor.unsqueeze(2)).squeeze(2)

    def __truediv__(self, other: Union["Vector", torch.Tensor]) -> "Vector":
        if isinstance(other, Vector):
            return self.__class__(tensor=torch.div(self.tensor, other.tensor))
        elif isinstance(other, torch.Tensor):
            return self.__class__(tensor=torch.div(self.tensor, other))
        else:
            raise ValueError(
                f"expected type 'Vector' or 'Tensor', got {type(other).__name__}"
            )

    def dot(self, other: "Vector") -> torch.Tensor:
        return torch.mul(self.tensor, other.tensor).sum(-1)

    inner = dot

    def abs(self) -> "Vector":
        return self.__class__(tensor=torch.abs(self.tensor))

    def outer(self, other: "Vector") -> torch.Tensor:
        return torch.matmul(self.tensor.unsqueeze(2), other.tensor.unsqueeze(1))

    def norm(self, *args, **kwargs) -> torch.Tensor:
        return torch.norm(self.tensor, *args, **kwargs)

    def cat(self, vecs: Union["Vector", Tuple["Vector"], List["Vector"]]) -> "Vector":
        if isinstance(vecs, Vector):
            result = torch.cat([self.tensor, vecs.tensor], 1)
        else:
            result = torch.cat([self.tensor] + [vec.tensor for vec in vecs], 1)
        return Vector(tensor=result)

    def to_matrix(self) -> torch.Tensor:
        return self.tensor.clone()

    def _local_impl(
        self, vec2: Manifold, jacobians: List[torch.Tensor] = None
    ) -> torch.Tensor:
        if not isinstance(vec2, Vector):
            raise ValueError("Non-vector inputs for Vector.local()")
        else:
            return LieGroup._local_impl(self, vec2, jacobians)

    def _retract_impl(self, delta: torch.Tensor) -> "Vector":
        return self.__class__(tensor=torch.add(self.tensor, delta))

    def _compose_impl(self, vec2: LieGroup) -> "Vector":
        return self.__class__(tensor=torch.add(self.tensor, vec2.tensor))

    def _compose_jacobian(self, _: LieGroup) -> Tuple[torch.Tensor, torch.Tensor]:
        identity = torch.eye(self.dof(), dtype=self.dtype, device=self.device).repeat(
            self.shape[0], 1, 1
        )
        return (identity, identity.clone())

    def _inverse_impl(self) -> "Vector":
        return self.__class__(tensor=-self.tensor)

    def _adjoint_impl(self) -> torch.Tensor:
        return (
            torch.eye(self.dof(), dtype=self.dtype, device=self.device)
            .repeat(self.shape[0], 1, 1)
            .to(self.device)
        )

    def _project_impl(
        self, euclidean_grad: torch.Tensor, is_sparse: bool = False
    ) -> torch.Tensor:
        self._project_check(euclidean_grad, is_sparse)
        return euclidean_grad.clone()

    @staticmethod
    def exp_map(
        tangent_vector: torch.Tensor, jacobians: Optional[List[torch.Tensor]] = None
    ) -> "Vector":
        if tangent_vector.ndim != 2:
            raise ValueError("The dimension of tangent vectors should be 2.")

        Vector._exp_map_jacobian_impl(tangent_vector, jacobians)

        return Vector(tensor=tangent_vector.clone())

    @staticmethod
    def _check_tensor_impl(tensor: torch.Tensor) -> bool:
        if tensor.ndim != 2:
            raise ValueError("Vector variables expect tensors with ndim=2.")

        return True

    @staticmethod
    def normalize(tensor: torch.Tensor) -> torch.Tensor:
        if tensor.ndim != 2:
            raise ValueError("Vector variables expect tensors with ndim=2.")

        return tensor

    @staticmethod
    def _exp_map_jacobian_impl(
        tangent_vector: torch.Tensor, jacobians: Optional[List[torch.Tensor]]
    ):
        if jacobians is not None:
            shape = tangent_vector.shape
            LieGroup._check_jacobians_list(jacobians)
            jacobians.append(
                torch.eye(
                    shape[1], dtype=tangent_vector.dtype, device=tangent_vector.device
                ).repeat(shape[0], 1, 1)
            )

    def _log_map_impl(
        self, jacobians: Optional[List[torch.Tensor]] = None
    ) -> torch.Tensor:
        self._log_map_jacobian_impl(jacobians)
        return self.tensor.clone()

    def _log_map_jacobian_impl(self, jacobians: Optional[List[torch.Tensor]] = None):
        if jacobians is not None:
            shape = self.shape
            Vector._check_jacobians_list(jacobians)
            jacobians.append(
                torch.eye(shape[1], dtype=self.dtype, device=self.device).repeat(
                    shape[0], 1, 1
                )
            )

    def __hash__(self):
        return id(self)

    def _copy_impl(self, new_name: Optional[str] = None) -> "Vector":
        return self.__class__(tensor=self.tensor.clone(), name=new_name)

    # added to avoid casting downstream
    def copy(self, new_name: Optional[str] = None) -> "Vector":
        return cast(Vector, super().copy(new_name=new_name))


rand_vector = Vector.rand
randn_vector = Vector.randn
