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

_LOSS_EPS = 1e-20


class RobustLoss(abc.ABC):
    @classmethod
    def evaluate(cls, x: torch.Tensor, log_radius: torch.Tensor) -> torch.Tensor:
        return cls._evaluate_impl(x, log_radius.exp())

    @classmethod
    def linearize(cls, x: torch.Tensor, log_radius: torch.Tensor) -> torch.Tensor:
        return cls._linearize_impl(x, log_radius.exp())

    @staticmethod
    @abc.abstractmethod
    def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
        pass

    @staticmethod
    @abc.abstractmethod
    def _linearize_impl(x: torch.Tensor, raidus: torch.Tensor) -> torch.Tensor:
        pass


class WelschLoss(RobustLoss):
    @staticmethod
    def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
        return radius - radius * torch.exp(-x / (radius + _LOSS_EPS))

    @staticmethod
    def _linearize_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
        return torch.exp(-x / (radius + _LOSS_EPS))


class HuberLoss(RobustLoss):
    @staticmethod
    def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
        return torch.where(
            x > radius, 2 * torch.sqrt(radius * x.max(radius) + _LOSS_EPS) - radius, x
        )

    @staticmethod
    def _linearize_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
        return torch.sqrt(radius / torch.max(x, radius) + _LOSS_EPS)


class HingeLoss(RobustLoss):
    @staticmethod
    def _evaluate_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
        return torch.where(x > radius, torch.sqrt(x) - torch.sqrt(radius), _LOSS_EPS)

    @staticmethod
    def _linearize_impl(x: torch.Tensor, radius: torch.Tensor) -> torch.Tensor:
        return torch.where(x > radius, 1.0 / (2 * torch.sqrt(x) + _LOSS_EPS), 0.0)


class GNCRobustLoss(RobustLoss, abc.ABC):
    @classmethod
    def evaluate(  # type: ignore
        cls, x: torch.Tensor, log_radius: torch.Tensor, mu: torch.Tensor
    ) -> torch.Tensor:
        return cls._evaluate_impl(x, log_radius.exp(), mu)

    @classmethod
    def linearize(  # type: ignore
        cls, x: torch.Tensor, log_radius: torch.Tensor, mu: torch.Tensor
    ) -> torch.Tensor:
        return cls._linearize_impl(x, log_radius.exp(), mu)

    @staticmethod
    @abc.abstractmethod
    def _evaluate_impl(  # type: ignore
        x: torch.Tensor, radius: torch.Tensor, mu: torch.Tensor
    ) -> torch.Tensor:
        pass

    @staticmethod
    @abc.abstractmethod
    def _linearize_impl(  # type: ignore
        x: torch.Tensor, raidus: torch.Tensor, mu: torch.Tensor
    ) -> torch.Tensor:
        pass


# Geman-McClure loss is one common robust loss that can be adaptively tuned
# with a continous control parameter `mu`, hence very suited for GNC.
# More context on adaptive robust loss at https://arxiv.org/pdf/1701.03077.pdf
class GemanMcClureLoss(GNCRobustLoss):
    @staticmethod
    def _evaluate_impl(  # type: ignore
        x: torch.Tensor, radius: torch.Tensor, mu: torch.Tensor
    ) -> torch.Tensor:
        """
        x: squared loss (i.e., r^2)
        radius: scale of quadratic bowl (also squared)
        mu: GNC control parameter ranging from 1 (full Geman McClure) to
            positive infinity (full Quadratic)
        """
        return mu * radius * x / (mu * radius + x + _LOSS_EPS)

    @staticmethod
    def _linearize_impl(  # type: ignore
        x: torch.Tensor, radius: torch.Tensor, mu: torch.Tensor
    ) -> torch.Tensor:
        return (mu * radius) ** 2 / ((mu * radius + x) ** 2 + _LOSS_EPS)
