"""End-to-end tests for realistic oil spill segmentation scenarios."""

import pytest
import torch
from sdf_loss import SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss


# Skip GPU tests if CUDA is not available
cuda_available = torch.cuda.is_available()
skip_no_cuda = pytest.mark.skipif(not cuda_available, reason="CUDA not available")


def create_oil_spill_like_masks(batch_size, height, width, device="cpu"):
    """Create realistic oil spill-like segmentation masks.

    Oil spills typically have:
    - Irregular blob shapes
    - Multiple disconnected regions
    - Varying sizes (small to large coverage)
    """
    target = torch.zeros(batch_size, 1, height, width, device=device)

    for b in range(batch_size):
        # Create 1-3 oil spill blobs per image
        num_blobs = torch.randint(1, 4, (1,)).item()
        for _ in range(num_blobs):
            # Random blob center and size
            cx = torch.randint(height // 4, 3 * height // 4, (1,)).item()
            cy = torch.randint(width // 4, 3 * width // 4, (1,)).item()
            rx = torch.randint(height // 10, height // 4, (1,)).item()
            ry = torch.randint(width // 10, width // 4, (1,)).item()

            # Create elliptical blob
            y, x = torch.meshgrid(
                torch.arange(height, device=device),
                torch.arange(width, device=device),
                indexing="ij",
            )
            mask = ((x - cx) ** 2 / (rx**2 + 1) + (y - cy) ** 2 / (ry**2 + 1)) < 1
            target[b, 0] = torch.maximum(target[b, 0], mask.float())

    return target


def create_prediction_from_target(target, noise_level=0.3):
    """Create a noisy prediction from ground truth (simulates model output)."""
    # Add noise and convert to logits
    noise = torch.randn_like(target) * noise_level
    pred_prob = torch.clamp(target + noise, 0, 1)
    # Convert to logits (inverse sigmoid)
    pred_prob = torch.clamp(pred_prob, 0.01, 0.99)
    pred_logits = torch.log(pred_prob / (1 - pred_prob))
    return pred_logits


class TestRealisticOilSpillScenarios:
    """Tests simulating realistic oil spill segmentation training."""

    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_typical_training_batch(self, loss_class):
        """Test with typical training batch size and image dimensions."""
        # Typical satellite image patch size for segmentation
        batch_size = 8
        height, width = 256, 256

        target = create_oil_spill_like_masks(batch_size, height, width)
        pred = create_prediction_from_target(target)

        loss_fn = loss_class()
        loss = loss_fn(pred, target)

        assert torch.isfinite(loss), f"{loss_class.__name__} should produce finite loss"
        assert loss.item() >= 0, (
            f"{loss_class.__name__} should produce non-negative loss"
        )

    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_large_image(self, loss_class):
        """Test with larger image sizes (e.g., full satellite tiles)."""
        batch_size = 2
        height, width = 512, 512

        target = create_oil_spill_like_masks(batch_size, height, width)
        pred = create_prediction_from_target(target)

        loss_fn = loss_class()
        loss = loss_fn(pred, target)

        assert torch.isfinite(loss), f"{loss_class.__name__} should handle large images"

    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_sparse_oil_spill(self, loss_class):
        """Test with sparse oil spill (small coverage area)."""
        batch_size = 4
        height, width = 256, 256

        # Create sparse masks (small oil spills)
        target = torch.zeros(batch_size, 1, height, width)
        for b in range(batch_size):
            # Single small blob
            cx, cy = height // 2, width // 2
            r = 10  # Small radius
            y, x = torch.meshgrid(
                torch.arange(height), torch.arange(width), indexing="ij"
            )
            target[b, 0] = ((x - cx) ** 2 + (y - cy) ** 2 < r**2).float()

        pred = create_prediction_from_target(target)

        loss_fn = loss_class()
        loss = loss_fn(pred, target)

        assert torch.isfinite(loss), f"{loss_class.__name__} should handle sparse masks"

    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_dense_oil_spill(self, loss_class):
        """Test with dense oil spill (large coverage area)."""
        batch_size = 4
        height, width = 256, 256

        # Create dense masks (large oil spills covering most of image)
        target = torch.zeros(batch_size, 1, height, width)
        target[:, :, 20:-20, 20:-20] = 1.0

        pred = create_prediction_from_target(target)

        loss_fn = loss_class()
        loss = loss_fn(pred, target)

        assert torch.isfinite(loss), f"{loss_class.__name__} should handle dense masks"

    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_no_oil_spill(self, loss_class):
        """Test with no oil spill in image (all background)."""
        batch_size = 4
        height, width = 256, 256

        target = torch.zeros(batch_size, 1, height, width)
        pred = torch.randn(batch_size, 1, height, width) * 0.1 - 2  # Low logits

        loss_fn = loss_class()
        loss = loss_fn(pred, target)

        assert torch.isfinite(loss), f"{loss_class.__name__} should handle empty masks"

    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_training_loop_simulation(self, loss_class):
        """Simulate a few training iterations."""
        batch_size = 4
        height, width = 128, 128
        num_iterations = 5

        # Simple "model" - just learnable parameters
        model_output = torch.randn(batch_size, 1, height, width, requires_grad=True)
        optimizer = torch.optim.Adam([model_output], lr=0.1)

        loss_fn = loss_class()
        target = create_oil_spill_like_masks(batch_size, height, width)

        losses = []
        for _ in range(num_iterations):
            optimizer.zero_grad()
            loss = loss_fn(model_output, target)
            loss.backward()
            optimizer.step()
            losses.append(loss.item())

        # All losses should be finite
        assert all(torch.isfinite(torch.tensor(loss_val)) for loss_val in losses), (
            f"{loss_class.__name__} should produce finite losses throughout training"
        )


class TestGPUTraining:
    """Tests for GPU training scenarios."""

    @skip_no_cuda
    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_gpu_training_batch(self, loss_class):
        """Test training on GPU with typical batch."""
        device = "cuda"
        batch_size = 8
        height, width = 256, 256

        target = create_oil_spill_like_masks(batch_size, height, width, device=device)
        pred = create_prediction_from_target(target)

        loss_fn = loss_class()
        loss = loss_fn(pred, target)

        assert loss.device.type == "cuda", "Loss should be on GPU"
        assert torch.isfinite(loss), "Loss should be finite on GPU"

    @skip_no_cuda
    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_gpu_gradient_flow(self, loss_class):
        """Test that gradients flow correctly on GPU."""
        device = "cuda"
        batch_size = 4
        height, width = 128, 128

        target = create_oil_spill_like_masks(batch_size, height, width, device=device)
        pred = torch.randn(
            batch_size, 1, height, width, device=device, requires_grad=True
        )

        loss_fn = loss_class()
        loss = loss_fn(pred, target)
        loss.backward()

        assert pred.grad is not None, "Gradients should flow on GPU"
        assert pred.grad.device.type == "cuda", "Gradients should be on GPU"
        assert torch.isfinite(pred.grad).all(), "Gradients should be finite on GPU"

    @skip_no_cuda
    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_no_cpu_transfer_during_forward(self, loss_class):
        """Verify no CPU transfers happen during forward pass."""
        device = "cuda"
        batch_size = 4
        height, width = 128, 128

        target = create_oil_spill_like_masks(batch_size, height, width, device=device)
        pred = torch.randn(batch_size, 1, height, width, device=device)

        # Clear CUDA cache and synchronize
        torch.cuda.empty_cache()
        torch.cuda.synchronize()

        loss_fn = loss_class()

        # The forward pass should not trigger any CPU operations
        # We can't directly detect CPU transfers, but we verify the output stays on GPU
        loss = loss_fn(pred, target)

        assert loss.device.type == "cuda", "Loss computation should stay on GPU"

    @skip_no_cuda
    def test_gpu_training_loop_memory_stable(self):
        """Test that memory usage is stable during training loop."""
        device = "cuda"
        batch_size = 4
        height, width = 128, 128
        num_iterations = 10

        torch.cuda.empty_cache()
        torch.cuda.synchronize()

        model_output = torch.randn(
            batch_size, 1, height, width, device=device, requires_grad=True
        )
        optimizer = torch.optim.Adam([model_output], lr=0.1)

        loss_fn = DiSCoLoss()
        target = create_oil_spill_like_masks(batch_size, height, width, device=device)

        # Warm up
        for _ in range(3):
            optimizer.zero_grad()
            loss = loss_fn(model_output, target)
            loss.backward()
            optimizer.step()

        torch.cuda.synchronize()
        memory_after_warmup = torch.cuda.memory_allocated()

        # Run more iterations
        for _ in range(num_iterations):
            optimizer.zero_grad()
            loss = loss_fn(model_output, target)
            loss.backward()
            optimizer.step()

        torch.cuda.synchronize()
        memory_after_training = torch.cuda.memory_allocated()

        # Memory should not grow significantly (allow 10% tolerance)
        memory_growth = memory_after_training - memory_after_warmup
        assert memory_growth < memory_after_warmup * 0.1, (
            f"Memory grew by {memory_growth} bytes during training loop"
        )


class TestMixedPrecision:
    """Tests for mixed precision (fp16) training compatibility."""

    @skip_no_cuda
    @pytest.mark.parametrize(
        "loss_class", [SDFWeightedBCELoss, SDFWeightedDiceLoss, DiSCoLoss]
    )
    def test_autocast_compatibility(self, loss_class):
        """Test compatibility with torch.cuda.amp.autocast."""
        device = "cuda"
        batch_size = 4
        height, width = 128, 128

        target = create_oil_spill_like_masks(batch_size, height, width, device=device)
        pred = torch.randn(
            batch_size, 1, height, width, device=device, requires_grad=True
        )

        loss_fn = loss_class()

        with torch.cuda.amp.autocast():
            loss = loss_fn(pred, target)

        assert torch.isfinite(loss), f"{loss_class.__name__} should work with autocast"

        # Backward pass outside autocast
        loss.backward()
        assert pred.grad is not None, "Gradients should flow with autocast"


if __name__ == "__main__":
    pytest.main([__file__, "-v"])
