# Copyright The PyTorch Lightning team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any, List, Optional, Tuple, Union

import torch
from torch import Tensor

from torchmetrics.functional.classification.average_precision import _average_precision_compute_with_precision_recall
from torchmetrics.metric import Metric
from torchmetrics.utilities.data import METRIC_EPS, to_onehot
from torchmetrics.utilities.prints import rank_zero_warn


def _recall_at_precision(
    precision: Tensor,
    recall: Tensor,
    thresholds: Tensor,
    min_precision: float,
) -> Tuple[Tensor, Tensor]:
    try:
        max_recall, _, best_threshold = max(
            (r, p, t) for p, r, t in zip(precision, recall, thresholds) if p >= min_precision
        )

    except ValueError:
        max_recall = torch.tensor(0.0, device=recall.device, dtype=recall.dtype)
        best_threshold = torch.tensor(0)

    if max_recall == 0.0:
        best_threshold = torch.tensor(1e6, device=thresholds.device, dtype=thresholds.dtype)

    return max_recall, best_threshold


class BinnedPrecisionRecallCurve(Metric):
    """Computes precision-recall pairs for different thresholds. Works for both binary and multiclass problems. In
    the case of multiclass, the values will be calculated based on a one-vs-the-rest approach.

    Computation is performed in constant-memory by computing precision and recall
    for ``thresholds`` buckets/thresholds (evenly distributed between 0 and 1).

    .. warn:
        This metric has been deprecated in v0.10 and will be removed in v0.11.
        Instead use `PrecisionRecallCurve` metric with the `thresholds` argument set accordingly.

    Forward accepts

    - ``preds`` (float tensor): ``(N, ...)`` (binary) or ``(N, C, ...)`` (multiclass) tensor
      with probabilities, where C is the number of classes.

    - ``target`` (long tensor): ``(N, ...)`` or ``(N, C, ...)`` with integer labels

    Args:
        num_classes: integer with number of classes. For binary, set to 1.
        thresholds: list or tensor with specific thresholds or a number of bins from linear sampling.
            It is used for computation will lead to more detailed curve and accurate estimates,
            but will be slower and consume more memory.

        kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.

    Raises:
        ValueError:
            If ``thresholds`` is not a ``int``, ``list`` or ``tensor``

    Example (binary case):
        >>> from torchmetrics import BinnedPrecisionRecallCurve
        >>> pred = torch.tensor([0, 0.1, 0.8, 0.4])
        >>> target = torch.tensor([0, 1, 1, 0])
        >>> pr_curve = BinnedPrecisionRecallCurve(num_classes=1, thresholds=5)
        >>> precision, recall, thresholds = pr_curve(pred, target)
        >>> precision
        tensor([0.5000, 0.5000, 1.0000, 1.0000, 1.0000, 1.0000])
        >>> recall
        tensor([1.0000, 0.5000, 0.5000, 0.5000, 0.0000, 0.0000])
        >>> thresholds
        tensor([0.0000, 0.2500, 0.5000, 0.7500, 1.0000])

    Example (multiclass case):
        >>> pred = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],
        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],
        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],
        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])
        >>> target = torch.tensor([0, 1, 3, 2])
        >>> pr_curve = BinnedPrecisionRecallCurve(num_classes=5, thresholds=3)
        >>> precision, recall, thresholds = pr_curve(pred, target)
        >>> precision
        [tensor([0.2500, 1.0000, 1.0000, 1.0000]),
        tensor([0.2500, 1.0000, 1.0000, 1.0000]),
        tensor([2.5000e-01, 1.0000e-06, 1.0000e+00, 1.0000e+00]),
        tensor([2.5000e-01, 1.0000e-06, 1.0000e+00, 1.0000e+00]),
        tensor([2.5000e-07, 1.0000e+00, 1.0000e+00, 1.0000e+00])]
        >>> recall
        [tensor([1.0000, 1.0000, 0.0000, 0.0000]),
        tensor([1.0000, 1.0000, 0.0000, 0.0000]),
        tensor([1.0000, 0.0000, 0.0000, 0.0000]),
        tensor([1.0000, 0.0000, 0.0000, 0.0000]),
        tensor([0., 0., 0., 0.])]
        >>> thresholds
        [tensor([0.0000, 0.5000, 1.0000]),
        tensor([0.0000, 0.5000, 1.0000]),
        tensor([0.0000, 0.5000, 1.0000]),
        tensor([0.0000, 0.5000, 1.0000]),
        tensor([0.0000, 0.5000, 1.0000])]
    """

    is_differentiable: bool = False
    higher_is_better: Optional[bool] = None
    full_state_update: bool = False
    TPs: Tensor
    FPs: Tensor
    FNs: Tensor

    def __init__(
        self,
        num_classes: int,
        thresholds: Union[int, Tensor, List[float]] = 100,
        **kwargs: Any,
    ) -> None:
        rank_zero_warn(
            "Metric `BinnedPrecisionRecallCurve` has been deprecated in v0.10 and will be completly removed in v0.11."
            " Instead, use the refactored version of `PrecisionRecallCurve` by specifying the `thresholds` argument.",
            DeprecationWarning,
        )
        super().__init__(**kwargs)

        self.num_classes = num_classes
        if isinstance(thresholds, int):
            self.num_thresholds = thresholds
            self.thresholds = torch.linspace(0, 1.0, thresholds)

        elif thresholds is not None:
            if not isinstance(thresholds, (list, Tensor)):
                raise ValueError("Expected argument `thresholds` to either be an integer, list of floats or a tensor")
            self.thresholds = torch.tensor(thresholds) if isinstance(thresholds, list) else thresholds
            self.num_thresholds = self.thresholds.numel()

        for name in ("TPs", "FPs", "FNs"):
            self.add_state(
                name=name,
                default=torch.zeros(num_classes, self.num_thresholds, dtype=torch.float32),
                dist_reduce_fx="sum",
            )

    def update(self, preds: Tensor, target: Tensor) -> None:  # type: ignore
        """
        Args
            preds: (n_samples, n_classes) tensor
            target: (n_samples, n_classes) tensor
        """
        # binary case
        if len(preds.shape) == len(target.shape) == 1:
            preds = preds.reshape(-1, 1)
            target = target.reshape(-1, 1)

        if len(preds.shape) == len(target.shape) + 1:
            target = to_onehot(target, num_classes=self.num_classes)

        target = target == 1
        # Iterate one threshold at a time to conserve memory
        for i in range(self.num_thresholds):
            predictions = preds >= self.thresholds[i]
            self.TPs[:, i] += (target & predictions).sum(dim=0)
            self.FPs[:, i] += ((~target) & predictions).sum(dim=0)
            self.FNs[:, i] += (target & (~predictions)).sum(dim=0)

    def compute(self) -> Union[Tuple[Tensor, Tensor, Tensor], Tuple[List[Tensor], List[Tensor], List[Tensor]]]:
        """Returns float tensor of size n_classes."""
        precisions = (self.TPs + METRIC_EPS) / (self.TPs + self.FPs + METRIC_EPS)
        recalls = self.TPs / (self.TPs + self.FNs + METRIC_EPS)

        # Need to guarantee that last precision=1 and recall=0, similar to precision_recall_curve
        t_ones = torch.ones(self.num_classes, 1, dtype=precisions.dtype, device=precisions.device)
        precisions = torch.cat([precisions, t_ones], dim=1)
        t_zeros = torch.zeros(self.num_classes, 1, dtype=recalls.dtype, device=recalls.device)
        recalls = torch.cat([recalls, t_zeros], dim=1)
        if self.num_classes == 1:
            return precisions[0, :], recalls[0, :], self.thresholds
        return list(precisions), list(recalls), [self.thresholds for _ in range(self.num_classes)]


class BinnedAveragePrecision(BinnedPrecisionRecallCurve):
    """Computes the average precision score, which summarises the precision recall curve into one number. Works for
    both binary and multiclass problems. In the case of multiclass, the values will be calculated based on a one-
    vs-the-rest approach.

    Computation is performed in constant-memory by computing precision and recall
    for ``thresholds`` buckets/thresholds (evenly distributed between 0 and 1).

    .. warn:
        This metric has been deprecated in v0.10 and will be removed in v0.11.
        Instead use `AveragePrecision` metric with the `thresholds` argument set accordingly.

    Forward accepts

    - ``preds`` (float tensor): ``(N, ...)`` (binary) or ``(N, C, ...)`` (multiclass) tensor
      with probabilities, where C is the number of classes.

    - ``target`` (long tensor): ``(N, ...)`` with integer labels

    Args:
        num_classes: integer with number of classes. Not nessesary to provide for binary problems.
        thresholds: list or tensor with specific thresholds or a number of bins from linear sampling.
            It is used for computation will lead to more detailed curve and accurate estimates,
            but will be slower and consume more memory

        kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.

    Raises:
        ValueError:
            If ``thresholds`` is not a ``list`` or ``tensor``

    Example (binary case):
        >>> from torchmetrics import BinnedAveragePrecision
        >>> pred = torch.tensor([0, 1, 2, 3])
        >>> target = torch.tensor([0, 1, 1, 1])
        >>> average_precision = BinnedAveragePrecision(num_classes=1, thresholds=10)
        >>> average_precision(pred, target)
        tensor(1.0000)

    Example (multiclass case):
        >>> pred = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],
        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],
        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],
        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])
        >>> target = torch.tensor([0, 1, 3, 2])
        >>> average_precision = BinnedAveragePrecision(num_classes=5, thresholds=10)
        >>> average_precision(pred, target)
        [tensor(1.0000), tensor(1.0000), tensor(0.2500), tensor(0.2500), tensor(-0.)]
    """

    def __init__(
        self,
        num_classes: int,
        thresholds: Union[int, Tensor, List[float]] = 100,
        **kwargs: Any,
    ) -> None:
        rank_zero_warn(
            "Metric `BinnedAveragePrecision` has been deprecated in v0.10 and will be completly removed in v0.11."
            " Instead, use the refactored version of `AveragePrecision` by specifying the `thresholds` argument.",
            DeprecationWarning,
        )
        super().__init__(num_classes=num_classes, thresholds=thresholds, **kwargs)

    def compute(self) -> Union[List[Tensor], Tensor]:  # type: ignore
        precisions, recalls, _ = super().compute()
        return _average_precision_compute_with_precision_recall(precisions, recalls, self.num_classes, average=None)


class BinnedRecallAtFixedPrecision(BinnedPrecisionRecallCurve):
    """Computes the higest possible recall value given the minimum precision thresholds provided.

    Computation is performed in constant-memory by computing precision and recall
    for ``thresholds`` buckets/thresholds (evenly distributed between 0 and 1).

    .. warn:
        This metric has been deprecated in v0.10 and will be removed in v0.11.
        Instead use `RecallAtFixedPrecision` metric with the `thresholds` argument set accordingly.

    Forward accepts

    - ``preds`` (float tensor): ``(N, ...)`` (binary) or ``(N, C, ...)`` (multiclass) tensor
      with probabilities, where C is the number of classes.

    - ``target`` (long tensor): ``(N, ...)`` with integer labels

    Args:
        num_classes: integer with number of classes. Provide 1 for binary problems.
        min_precision: float value specifying minimum precision threshold.
        thresholds: list or tensor with specific thresholds or a number of bins from linear sampling.
            It is used for computation will lead to more detailed curve and accurate estimates,
            but will be slower and consume more memory

        kwargs: Additional keyword arguments, see :ref:`Metric kwargs` for more info.

    Raises:
        ValueError:
            If ``thresholds`` is not a list or tensor

    Example (binary case):
        >>> from torchmetrics import BinnedRecallAtFixedPrecision
        >>> pred = torch.tensor([0, 0.2, 0.5, 0.8])
        >>> target = torch.tensor([0, 1, 1, 0])
        >>> average_precision = BinnedRecallAtFixedPrecision(num_classes=1, thresholds=10, min_precision=0.5)
        >>> average_precision(pred, target)
        (tensor(1.0000), tensor(0.1111))

    Example (multiclass case):
        >>> pred = torch.tensor([[0.75, 0.05, 0.05, 0.05, 0.05],
        ...                      [0.05, 0.75, 0.05, 0.05, 0.05],
        ...                      [0.05, 0.05, 0.75, 0.05, 0.05],
        ...                      [0.05, 0.05, 0.05, 0.75, 0.05]])
        >>> target = torch.tensor([0, 1, 3, 2])
        >>> average_precision = BinnedRecallAtFixedPrecision(num_classes=5, thresholds=10, min_precision=0.5)
        >>> average_precision(pred, target)
        (tensor([1.0000, 1.0000, 0.0000, 0.0000, 0.0000]),
        tensor([6.6667e-01, 6.6667e-01, 1.0000e+06, 1.0000e+06, 1.0000e+06]))
    """

    def __init__(
        self,
        num_classes: int,
        min_precision: float,
        thresholds: Union[int, Tensor, List[float]] = 100,
        **kwargs: Any,
    ) -> None:
        rank_zero_warn(
            "Metric `BinnedRecallAtFixedPrecision` has been deprecated in v0.10 and will be completly removed in v0.11."
            " Instead, use the refactored version of `RecallAtFixedPrecision` by specifying the `thresholds` argument.",
            DeprecationWarning,
        )
        super().__init__(num_classes=num_classes, thresholds=thresholds, **kwargs)
        self.min_precision = min_precision

    def compute(self) -> Tuple[Tensor, Tensor]:  # type: ignore
        """Returns float tensor of size n_classes."""
        precisions, recalls, thresholds = super().compute()

        if self.num_classes == 1:
            return _recall_at_precision(precisions, recalls, thresholds, self.min_precision)

        recalls_at_p = torch.zeros(self.num_classes, device=recalls[0].device, dtype=recalls[0].dtype)
        thresholds_at_p = torch.zeros(self.num_classes, device=thresholds[0].device, dtype=thresholds[0].dtype)
        for i in range(self.num_classes):
            recalls_at_p[i], thresholds_at_p[i] = _recall_at_precision(
                precisions[i], recalls[i], thresholds[i], self.min_precision
            )
        return recalls_at_p, thresholds_at_p
