import warnings

import pytest
import numpy as np

from multiprocessing.managers import SharedMemoryManager

from tme.backends import MatchingBackend, NumpyFFTWBackend, BackendManager, backend

BACKENDS_TO_TEST = []
for backend_class in backend._BACKEND_REGISTRY.values():
    try:
        BACKENDS_TO_TEST.append(backend_class(device="cpu"))
    except ImportError:
        print(f"Couldn't import {backend_class}. Skipping...")

METHODS_TO_TEST = MatchingBackend.__abstractmethods__


class TestBackendManager:
    def setup_method(self):
        self.manager = BackendManager()

    def test_initialization(self):
        manager = BackendManager()
        backend_name = manager._backend_name
        assert f"<BackendManager: using {backend_name}>" == str(manager)

    def test_dir(self):
        _ = dir(self.manager)
        for method in METHODS_TO_TEST:
            assert hasattr(self.manager, method)

    def test_add_backend(self):
        self.manager.add_backend(backend_name="test", backend_class=NumpyFFTWBackend)

    def test_add_backend_error(self):
        class _Bar:
            def __init__(self):
                pass

        with pytest.raises(ValueError):
            self.manager.add_backend(backend_name="test", backend_class=_Bar)

    def test_change_backend_error(self):
        with pytest.raises(NotImplementedError):
            self.manager.change_backend(backend_name=None)

    def test_available_backends(self):
        available = self.manager.available_backends()
        assert isinstance(available, list)
        for be in available:
            assert be in self.manager._BACKEND_REGISTRY


class TestBackends:
    def setup_method(self):
        self.backend = NumpyFFTWBackend()
        self.x1 = np.random.rand(30, 30).astype(np.float32)
        self.x2 = np.random.rand(30, 30).astype(np.float32)

    def teardown_method(self):
        self.backend = None

    def test_initialization_errors(self):
        with pytest.raises(TypeError):
            _ = MatchingBackend()

    @pytest.mark.parametrize("backend", [type(x) for x in BACKENDS_TO_TEST])
    def test_initialization(self, backend):
        _ = backend()

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize(
        "method_name",
        ("add", "subtract", "multiply", "divide", "minimum", "maximum", "mod"),
    )
    def test_arithmetic_operations(self, method_name, backend):
        base = getattr(self.backend, method_name)(self.x1, self.x2)
        x1 = backend.to_backend_array(self.x1)
        x2 = backend.to_backend_array(self.x2)
        other = getattr(backend, method_name)(x1, x2)

        assert np.allclose(base, backend.to_numpy_array(other))

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize(
        "method_name", ("sum", "mean", "std", "max", "min", "unique")
    )
    @pytest.mark.parametrize("axis", ((0), (1)))
    def test_reduction_operations(self, method_name, backend, axis):
        base = getattr(self.backend, method_name)(self.x1, axis=axis)
        other = getattr(backend, method_name)(
            backend.to_backend_array(self.x1), axis=axis
        )
        # Account for bessel function correction in pytorch
        rtol = 0.01 if method_name != "std" else 0.5
        assert np.allclose(base, backend.to_numpy_array(other), rtol=rtol)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize(
        "method_name",
        ("sqrt", "square", "abs", "transpose", "tobytes", "size"),
    )
    def test_array_manipulation(self, method_name, backend):
        base = getattr(self.backend, method_name)(self.x1)
        other = getattr(backend, method_name)(backend.to_backend_array(self.x1))

        if isinstance(base, np.ndarray):
            assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)
        else:
            assert base == other

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("shape", ((10, 15), (10, 15, 20)))
    @pytest.mark.parametrize(
        "dtype", (("_float_dtype", "_complex_dtype", "_int_dtype"))
    )
    def test_zeros(self, shape, backend, dtype):
        dtype_base = getattr(self.backend, dtype)
        dtype_backend = getattr(backend, dtype)
        base = self.backend.zeros(shape, dtype=dtype_base)
        other = backend.zeros(shape, dtype=dtype_backend)
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("shape", ((10, 15), (10, 15, 20)))
    @pytest.mark.parametrize("fill_value", (-1, 0, 1))
    def test_full(self, shape, backend, fill_value):
        base = self.backend.full(shape, fill_value=fill_value)
        other = backend.full(shape, fill_value=fill_value)
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("power", (0.5, 1, 2))
    def test_power(self, backend, power):
        base = self.backend.power(self.x1, power)
        other = backend.power(backend.to_backend_array(self.x1), power)
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("shift", (-5, 0, 10))
    @pytest.mark.parametrize("axis", (0, 1))
    def test_roll(self, backend, shift, axis):
        base = self.backend.roll(self.x1, (shift,), (axis,))
        other = backend.roll(backend.to_backend_array(self.x1), (shift,), (axis,))
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("shape", ((10, 15), (10, 15, 20)))
    @pytest.mark.parametrize("fill_value", (-1, 0, 1))
    def test_fill(self, shape, backend, fill_value):
        base = self.backend.full(shape, fill_value=fill_value)
        other = backend.zeros(shape)
        other = backend.fill(other, fill_value)
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("min_distance", (1, 5, 10))
    def test_max_filter_coordinates(self, backend, min_distance):
        coordinates = backend.max_filter_coordinates(
            backend.to_backend_array(self.x1), min_distance=min_distance
        )
        if len(coordinates):
            assert coordinates.shape[1] == self.x1.ndim
        assert True

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize(
        "dtype", (("_float_dtype", "_complex_dtype", "_int_dtype"))
    )
    @pytest.mark.parametrize(
        "dtype_target", (("_int_dtype", "_complex_dtype", "_float_dtype"))
    )
    def test_astype(self, dtype, backend, dtype_target):
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            dtype_base = getattr(backend, dtype)
            dtype_target = getattr(backend, dtype_target)

            base = backend.zeros((20, 20, 20), dtype=dtype_base)
            arr = backend.astype(base, dtype_target)

            assert arr.dtype == dtype_target

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("N", (0, 15, 30))
    def test_arange(self, backend, N):
        base = self.backend.arange(N)
        other = getattr(backend, "arange")(
            N,
        )
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.1)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("return_inverse", (False, True))
    @pytest.mark.parametrize("return_counts", (False, True))
    @pytest.mark.parametrize("return_index", (False, True))
    def test_unique(self, backend, return_inverse, return_counts, return_index):
        base = self.backend.unique(
            self.x1,
            return_inverse=return_inverse,
            return_counts=return_counts,
            return_index=return_index,
        )
        other = backend.unique(
            backend.to_backend_array(self.x1),
            return_inverse=return_inverse,
            return_counts=return_counts,
            return_index=return_index,
        )
        if isinstance(base, tuple):
            base, other = tuple(base), tuple(other)
        for k in range(len(base)):
            print(
                k,
                base[k].shape,
                other[k].shape,
                return_inverse,
                return_counts,
                return_index,
            )
            assert np.allclose(
                base[k].ravel(), backend.to_numpy_array(other[k]).ravel(), rtol=0.1
            )

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("k", (0, 15, 30))
    def test_repeat(self, backend, k):
        base = self.backend.repeat(self.x1, k)
        other = backend.repeat(backend.to_backend_array(self.x1), k)
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.1)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("dim", (1, 3))
    @pytest.mark.parametrize("k", (0, 15, 30))
    def test_topk_indices(self, backend, k: int, dim: int):
        data = np.random.rand(*(50 for _ in range(dim)))
        base = self.backend.topk_indices(data, k)
        other = backend.topk_indices(backend.to_backend_array(data), k)

        for i in range(len(base)):
            np.allclose(
                base[i],
                backend.to_numpy_array(backend.to_backend_array(other[i])),
                rtol=0.1,
            )

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    def test_indices(self, backend):
        base = self.backend.indices(self.x1.shape)
        other = backend.indices(backend.to_backend_array(self.x1).shape)
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.1)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    def test_get_available_memory(self, backend):
        mem = backend.get_available_memory()
        assert isinstance(mem, int)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    def test_shared_memory(self, backend):
        shared_memory_handler = None
        base = backend.to_backend_array(self.x1)
        shared = backend.to_sharedarr(
            arr=base, shared_memory_handler=shared_memory_handler
        )
        arr = backend.from_sharedarr(shared)
        assert np.allclose(backend.to_numpy_array(arr), backend.to_numpy_array(base))

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    def test_shared_memory_managed(self, backend):
        with SharedMemoryManager() as shared_memory_handler:
            base = backend.to_backend_array(self.x1)
            shared = backend.to_sharedarr(
                arr=base, shared_memory_handler=shared_memory_handler
            )
            arr = backend.from_sharedarr(shared)
            assert np.allclose(
                backend.to_numpy_array(arr), backend.to_numpy_array(base)
            )

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("shape", ((10, 15, 100), (10, 15, 20)))
    @pytest.mark.parametrize("padval", (-1, 0, 1))
    def test_topleft_pad(self, backend, shape, padval):
        arr = np.random.rand(30, 30, 30)
        base = self.backend.topleft_pad(arr, shape=shape, padval=padval)
        other = backend.topleft_pad(
            backend.to_backend_array(arr), shape=shape, padval=padval
        )
        assert np.allclose(base, backend.to_numpy_array(other), rtol=0.01)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("fast_shape", ((10, 15, 100), (55, 23, 17)))
    def test_fft(self, backend, fast_shape):
        _, fast_shape, fast_ft_shape = backend.compute_convolution_shapes(
            fast_shape, (1 for _ in range(len(fast_shape)))
        )
        arr = np.random.rand(*fast_shape)
        out = np.zeros(fast_ft_shape)

        real_arr = backend.astype(backend.to_backend_array(arr), backend._float_dtype)
        complex_arr = backend.astype(
            backend.to_backend_array(out), backend._complex_dtype
        )

        backend.rfftn(
            backend.astype(backend.to_backend_array(arr), backend._float_dtype),
            complex_arr,
        )
        backend.irfftn(complex_arr, real_arr)
        assert np.allclose(arr, backend.to_numpy_array(real_arr), rtol=0.3)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("fast_shape", ((10, 15, 100), (55, 23, 17)))
    def test_fftn(self, backend, fast_shape):
        _, fast_shape, fast_ft_shape = backend.compute_convolution_shapes(
            fast_shape, (1 for _ in range(len(fast_shape)))
        )
        arr = np.random.rand(*fast_shape)
        out = np.zeros(fast_ft_shape)

        real_arr = backend.astype(backend.to_backend_array(arr), backend._float_dtype)
        complex_arr = backend.astype(
            backend.to_backend_array(out), backend._complex_dtype
        )

        complex_arr = backend.rfftn(
            backend.astype(backend.to_backend_array(arr), backend._float_dtype),
        )
        real_arr = backend.irfftn(complex_arr)
        assert np.allclose(arr, backend.to_numpy_array(real_arr), rtol=0.3)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    def test_compute_convolution_shapes(self, backend):
        base = self.backend.compute_convolution_shapes(self.x1.shape, self.x2.shape)
        other = backend.compute_convolution_shapes(self.x1.shape, self.x2.shape)

        assert base == other

    @pytest.mark.parametrize("dim", (2, 3))
    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("create_mask", (False, True))
    def test_rigid_transform(self, backend, dim, create_mask):
        shape = tuple(50 for _ in range(dim))
        arr = np.zeros(shape)
        if dim == 2:
            arr[20:25, 21:26] = 1
        elif dim == 3:
            arr[20:25, 21:26, 26:31] = 1

        from tme.rotations import get_rotation_matrices

        np.random.seed(42)
        rotation_matrix = get_rotation_matrices(
            dim=dim, angular_sampling=10, use_optimized_set=False
        )[-1]

        out = np.zeros_like(arr)
        out.setflags(write=True)

        arr_mask, out_mask = None, None
        if create_mask:
            arr_mask = np.multiply(np.random.rand(*arr.shape) > 0.5, 1.0)
            out_mask = np.zeros_like(arr_mask)
            out_mask.setflags(write=True)

        out, _ = NumpyFFTWBackend().rigid_transform(
            arr=arr,
            arr_mask=arr_mask,
            rotation_matrix=rotation_matrix,
            out=out,
            out_mask=out_mask,
            order=1,
            use_geometric_center=True,
        )

        arr = backend.to_backend_array(arr.copy())
        out_be = backend.to_backend_array(out.copy())
        if create_mask:
            arr_mask = backend.to_backend_array(arr_mask)
            out_mask = backend.to_backend_array(out_mask)

        rotation_matrix = backend.to_backend_array(rotation_matrix)

        out_be, _ = backend.rigid_transform(
            arr=arr,
            arr_mask=arr_mask,
            rotation_matrix=rotation_matrix,
            out=out_be,
            out_mask=out_mask,
            order=1,
            use_geometric_center=True,
        )
        out_be = backend.to_numpy_array(out_be)
        assert np.allclose(out, out_be, atol=0.3)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("create_mask", (False, True))
    def test_rigid_transform_identity(self, backend, create_mask):
        dim = 3
        shape = tuple(50 for _ in range(dim))

        arr = np.zeros(shape)
        arr[20:25, 21:26, 26:31] = 1

        rotation_matrix = backend.to_backend_array(np.eye(dim))
        out, _ = backend.rigid_transform(
            arr=backend.to_backend_array(arr),
            rotation_matrix=backend.to_backend_array(rotation_matrix),
            order=1,
            use_geometric_center=True,
        )
        assert np.allclose(out, arr, atol=0.01)

    @pytest.mark.parametrize("dim", (2, 3))
    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    @pytest.mark.parametrize("create_mask", (False, True))
    def test_rigid_transform_batch(self, backend, dim, create_mask):
        shape = tuple(30 for _ in range(dim))
        arr = np.random.rand(5, *shape)
        out = np.zeros_like(arr)

        arr_b = arr.copy()
        out_b = np.zeros_like(arr)

        rotation_matrix = np.eye(dim)
        rotation_matrix[0, 0] = -1

        arr_mask, out_mask = None, None
        if create_mask:
            arr_mask = np.multiply(np.random.rand(*arr.shape) > 0.5, 1.0)
            out_mask = np.zeros_like(arr_mask)
            arr_mask = backend.to_backend_array(arr_mask)
            out_mask = backend.to_backend_array(out_mask)

        arr = backend.to_backend_array(arr)
        out = backend.to_backend_array(out)

        rotation_matrix = backend.to_backend_array(rotation_matrix)
        backend.rigid_transform(
            arr=arr,
            arr_mask=arr_mask,
            rotation_matrix=rotation_matrix,
            out=out,
            out_mask=out_mask,
            batched=True,
            order=1,
        )

        arr_b = backend.to_backend_array(arr_b)
        out_b = backend.to_backend_array(out_b)

        for i in range(arr_b.shape[0]):
            backend.rigid_transform(
                arr=arr[i],
                arr_mask=arr_mask if arr_mask is None else arr_mask[i],
                rotation_matrix=rotation_matrix,
                out=out_b[i],
                out_mask=out_mask if out_mask is None else out_mask[i],
                order=1,
            )

        assert np.allclose(arr, arr_b)

    @pytest.mark.parametrize("backend", BACKENDS_TO_TEST)
    def test_datatype_bytes(self, backend):
        assert isinstance(backend.datatype_bytes(backend._float_dtype), int)
        assert isinstance(backend.datatype_bytes(backend._complex_dtype), int)
        assert isinstance(backend.datatype_bytes(backend._int_dtype), int)
