# 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")

[((12, 0, 12, 21), 'torch.manual_seed', 'torch.manual_seed', ({(12, 18, 12, 20): '(42)'}, {}), '(42)', False, 'import torch\n'), ((18, 13, 18, 106), 'torchflare.metrics.recall_meter.Recall', 'Recall', (), '', False, 'from torchflare.metrics.recall_meter import Recall\n'), ((19, 18, 19, 37), 'torch.randn', 'torch.randn', ({(19, 30, 19, 33): '100', (19, 35, 19, 36): '1'}, {}), '(100, 1)', False, 'import torch\n'), ((20, 18, 20, 50), 'torch.randint', 'torch.randint', (), '', False, 'import torch\n'), ((22, 13, 22, 111), 'torchflare.metrics.meters._BaseInputHandler', '_BaseInputHandler', (), '', False, 'from torchflare.metrics.meters import _BaseInputHandler\n'), ((52, 13, 52, 106), 'torchflare.metrics.recall_meter.Recall', 'Recall', (), '', False, 'from torchflare.metrics.recall_meter import Recall\n'), ((54, 18, 54, 37), 'torch.randn', 'torch.randn', ({(54, 30, 54, 33): '100', (54, 35, 54, 36): '4'}, {}), '(100, 4)', False, 'import torch\n'), ((55, 18, 55, 50), 'torch.randint', 'torch.randint', (), '', False, 'import torch\n'), ((57, 13, 57, 111), 'torchflare.metrics.meters._BaseInputHandler', '_BaseInputHandler', (), '', False, 'from torchflare.metrics.meters import _BaseInputHandler\n'), ((91, 13, 91, 106), 'torchflare.metrics.recall_meter.Recall', 'Recall', (), '', False, 'from torchflare.metrics.recall_meter import Recall\n'), ((93, 18, 93, 37), 'torch.randn', 'torch.randn', ({(93, 30, 93, 33): '100', (93, 35, 93, 36): '4'}, {}), '(100, 4)', False, 'import torch\n'), ((94, 18, 94, 52), 'torch.randint', 'torch.randint', (), '', False, 'import torch\n'), ((96, 13, 96, 111), 'torchflare.metrics.meters._BaseInputHandler', '_BaseInputHandler', (), '', False, 'from torchflare.metrics.meters import _BaseInputHandler\n'), ((65, 13, 65, 38), 'warnings.catch_warnings', 'warnings.catch_warnings', ({}, {}), '()', False, 'import warnings\n'), ((66, 12, 66, 76), 'warnings.simplefilter', 'warnings.simplefilter', (), '', False, 'import warnings\n'), ((104, 13, 104, 38), 'warnings.catch_warnings', 'warnings.catch_warnings', ({}, {}), '()', False, 'import warnings\n'), ((105, 12, 105, 76), 'warnings.simplefilter', 'warnings.simplefilter', (), '', False, 'import warnings\n'), ((107, 22, 107, 75), 'sklearn.metrics.recall_score', 'recall_score', (), '', False, 'from sklearn.metrics import recall_score\n')]