# Copyright The 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.
import sys

import numpy as np
import pytest
import torch
from torch import tensor
from torchmetrics.regression import MeanSquaredError, PearsonCorrCoef
from torchmetrics.utilities import check_forward_full_state_property, rank_zero_debug, rank_zero_info, rank_zero_warn
from torchmetrics.utilities.checks import _allclose_recursive
from torchmetrics.utilities.data import _bincount, _cumsum, _flatten, _flatten_dict, to_categorical, to_onehot
from torchmetrics.utilities.distributed import class_reduce, reduce
from torchmetrics.utilities.exceptions import TorchMetricsUserWarning
from torchmetrics.utilities.imports import _TORCH_GREATER_EQUAL_1_13


def test_prints():
    """Test that the different rank zero only functions works as expected."""
    rank_zero_debug("DEBUG")
    rank_zero_info("INFO")
    rank_zero_warn("WARN")


def test_reduce():
    """Test that reduction function works as expected and also raises error on wrong input."""
    start_tensor = torch.rand(50, 40, 30)

    assert torch.allclose(reduce(start_tensor, "elementwise_mean"), torch.mean(start_tensor))
    assert torch.allclose(reduce(start_tensor, "sum"), torch.sum(start_tensor))
    assert torch.allclose(reduce(start_tensor, "none"), start_tensor)

    with pytest.raises(ValueError, match="Reduction parameter unknown."):
        reduce(start_tensor, "error_reduction")


def test_class_reduce():
    """Test that class reduce function works as expected."""
    num = torch.randint(1, 10, (100,)).float()
    denom = torch.randint(10, 20, (100,)).float()
    weights = torch.randint(1, 100, (100,)).float()

    assert torch.allclose(class_reduce(num, denom, weights, "micro"), torch.sum(num) / torch.sum(denom))
    assert torch.allclose(class_reduce(num, denom, weights, "macro"), torch.mean(num / denom))
    assert torch.allclose(
        class_reduce(num, denom, weights, "weighted"), torch.sum(num / denom * (weights / torch.sum(weights)))
    )
    assert torch.allclose(class_reduce(num, denom, weights, "none"), num / denom)


def test_onehot():
    """Test that casting to onehot works as expected."""
    test_tensor = tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
    expected = torch.stack(
        [
            torch.cat([torch.eye(5, dtype=int), torch.zeros((5, 5), dtype=int)]),
            torch.cat([torch.zeros((5, 5), dtype=int), torch.eye(5, dtype=int)]),
        ]
    )

    assert test_tensor.shape == (2, 5)
    assert expected.shape == (2, 10, 5)

    onehot_classes = to_onehot(test_tensor, num_classes=10)
    onehot_no_classes = to_onehot(test_tensor)

    assert torch.allclose(onehot_classes, onehot_no_classes)

    assert onehot_classes.shape == expected.shape
    assert onehot_no_classes.shape == expected.shape

    assert torch.allclose(expected.to(onehot_no_classes), onehot_no_classes)
    assert torch.allclose(expected.to(onehot_classes), onehot_classes)


def test_to_categorical():
    """Test that casting to categorical works as expected."""
    test_tensor = torch.stack(
        [
            torch.cat([torch.eye(5, dtype=int), torch.zeros((5, 5), dtype=int)]),
            torch.cat([torch.zeros((5, 5), dtype=int), torch.eye(5, dtype=int)]),
        ]
    ).to(torch.float)

    expected = tensor([[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]])
    assert expected.shape == (2, 5)
    assert test_tensor.shape == (2, 10, 5)

    result = to_categorical(test_tensor)

    assert result.shape == expected.shape
    assert torch.allclose(result, expected.to(result.dtype))


def test_flatten_list():
    """Check that _flatten utility function works as expected."""
    inp = [[1, 2, 3], [4, 5], [6]]
    out = _flatten(inp)
    assert out == [1, 2, 3, 4, 5, 6]


def test_flatten_dict():
    """Check that _flatten_dict utility function works as expected."""
    inp = {"a": {"b": 1, "c": 2}, "d": 3}
    out = _flatten_dict(inp)
    assert out == {"b": 1, "c": 2, "d": 3}


@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires gpu")
def test_bincount():
    """Test that bincount works in deterministic setting on GPU."""
    torch.use_deterministic_algorithms(True)

    x = torch.randint(10, size=(100,))
    # uses custom implementation
    res1 = _bincount(x, minlength=10)

    torch.use_deterministic_algorithms(False)

    # uses torch.bincount
    res2 = _bincount(x, minlength=10)

    # explicit call to make sure, that res2 is not by accident using our manual implementation
    res3 = torch.bincount(x, minlength=10)

    # check for correctness
    assert torch.allclose(res1, res2)
    assert torch.allclose(res1, res3)


@pytest.mark.parametrize(("metric_class", "expected"), [(MeanSquaredError, False), (PearsonCorrCoef, True)])
def test_check_full_state_update_fn(capsys, metric_class, expected):
    """Test that the check function works as it should."""
    check_forward_full_state_property(
        metric_class=metric_class,
        input_args={"preds": torch.randn(1000), "target": torch.randn(1000)},
        num_update_to_compare=[10000],
        reps=5,
    )
    captured = capsys.readouterr()
    assert f"Recommended setting `full_state_update={expected}`" in captured.out


@pytest.mark.parametrize(
    ("inputs", "expected"),
    [
        ((torch.ones(2), torch.ones(2)), True),
        ((torch.rand(2), torch.rand(2)), False),
        (([torch.ones(2) for _ in range(2)], [torch.ones(2) for _ in range(2)]), True),
        (([torch.rand(2) for _ in range(2)], [torch.rand(2) for _ in range(2)]), False),
        (({f"{i}": torch.ones(2) for i in range(2)}, {f"{i}": torch.ones(2) for i in range(2)}), True),
        (({f"{i}": torch.rand(2) for i in range(2)}, {f"{i}": torch.rand(2) for i in range(2)}), False),
    ],
)
def test_recursive_allclose(inputs, expected):
    """Test the recursive allclose works as expected."""
    res = _allclose_recursive(*inputs)
    assert res == expected


@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU")
@pytest.mark.xfail(sys.platform == "win32", reason="test will only fail on non-windows systems")
@pytest.mark.skipif(
    not _TORCH_GREATER_EQUAL_1_13, reason="earlier versions was silently non-deterministic, even in deterministic mode"
)
def test_cumsum_still_not_supported():
    """Make sure that cumsum on gpu and deterministic mode still fails.

    If this test begins to passes, it means newer Pytorch versions support this and we can drop internal support.
    """
    torch.use_deterministic_algorithms(True)
    with pytest.raises(RuntimeError, match="cumsum_cuda_kernel does not have a deterministic implementation.*"):
        torch.arange(10).float().cuda().cumsum(0)
    torch.use_deterministic_algorithms(False)


@pytest.mark.skipif(not torch.cuda.is_available(), reason="test requires GPU")
def test_custom_cumsum():
    """Test custom cumsum implementation."""
    torch.use_deterministic_algorithms(True)
    x = torch.arange(100).float().cuda()
    if sys.platform != "win32":
        with pytest.warns(
            TorchMetricsUserWarning, match="You are trying to use a metric in deterministic mode on GPU that.*"
        ):
            res = _cumsum(x, dim=0).cpu()
    else:
        res = _cumsum(x, dim=0).cpu()
    res2 = np.cumsum(x.cpu(), axis=0)
    assert torch.allclose(res, res2)
