# 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.

import abc
import warnings
from typing import Any, Dict, Optional, Type

import torch
import torch.linalg

from theseus.core import Objective
from theseus.optimizer import DenseLinearization, Linearization

from .linear_solver import LinearSolver


class DenseSolver(LinearSolver):
    def __init__(
        self,
        objective: Objective,
        linearization_cls: Optional[Type[Linearization]] = None,
        linearization_kwargs: Optional[Dict[str, Any]] = None,
        check_singular: bool = False,
    ):
        linearization_cls = linearization_cls or DenseLinearization
        if linearization_cls != DenseLinearization:
            raise RuntimeError(
                "DenseSolver only works with theseus.nonlinear.DenseLinearization, "
                f"but {linearization_cls} was provided."
            )
        super().__init__(objective, linearization_cls, linearization_kwargs)
        self.linearization: DenseLinearization = self.linearization
        self._check_singular = check_singular

    @staticmethod
    def _apply_damping(
        matrix: torch.Tensor,
        damping: float,
        ellipsoidal: bool = True,
        eps: float = 1e-8,
    ) -> torch.Tensor:
        if matrix.ndim != 3:
            raise ValueError(
                "Matrix must have a 3 dimensions, the first one being a batch dimension."
            )
        _, n, m = matrix.shape
        if n != m:
            raise ValueError("Matrix must be square.")

        # See Nocedal and Wright, Numerical Optimization, pp. 260 and 261
        # https://www.csie.ntu.edu.tw/~r97002/temp/num_optimization.pdf
        if ellipsoidal:
            # Add eps to guard against ill-conditioned matrix
            damped_D = torch.diag_embed(damping * matrix.diagonal(dim1=1, dim2=2) + eps)
        else:
            damped_D = damping * torch.eye(
                n, device=matrix.device, dtype=matrix.dtype
            ).unsqueeze(0)
        return matrix + damped_D

    def _apply_damping_and_solve(
        self,
        Atb: torch.Tensor,
        AtA: torch.Tensor,
        damping: Optional[float] = None,
        ellipsoidal_damping: bool = True,
        damping_eps: float = 1e-8,
    ) -> torch.Tensor:
        if damping:
            AtA = DenseSolver._apply_damping(
                AtA, damping, ellipsoidal=ellipsoidal_damping, eps=damping_eps
            )
        return self._solve_sytem(Atb, AtA)

    @abc.abstractmethod
    def _solve_sytem(self, Atb: torch.Tensor, AtA: torch.Tensor) -> torch.Tensor:
        pass

    def solve(
        self,
        damping: Optional[float] = None,
        ellipsoidal_damping: bool = True,
        damping_eps: float = 1e-8,
        **kwargs,
    ) -> torch.Tensor:
        if self._check_singular:
            AtA = self.linearization.AtA
            Atb = self.linearization.Atb
            with torch.no_grad():
                output = torch.zeros(AtA.shape[0], AtA.shape[1]).to(AtA.device)
                _, _, infos = torch.lu(AtA, get_infos=True)
                good_idx = infos.bool().logical_not()
                if not good_idx.all():
                    warnings.warn(
                        "Singular matrix found in batch, solution will be set "
                        "to all 0 for all singular matrices.",
                        RuntimeWarning,
                    )
            AtA = AtA[good_idx]
            Atb = Atb[good_idx]
            solution = self._apply_damping_and_solve(
                Atb,
                AtA,
                damping=damping,
                ellipsoidal_damping=ellipsoidal_damping,
                damping_eps=damping_eps,
            )
            output[good_idx] = solution
            return output
        else:
            return self._apply_damping_and_solve(
                self.linearization.Atb,
                self.linearization.AtA,
                damping=damping,
                ellipsoidal_damping=ellipsoidal_damping,
                damping_eps=damping_eps,
            )


class LUDenseSolver(DenseSolver):
    def __init__(
        self,
        objective: Objective,
        linearization_cls: Optional[Type[Linearization]] = DenseLinearization,
        linearization_kwargs: Optional[Dict[str, Any]] = None,
        check_singular: bool = False,
    ):
        super().__init__(
            objective,
            linearization_cls,
            linearization_kwargs,
            check_singular=check_singular,
        )

    def _solve_sytem(self, Atb: torch.Tensor, AtA: torch.Tensor) -> torch.Tensor:
        return torch.linalg.solve(AtA, Atb).squeeze(2)


class CholeskyDenseSolver(DenseSolver):
    def __init__(
        self,
        objective: Objective,
        linearization_cls: Optional[Type[Linearization]] = DenseLinearization,
        linearization_kwargs: Optional[Dict[str, Any]] = None,
        check_singular: bool = False,
    ):
        super().__init__(
            objective,
            linearization_cls,
            linearization_kwargs,
            check_singular=check_singular,
        )

    def _solve_sytem(self, Atb: torch.Tensor, AtA: torch.Tensor) -> torch.Tensor:
        lower = torch.linalg.cholesky(AtA)
        return torch.cholesky_solve(Atb, lower).squeeze(2)
