# flake8: noqa
import warnings

import pytest
import torch
from sklearn.exceptions import UndefinedMetricWarning
from sklearn.metrics import recall_score

from torchflare.metrics.recall_meter import Recall
from torchflare.metrics.meters import _BaseInputHandler

torch.manual_seed(42)


def test_binary_inputs():
    def _test(num_classes, threshold, multilabel, average):

        rc = Recall(num_classes=num_classes, threshold=threshold, multilabel=multilabel, average=average,)
        outputs = torch.randn(100, 1)
        targets = torch.randint(0, 2, size=(100,))

        bs = _BaseInputHandler(num_classes=num_classes, average=average, threshold=0.5, multilabel=multilabel,)

        np_outputs, np_targets = bs._compute(outputs=outputs, targets=targets)
        rc.accumulate(outputs=outputs, targets=targets)
        rec_val = rc.value
        assert rc.case_type == "binary"

        rec_skm = recall_score(np_targets.numpy(), np_outputs.numpy(), average="binary")

        assert rec_skm == pytest.approx(rec_val.item())

        rc.reset()
        bs = 16
        iters = targets.shape[0] // bs + 1

        for i in range(iters):
            idx = i * bs
            rc.accumulate(outputs=outputs[idx : idx + bs], targets=targets[idx : idx + bs])

        m_rc = rc.value
        assert rec_skm == pytest.approx(m_rc.item())

    for _ in range(10):
        _test(num_classes=2, threshold=0.5, multilabel=False, average="macro")
        _test(num_classes=2, threshold=0.5, multilabel=False, average="micro")


def test_multiclass_inputs():
    def _test(num_classes, threshold, multilabel, average):

        rc = Recall(num_classes=num_classes, threshold=threshold, multilabel=multilabel, average=average,)

        outputs = torch.randn(100, 4)
        targets = torch.randint(0, 4, size=(100,))

        bs = _BaseInputHandler(num_classes=num_classes, average=average, threshold=0.5, multilabel=multilabel,)

        np_outputs, np_targets = bs._compute(outputs=outputs, targets=targets)

        rc.accumulate(outputs=outputs, targets=targets)
        rec_val = rc.value
        assert rc.case_type == "multiclass"

        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=UndefinedMetricWarning)

            rec_skm = recall_score(np_targets.numpy(), np_outputs.numpy(), average=average)

        assert rec_skm == pytest.approx(rec_val.item())

        rc.reset()
        bs = 16
        iters = targets.shape[0] // bs + 1

        for i in range(iters):
            idx = i * bs
            rc.accumulate(outputs=outputs[idx : idx + bs], targets=targets[idx : idx + bs])

        rec_m = rc.value
        assert rec_skm == pytest.approx(rec_m.item())

    for _ in range(10):
        _test(num_classes=4, threshold=0.5, multilabel=False, average="macro")
        _test(num_classes=4, threshold=0.5, multilabel=False, average="micro")


def test_multilabel_inputs():
    def _test(num_classes, threshold, multilabel, average):

        rc = Recall(num_classes=num_classes, threshold=threshold, multilabel=multilabel, average=average,)

        outputs = torch.randn(100, 4)
        targets = torch.randint(0, 2, size=(100, 4))

        bs = _BaseInputHandler(num_classes=num_classes, average=average, threshold=0.5, multilabel=multilabel,)

        np_outputs, np_targets = bs._compute(outputs=outputs, targets=targets)

        rc.accumulate(outputs=outputs, targets=targets)
        rec_val = rc.value

        assert rc.case_type == "multilabel"
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", category=UndefinedMetricWarning)

            rec_skm = recall_score(np_targets, np_outputs, average=average)

        assert rec_skm == pytest.approx(rec_val.item())
        rc.reset()
        bs = 16
        iters = targets.shape[0] // bs + 1

        for i in range(iters):
            idx = i * bs
            rc.accumulate(
                outputs=outputs[idx : idx + bs], targets=targets[idx : idx + bs],
            )

        m_rc = rc.value
        assert rec_skm == pytest.approx(m_rc.item())

    for _ in range(10):
        _test(num_classes=4, threshold=0.5, multilabel=True, average="macro")
        _test(num_classes=4, threshold=0.5, multilabel=True, average="micro")

[((263, 284), 'torch.manual_seed', 'torch.manual_seed', (['(42)'], {}), '(42)\n', (280, 284), False, 'import torch\n'), ((387, 483), 'torchflare.metrics.recall_meter.Recall', 'Recall', ([], {'num_classes': 'num_classes', 'threshold': 'threshold', 'multilabel': 'multilabel', 'average': 'average'}), '(num_classes=num_classes, threshold=threshold, multilabel=multilabel,\n    average=average)\n', (393, 483), False, 'from torchflare.metrics.recall_meter import Recall\n'), ((499, 518), 'torch.randn', 'torch.randn', (['(100)', '(1)'], {}), '(100, 1)\n', (510, 518), False, 'import torch\n'), ((537, 569), 'torch.randint', 'torch.randint', (['(0)', '(2)'], {'size': '(100,)'}), '(0, 2, size=(100,))\n', (550, 569), False, 'import torch\n'), ((584, 685), 'torchflare.metrics.meters._BaseInputHandler', '_BaseInputHandler', ([], {'num_classes': 'num_classes', 'average': 'average', 'threshold': '(0.5)', 'multilabel': 'multilabel'}), '(num_classes=num_classes, average=average, threshold=0.5,\n    multilabel=multilabel)\n', (601, 685), False, 'from torchflare.metrics.meters import _BaseInputHandler\n'), ((1628, 1724), 'torchflare.metrics.recall_meter.Recall', 'Recall', ([], {'num_classes': 'num_classes', 'threshold': 'threshold', 'multilabel': 'multilabel', 'average': 'average'}), '(num_classes=num_classes, threshold=threshold, multilabel=multilabel,\n    average=average)\n', (1634, 1724), False, 'from torchflare.metrics.recall_meter import Recall\n'), ((1741, 1760), 'torch.randn', 'torch.randn', (['(100)', '(4)'], {}), '(100, 4)\n', (1752, 1760), False, 'import torch\n'), ((1779, 1811), 'torch.randint', 'torch.randint', (['(0)', '(4)'], {'size': '(100,)'}), '(0, 4, size=(100,))\n', (1792, 1811), False, 'import torch\n'), ((1826, 1927), 'torchflare.metrics.meters._BaseInputHandler', '_BaseInputHandler', ([], {'num_classes': 'num_classes', 'average': 'average', 'threshold': '(0.5)', 'multilabel': 'multilabel'}), '(num_classes=num_classes, average=average, threshold=0.5,\n    multilabel=multilabel)\n', (1843, 1927), False, 'from torchflare.metrics.meters import _BaseInputHandler\n'), ((2998, 3094), 'torchflare.metrics.recall_meter.Recall', 'Recall', ([], {'num_classes': 'num_classes', 'threshold': 'threshold', 'multilabel': 'multilabel', 'average': 'average'}), '(num_classes=num_classes, threshold=threshold, multilabel=multilabel,\n    average=average)\n', (3004, 3094), False, 'from torchflare.metrics.recall_meter import Recall\n'), ((3111, 3130), 'torch.randn', 'torch.randn', (['(100)', '(4)'], {}), '(100, 4)\n', (3122, 3130), False, 'import torch\n'), ((3149, 3183), 'torch.randint', 'torch.randint', (['(0)', '(2)'], {'size': '(100, 4)'}), '(0, 2, size=(100, 4))\n', (3162, 3183), False, 'import torch\n'), ((3198, 3299), 'torchflare.metrics.meters._BaseInputHandler', '_BaseInputHandler', ([], {'num_classes': 'num_classes', 'average': 'average', 'threshold': '(0.5)', 'multilabel': 'multilabel'}), '(num_classes=num_classes, average=average, threshold=0.5,\n    multilabel=multilabel)\n', (3215, 3299), False, 'from torchflare.metrics.meters import _BaseInputHandler\n'), ((2147, 2172), 'warnings.catch_warnings', 'warnings.catch_warnings', ([], {}), '()\n', (2170, 2172), False, 'import warnings\n'), ((2186, 2250), 'warnings.simplefilter', 'warnings.simplefilter', (['"""ignore"""'], {'category': 'UndefinedMetricWarning'}), "('ignore', category=UndefinedMetricWarning)\n", (2207, 2250), False, 'import warnings\n'), ((3519, 3544), 'warnings.catch_warnings', 'warnings.catch_warnings', ([], {}), '()\n', (3542, 3544), False, 'import warnings\n'), ((3558, 3622), 'warnings.simplefilter', 'warnings.simplefilter', (['"""ignore"""'], {'category': 'UndefinedMetricWarning'}), "('ignore', category=UndefinedMetricWarning)\n", (3579, 3622), False, 'import warnings\n'), ((3646, 3699), 'sklearn.metrics.recall_score', 'recall_score', (['np_targets', 'np_outputs'], {'average': 'average'}), '(np_targets, np_outputs, average=average)\n', (3658, 3699), False, 'from sklearn.metrics import recall_score\n')]