import torch

from pytorch_widedeep.callbacks import Callback

from .wdtypes import *  # noqa: F403


class Metric(object):
    def __init__(self):
        self._name = ""

    def reset(self):
        raise NotImplementedError("Custom Metrics must implement this function")

    def __call__(self, y_pred: Tensor, y_true: Tensor):
        raise NotImplementedError("Custom Metrics must implement this function")


class MultipleMetrics(object):
    def __init__(self, metrics: List[Metric], prefix: str = ""):

        instantiated_metrics = []
        for metric in metrics:
            if isinstance(metric, type):
                instantiated_metrics.append(metric())
            else:
                instantiated_metrics.append(metric)
        self._metrics = instantiated_metrics
        self.prefix = prefix

    def reset(self):
        for metric in self._metrics:
            metric.reset()

    def __call__(self, y_pred: Tensor, y_true: Tensor) -> Dict:
        logs = {}
        for metric in self._metrics:
            logs[self.prefix + metric._name] = metric(y_pred, y_true)
        return logs


class MetricCallback(Callback):
    def __init__(self, container: MultipleMetrics):
        self.container = container

    def on_epoch_begin(self, epoch: int, logs: Optional[Dict] = None):
        self.container.reset()


class Accuracy(Metric):
    def __init__(self, top_k: int = 1):
        r"""Class to calculate the accuracy for both binary and categorical problems

        Parameters
        ----------
        top_k: int, default = 1
            Accuracy will be computed using the top k most likely classes in
            multiclass problems

        Examples
        --------
        >>> import torch
        >>>
        >>> from pytorch_widedeep.metrics import Accuracy
        >>>
        >>> acc = Accuracy()
        >>> y_true = torch.tensor([0, 1, 0, 1]).view(-1, 1)
        >>> y_pred = torch.tensor([[0.3, 0.2, 0.6, 0.7]]).view(-1, 1)
        >>> acc(y_pred, y_true)
        0.5
        >>>
        >>> acc = Accuracy(top_k=2)
        >>> y_true = torch.tensor([0, 1, 2])
        >>> y_pred = torch.tensor([[0.3, 0.5, 0.2], [0.1, 0.1, 0.8], [0.1, 0.5, 0.4]])
        >>> acc(y_pred, y_true)
        0.6666666666666666
        """
        super(Accuracy, self).__init__()

        self.top_k = top_k
        self.correct_count = 0
        self.total_count = 0
        self._name = "acc"

    def reset(self):
        """
        resets counters to 0
        """
        self.correct_count = 0
        self.total_count = 0

    def __call__(self, y_pred: Tensor, y_true: Tensor) -> float:
        num_classes = y_pred.size(1)

        if num_classes == 1:
            y_pred = y_pred.round()
            y_true = y_true
        elif num_classes > 1:
            y_pred = y_pred.topk(self.top_k, 1)[1]
            y_true = y_true.view(-1, 1).expand_as(y_pred)

        self.correct_count += y_pred.eq(y_true).sum().item()  # type: ignore[assignment]
        self.total_count += len(y_pred)
        accuracy = float(self.correct_count) / float(self.total_count)
        return accuracy


class Precision(Metric):
    def __init__(self, average: bool = True):
        r"""Class to calculate the precision for both binary and categorical problems

        Parameters
        ----------
        average: bool, default = True
            This applies only to multiclass problems. if ``True`` calculate
            precision for each label, and finds their unweighted mean.

        Examples
        --------
        >>> import torch
        >>>
        >>> from pytorch_widedeep.metrics import Precision
        >>>
        >>> prec = Precision()
        >>> y_true = torch.tensor([0, 1, 0, 1]).view(-1, 1)
        >>> y_pred = torch.tensor([[0.3, 0.2, 0.6, 0.7]]).view(-1, 1)
        >>> prec(y_pred, y_true)
        0.5
        >>>
        >>> prec = Precision(average=True)
        >>> y_true = torch.tensor([0, 1, 2])
        >>> y_pred = torch.tensor([[0.7, 0.1, 0.2], [0.1, 0.1, 0.8], [0.1, 0.5, 0.4]])
        >>> prec(y_pred, y_true)
        0.3333333432674408
        """
        super(Precision, self).__init__()

        self.average = average
        self.true_positives = 0
        self.all_positives = 0
        self.eps = 1e-20
        self._name = "prec"

    def reset(self):
        """
        resets counters to 0
        """
        self.true_positives = 0
        self.all_positives = 0

    def __call__(self, y_pred: Tensor, y_true: Tensor) -> float:
        num_class = y_pred.size(1)

        if num_class == 1:
            y_pred = y_pred.round()
            y_true = y_true
        elif num_class > 1:
            y_true = torch.eye(num_class)[y_true.squeeze().long()]
            y_pred = y_pred.topk(1, 1)[1].view(-1)
            y_pred = torch.eye(num_class)[y_pred.long()]

        self.true_positives += (y_true * y_pred).sum(dim=0)  # type:ignore
        self.all_positives += y_pred.sum(dim=0)  # type:ignore

        precision = self.true_positives / (self.all_positives + self.eps)

        if self.average:
            return precision.mean().item()  # type:ignore
        else:
            return precision


class Recall(Metric):
    def __init__(self, average: bool = True):
        r"""Class to calculate the recall for both binary and categorical problems

        Parameters
        ----------
        average: bool, default = True
            This applies only to multiclass problems. if ``True`` calculate recall
            for each label, and finds their unweighted mean.

        Examples
        --------
        >>> import torch
        >>>
        >>> from pytorch_widedeep.metrics import Recall
        >>>
        >>> rec = Recall()
        >>> y_true = torch.tensor([0, 1, 0, 1]).view(-1, 1)
        >>> y_pred = torch.tensor([[0.3, 0.2, 0.6, 0.7]]).view(-1, 1)
        >>> rec(y_pred, y_true)
        0.5
        >>>
        >>> rec = Recall(average=True)
        >>> y_true = torch.tensor([0, 1, 2])
        >>> y_pred = torch.tensor([[0.7, 0.1, 0.2], [0.1, 0.1, 0.8], [0.1, 0.5, 0.4]])
        >>> rec(y_pred, y_true)
        0.3333333432674408
        """
        super(Recall, self).__init__()

        self.average = average
        self.true_positives = 0
        self.actual_positives = 0
        self.eps = 1e-20
        self._name = "rec"

    def reset(self):
        """
        resets counters to 0
        """
        self.true_positives = 0
        self.actual_positives = 0

    def __call__(self, y_pred: Tensor, y_true: Tensor) -> float:
        num_class = y_pred.size(1)

        if num_class == 1:
            y_pred = y_pred.round()
            y_true = y_true
        elif num_class > 1:
            y_true = torch.eye(num_class)[y_true.squeeze().long()]
            y_pred = y_pred.topk(1, 1)[1].view(-1)
            y_pred = torch.eye(num_class)[y_pred.long()]

        self.true_positives += (y_true * y_pred).sum(dim=0)  # type: ignore
        self.actual_positives += y_true.sum(dim=0)  # type: ignore

        recall = self.true_positives / (self.actual_positives + self.eps)

        if self.average:
            return recall.mean().item()  # type:ignore
        else:
            return recall


class FBetaScore(Metric):
    def __init__(self, beta: int, average: bool = True):
        r"""Class to calculate the fbeta score for both binary and categorical problems

        :math:`F_{\beta} = ((1 + {\beta}^2) * \frac{(precision * recall)}{({\beta}^2 * precision + recall)}`

        Parameters
        ----------
        beta: int
            Coefficient to control the balance between precision and recall
        average: bool, default = True
            This applies only to multiclass problems. if ``True`` calculate fbeta
            for each label, and find their unweighted mean.

        Examples
        --------
        >>> import torch
        >>>
        >>> from pytorch_widedeep.metrics import FBetaScore
        >>>
        >>> fbeta = FBetaScore(beta=2)
        >>> y_true = torch.tensor([0, 1, 0, 1]).view(-1, 1)
        >>> y_pred = torch.tensor([[0.3, 0.2, 0.6, 0.7]]).view(-1, 1)
        >>> fbeta(y_pred, y_true)
        0.5
        >>>
        >>> fbeta = FBetaScore(beta=2)
        >>> y_true = torch.tensor([0, 1, 2])
        >>> y_pred = torch.tensor([[0.7, 0.1, 0.2], [0.1, 0.1, 0.8], [0.1, 0.5, 0.4]])
        >>> fbeta(y_pred, y_true)
        0.3333333432674408
        """
        super(FBetaScore, self).__init__()

        self.beta = beta
        self.average = average
        self.precision = Precision(average=False)
        self.recall = Recall(average=False)
        self.eps = 1e-20
        self._name = "".join(["f", str(self.beta)])

    def reset(self):
        """
        resets precision and recall
        """
        self.precision.reset()
        self.recall.reset()

    def __call__(self, y_pred: Tensor, y_true: Tensor) -> float:

        prec = self.precision(y_pred, y_true)
        rec = self.recall(y_pred, y_true)
        beta2 = self.beta ** 2

        fbeta = ((1 + beta2) * prec * rec) / (beta2 * prec + rec + self.eps)

        if self.average:
            return fbeta.mean().item()  # type: ignore[attr-defined]
        else:
            return fbeta


class F1Score(Metric):
    def __init__(self, average: bool = True):
        r"""Class to calculate the f1 score for both binary and categorical problems

        Parameters
        ----------
        average: bool, default = True
            This applies only to multiclass problems. if ``True`` calculate f1 for
            each label, and find their unweighted mean.

        Examples
        --------
        >>> import torch
        >>>
        >>> from pytorch_widedeep.metrics import F1Score
        >>>
        >>> f1 = F1Score()
        >>> y_true = torch.tensor([0, 1, 0, 1]).view(-1, 1)
        >>> y_pred = torch.tensor([[0.3, 0.2, 0.6, 0.7]]).view(-1, 1)
        >>> f1(y_pred, y_true)
        0.5
        >>>
        >>> f1 = F1Score()
        >>> y_true = torch.tensor([0, 1, 2])
        >>> y_pred = torch.tensor([[0.7, 0.1, 0.2], [0.1, 0.1, 0.8], [0.1, 0.5, 0.4]])
        >>> f1(y_pred, y_true)
        0.3333333432674408
        """
        super(F1Score, self).__init__()

        self.average = average
        self.f1 = FBetaScore(beta=1, average=self.average)
        self._name = self.f1._name

    def reset(self):
        """
        resets counters to 0
        """
        self.f1.reset()

    def __call__(self, y_pred: Tensor, y_true: Tensor) -> float:
        return self.f1(y_pred, y_true)


class R2Score(Metric):
    def __init__(self):
        r"""
        Calculates the R-Squared, the
        `coefficient of determination <https://en.wikipedia.org/wiki/Coefficient_of_determination>`_:

        :math:`R^2 = 1 - \frac{\sum_{j=1}^n(y_j - \hat{y_j})^2}{\sum_{j=1}^n(y_j - \bar{y})^2}`,

        where :math:`\hat{y_j}` is the ground truth, :math:`y_j` is the predicted value and
        :math:`\bar{y}` is the mean of the ground truth.

        Examples
        --------
        >>> import torch
        >>>
        >>> from pytorch_widedeep.metrics import R2Score
        >>>
        >>> r2 = R2Score()
        >>> y_true = torch.tensor([3, -0.5, 2, 7]).view(-1, 1)
        >>> y_pred = torch.tensor([2.5, 0.0, 2, 8]).view(-1, 1)
        >>> r2(y_pred, y_true)
        0.9486081370449679
        """
        self.numerator = 0
        self.denominator = 0
        self.num_examples = 0
        self.y_true_sum = 0

        self._name = "r2"

    def reset(self):
        """
        resets counters to 0
        """
        self.numerator = 0
        self.denominator = 0
        self.num_examples = 0
        self.y_true_sum = 0

    def __call__(self, y_pred: Tensor, y_true: Tensor) -> float:

        self.numerator += ((y_pred - y_true) ** 2).sum().item()

        self.num_examples += y_true.shape[0]
        self.y_true_sum += y_true.sum().item()
        y_true_avg = self.y_true_sum / self.num_examples
        self.denominator += ((y_true - y_true_avg) ** 2).sum().item()
        return 1 - (self.numerator / self.denominator)
