"""
Tests for device-specific operations (CUDA, MPS).

These tests verify that operations work correctly when tensors
are on accelerator devices, including proper device handling
for scatter operations and reductions.
"""

import pytest
import numpy as np


# Check device availability
def cuda_available():
    try:
        import torch
        return torch.cuda.is_available()
    except ImportError:
        return False


def mps_available():
    try:
        import torch
        return torch.backends.mps.is_available()
    except (ImportError, AttributeError):
        return False


# Skip markers
requires_cuda = pytest.mark.skipif(
    not cuda_available(),
    reason="CUDA not available"
)

requires_mps = pytest.mark.skipif(
    not mps_available(),
    reason="MPS not available"
)


class TestDeviceOperations:
    """Test operations on different devices."""

    @pytest.fixture
    def polymer_torch(self):
        """Create a test polymer with torch backend."""
        from ciffy import load
        return load("tests/data/9GCM.cif", backend="torch")

    @requires_cuda
    def test_to_cuda(self, polymer_torch):
        """Test moving polymer to CUDA device."""
        p_cuda = polymer_torch.to("cuda")

        assert p_cuda.coordinates.device.type == "cuda"
        assert p_cuda.atoms.device.type == "cuda"
        assert p_cuda.elements.device.type == "cuda"
        assert p_cuda.sequence.device.type == "cuda"

    @requires_mps
    def test_to_mps(self, polymer_torch):
        """Test moving polymer to MPS device."""
        p_mps = polymer_torch.to("mps")

        assert p_mps.coordinates.device.type == "mps"
        assert p_mps.atoms.device.type == "mps"
        assert p_mps.elements.device.type == "mps"
        assert p_mps.sequence.device.type == "mps"

    @requires_cuda
    def test_reduce_on_cuda(self, polymer_torch):
        """Test reduction operations on CUDA."""
        from ciffy import Scale

        p_cuda = polymer_torch.to("cuda")

        # Test reduce (per-atom to per-chain)
        means = p_cuda.reduce(p_cuda.coordinates, Scale.CHAIN)
        assert means.device.type == "cuda"
        assert means.shape[0] == p_cuda.size(Scale.CHAIN)

    @requires_mps
    def test_reduce_on_mps(self, polymer_torch):
        """Test reduction operations on MPS."""
        from ciffy import Scale

        p_mps = polymer_torch.to("mps")

        # Test reduce (per-atom to per-chain)
        means = p_mps.reduce(p_mps.coordinates, Scale.CHAIN)
        assert means.device.type == "mps"
        assert means.shape[0] == p_mps.size(Scale.CHAIN)

    @requires_cuda
    def test_center_on_cuda(self, polymer_torch):
        """Test centering on CUDA (uses reduce internally)."""
        from ciffy import Scale

        p_cuda = polymer_torch.to("cuda")
        centered, _ = p_cuda.center(Scale.MOLECULE)

        assert centered.coordinates.device.type == "cuda"

    @requires_mps
    def test_center_on_mps(self, polymer_torch):
        """Test centering on MPS (uses reduce internally)."""
        from ciffy import Scale

        p_mps = polymer_torch.to("mps")
        centered, _ = p_mps.center(Scale.MOLECULE)

        assert centered.coordinates.device.type == "mps"

    @requires_cuda
    def test_rmsd_on_cuda(self, polymer_torch):
        """Test RMSD calculation on CUDA."""
        import ciffy

        p_cuda = polymer_torch.to("cuda")

        # Calculate RMSD against self (should be 0)
        rmsd = ciffy.rmsd(p_cuda, p_cuda, ciffy.MOLECULE)

        # Result should be on CUDA and close to 0
        assert rmsd.device.type == "cuda"
        assert rmsd.item() < 1e-5

    @requires_mps
    def test_rmsd_on_mps(self, polymer_torch):
        """Test RMSD calculation on MPS.

        Note: MPS doesn't support SVD operations, so RMSD (which uses Kabsch
        alignment with SVD) will fail. This test verifies the limitation is
        handled gracefully.
        """
        import ciffy
        import torch

        p_mps = polymer_torch.to("mps")

        # MPS doesn't support SVD, so RMSD will fail
        # Use PYTORCH_ENABLE_MPS_FALLBACK=1 env var to enable CPU fallback
        with pytest.raises(NotImplementedError, match="MPS device"):
            ciffy.rmsd(p_mps, p_mps, ciffy.MOLECULE)


class TestMixedDeviceHandling:
    """Test that operations handle mixed-device scenarios gracefully."""

    @pytest.fixture
    def polymer_torch(self):
        """Create a test polymer with torch backend."""
        from ciffy import load
        return load("tests/data/9GCM.cif", backend="torch")

    @requires_cuda
    def test_scatter_with_cpu_index_cuda_features(self):
        """Test scatter operations handle CPU index with CUDA features."""
        import torch
        from ciffy.backend.torch_ops import scatter_sum, scatter_mean

        # Create features on CUDA, index on CPU
        features = torch.randn(10, 3, device="cuda")
        index = torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2, 2])  # CPU

        # This should work (index automatically moved to CUDA)
        result = scatter_sum(features, index, dim_size=3)
        assert result.device.type == "cuda"

        result = scatter_mean(features, index, dim_size=3)
        assert result.device.type == "cuda"

    @requires_mps
    def test_scatter_with_cpu_index_mps_features(self):
        """Test scatter operations handle CPU index with MPS features."""
        import torch
        from ciffy.backend.torch_ops import scatter_sum, scatter_mean

        # Create features on MPS, index on CPU
        features = torch.randn(10, 3, device="mps")
        index = torch.tensor([0, 0, 1, 1, 1, 2, 2, 2, 2, 2])  # CPU

        # This should work (index automatically moved to MPS)
        result = scatter_sum(features, index, dim_size=3)
        assert result.device.type == "mps"

        result = scatter_mean(features, index, dim_size=3)
        assert result.device.type == "mps"

    @requires_cuda
    def test_reduce_with_mismatched_sizes_device(self, polymer_torch):
        """Test reduce works even if internal sizes tensor is on different device."""
        import torch
        from ciffy import Scale

        # Get polymer on CUDA
        p_cuda = polymer_torch.to("cuda")

        # Manually move coordinates to CPU but keep sizes on CUDA
        # (This simulates a potential edge case)
        coords_cpu = p_cuda.coordinates.cpu()

        # Create reduction with CPU features - should still work
        # because create_reduction_index gets device from features
        from ciffy.operations.reduction import create_reduction_index

        sizes = p_cuda._sizes[Scale.CHAIN]
        assert sizes.device.type == "cuda"

        # Pass CPU device explicitly
        index = create_reduction_index(
            p_cuda.size(Scale.CHAIN),
            sizes,
            device=torch.device("cpu")
        )
        assert index.device.type == "cpu"


class TestScatterOperations:
    """Test scatter operations directly."""

    @requires_cuda
    def test_scatter_sum_cuda(self):
        """Test scatter_sum on CUDA."""
        import torch
        from ciffy.backend.torch_ops import scatter_sum

        features = torch.tensor([[1., 2.], [3., 4.], [5., 6.]], device="cuda")
        index = torch.tensor([0, 0, 1], device="cuda")

        result = scatter_sum(features, index, dim_size=2)

        expected = torch.tensor([[4., 6.], [5., 6.]], device="cuda")
        assert torch.allclose(result, expected)

    @requires_cuda
    def test_scatter_mean_cuda(self):
        """Test scatter_mean on CUDA."""
        import torch
        from ciffy.backend.torch_ops import scatter_mean

        features = torch.tensor([[1., 2.], [3., 4.], [5., 6.]], device="cuda")
        index = torch.tensor([0, 0, 1], device="cuda")

        result = scatter_mean(features, index, dim_size=2)

        expected = torch.tensor([[2., 3.], [5., 6.]], device="cuda")
        assert torch.allclose(result, expected)

    @requires_mps
    def test_scatter_sum_mps(self):
        """Test scatter_sum on MPS."""
        import torch
        from ciffy.backend.torch_ops import scatter_sum

        features = torch.tensor([[1., 2.], [3., 4.], [5., 6.]], device="mps")
        index = torch.tensor([0, 0, 1], device="mps")

        result = scatter_sum(features, index, dim_size=2)

        expected = torch.tensor([[4., 6.], [5., 6.]], device="mps")
        assert torch.allclose(result.cpu(), expected.cpu())

    @requires_mps
    def test_scatter_mean_mps(self):
        """Test scatter_mean on MPS."""
        import torch
        from ciffy.backend.torch_ops import scatter_mean

        features = torch.tensor([[1., 2.], [3., 4.], [5., 6.]], device="mps")
        index = torch.tensor([0, 0, 1], device="mps")

        result = scatter_mean(features, index, dim_size=2)

        expected = torch.tensor([[2., 3.], [5., 6.]], device="mps")
        assert torch.allclose(result.cpu(), expected.cpu())
