"""Test mps-deform-conv: FP32, FP16, BF16 support and validation"""

import torch
import time
import pytest

# Build and load the extension
print("Loading mps_deform_conv...")
from mps_deform_conv import deform_conv2d, DeformConv2d, ModulatedDeformConv2d, is_available

print(f"MPS available: {is_available()}")


# =============================================================================
# Validation Tests
# =============================================================================

def test_stride_validation():
    """Test that stride=0 raises ValueError"""
    torch.manual_seed(42)
    B, C_in, H, W = 1, 4, 8, 8
    C_out, K = 4, 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=torch.float32)
    offset = torch.randn(B, 2 * K * K, H - K + 1, W - K + 1, device='mps', dtype=torch.float32) * 0.1

    try:
        deform_conv2d(input, offset, weight, stride=(0, 1))
        assert False, "Should have raised ValueError for stride=0"
    except ValueError as e:
        assert "stride must be positive" in str(e)
        print(f"  stride=0 correctly rejected: {e}")
        return True
    return False


def test_dilation_validation():
    """Test that dilation=0 raises ValueError"""
    torch.manual_seed(42)
    B, C_in, H, W = 1, 4, 8, 8
    C_out, K = 4, 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=torch.float32)
    offset = torch.randn(B, 2 * K * K, H - K + 1, W - K + 1, device='mps', dtype=torch.float32) * 0.1

    try:
        deform_conv2d(input, offset, weight, dilation=(0, 1))
        assert False, "Should have raised ValueError for dilation=0"
    except ValueError as e:
        assert "dilation must be positive" in str(e)
        print(f"  dilation=0 correctly rejected: {e}")
        return True
    return False


def test_negative_padding_validation():
    """Test that negative padding raises ValueError"""
    torch.manual_seed(42)
    B, C_in, H, W = 1, 4, 8, 8
    C_out, K = 4, 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=torch.float32)
    offset = torch.randn(B, 2 * K * K, H - K + 1, W - K + 1, device='mps', dtype=torch.float32) * 0.1

    try:
        deform_conv2d(input, offset, weight, padding=(-1, 0))
        assert False, "Should have raised ValueError for negative padding"
    except ValueError as e:
        assert "padding must be non-negative" in str(e)
        print(f"  negative padding correctly rejected: {e}")
        return True
    return False


def test_groups_validation():
    """Test that invalid channel/group combinations are rejected"""
    torch.manual_seed(42)
    B, C_in, H, W = 1, 4, 8, 8
    C_out, K = 4, 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    # Weight with mismatched in_channels (3 doesn't divide 4)
    weight = torch.randn(C_out, 3, K, K, device='mps', dtype=torch.float32)
    offset = torch.randn(B, 2 * K * K, H - K + 1, W - K + 1, device='mps', dtype=torch.float32) * 0.1

    try:
        deform_conv2d(input, offset, weight)
        assert False, "Should have raised ValueError for invalid groups"
    except ValueError as e:
        assert "divisible" in str(e).lower()
        print(f"  invalid groups correctly rejected: {e}")
        return True
    return False


def test_device_mismatch():
    """Test that device mismatch raises ValueError"""
    torch.manual_seed(42)
    B, C_in, H, W = 1, 4, 8, 8
    C_out, K = 4, 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    weight = torch.randn(C_out, C_in, K, K, device='cpu', dtype=torch.float32)  # CPU!
    offset = torch.randn(B, 2 * K * K, H - K + 1, W - K + 1, device='mps', dtype=torch.float32) * 0.1

    try:
        deform_conv2d(input, offset, weight)
        assert False, "Should have raised ValueError for device mismatch"
    except ValueError as e:
        assert "device" in str(e).lower()
        print(f"  device mismatch correctly rejected: {e}")
        return True
    return False


def test_invalid_output_dimensions():
    """Test that invalid output dimensions raise ValueError"""
    torch.manual_seed(42)
    B, C_in, H, W = 1, 4, 4, 4  # Small input
    C_out, K = 4, 5  # Large kernel

    input = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=torch.float32)
    # This would give negative output dimensions
    offset = torch.randn(B, 2 * K * K, 1, 1, device='mps', dtype=torch.float32) * 0.1

    try:
        # stride=2 with 4x4 input and 5x5 kernel gives negative output
        deform_conv2d(input, offset, weight, stride=(2, 2))
        # Might not fail if offset shape matches expected output
        print("  Output dimension test: passed or offset shape matched")
        return True
    except ValueError as e:
        print(f"  Invalid output dimensions correctly rejected: {e}")
        return True


def test_bf16_precision_acceptable():
    """Test that BF16 precision loss is within acceptable bounds.

    BF16 has only 7 bits of mantissa vs FP32's 23 bits, so there WILL be
    precision loss. This test verifies the loss is within acceptable bounds
    for typical deep learning use cases (< 1% relative error).
    """
    torch.manual_seed(42)
    B, C_in, H, W = 2, 32, 16, 16
    C_out, K = 32, 3

    # Create FP32 reference tensors
    input_fp32 = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    weight_fp32 = torch.randn(C_out, C_in, K, K, device='mps', dtype=torch.float32)
    out_h = H - K + 1
    out_w = W - K + 1
    offset_fp32 = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=torch.float32) * 0.1

    # FP32 reference output
    import mps_deform_conv
    mps_deform_conv._bf16_warned = True  # Suppress warning for test
    output_fp32 = deform_conv2d(input_fp32, offset_fp32, weight_fp32)

    # BF16 output (converted from same FP32 inputs)
    input_bf16 = input_fp32.to(torch.bfloat16)
    weight_bf16 = weight_fp32.to(torch.bfloat16)
    offset_bf16 = offset_fp32.to(torch.bfloat16)

    output_bf16 = deform_conv2d(input_bf16, offset_bf16, weight_bf16)

    # Compute precision metrics
    diff = (output_fp32 - output_bf16.to(torch.float32)).abs()
    max_diff = diff.max().item()
    mean_diff = diff.mean().item()
    rel_error = (diff / (output_fp32.abs() + 1e-6)).mean().item() * 100

    print(f"  BF16 precision: max_diff={max_diff:.6f}, mean_diff={mean_diff:.6f}, rel_error={rel_error:.2f}%")

    # BF16 has 7-bit mantissa (vs FP32's 23-bit), so ~2-3% relative error is expected.
    # We use 3% threshold to catch severe precision issues while accepting normal BF16 behavior.
    threshold = 3.0  # 3% relative error threshold (typical for BF16)
    ok = rel_error < threshold

    if not ok:
        print(f"  FAIL: BF16 relative error {rel_error:.2f}% exceeds {threshold}% threshold")
    else:
        print(f"  PASS: BF16 relative error {rel_error:.2f}% within {threshold}% threshold")

    return ok


def test_forward(dtype, name):
    """Test forward pass"""
    torch.manual_seed(42)

    B, C_in, H, W = 2, 64, 32, 32
    C_out = 64
    K = 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=dtype)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=dtype)

    out_h = H - K + 1
    out_w = W - K + 1
    offset = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=dtype) * 0.1

    output = deform_conv2d(input, offset, weight, stride=(1, 1), padding=(0, 0))

    has_nan = torch.isnan(output).any().item()
    has_inf = torch.isinf(output).any().item()
    ok = not has_nan and not has_inf

    print(f"  {name} forward: shape={output.shape}, dtype={output.dtype}, ok={ok}")
    return ok

def test_forward_with_mask(dtype, name):
    """Test forward pass with DCNv2 mask"""
    torch.manual_seed(42)

    B, C_in, H, W = 2, 32, 16, 16
    C_out = 32
    K = 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=dtype)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=dtype)

    out_h = H - K + 1
    out_w = W - K + 1
    offset = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=dtype) * 0.1
    mask = torch.sigmoid(torch.randn(B, K * K, out_h, out_w, device='mps', dtype=dtype))

    output = deform_conv2d(input, offset, weight, stride=(1, 1), padding=(0, 0), mask=mask)

    has_nan = torch.isnan(output).any().item()
    has_inf = torch.isinf(output).any().item()
    ok = not has_nan and not has_inf

    print(f"  {name} forward+mask: shape={output.shape}, ok={ok}")
    return ok

def test_backward(dtype, name):
    """Test backward pass"""
    torch.manual_seed(42)

    B, C_in, H, W = 2, 32, 16, 16
    C_out = 32
    K = 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=dtype, requires_grad=True)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=dtype, requires_grad=True)

    out_h = H - K + 1
    out_w = W - K + 1
    offset = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=dtype) * 0.1
    offset.requires_grad_(True)

    output = deform_conv2d(input, offset, weight, stride=(1, 1), padding=(0, 0))
    loss = output.sum()
    loss.backward()

    grad_input_ok = input.grad is not None and not torch.isnan(input.grad).any()
    grad_weight_ok = weight.grad is not None and not torch.isnan(weight.grad).any()
    grad_offset_ok = offset.grad is not None and not torch.isnan(offset.grad).any()
    ok = grad_input_ok and grad_weight_ok and grad_offset_ok

    print(f"  {name} backward: input={grad_input_ok}, weight={grad_weight_ok}, offset={grad_offset_ok}")
    return ok

def test_module(dtype, name):
    """Test DeformConv2d module"""
    torch.manual_seed(42)

    B, C_in, H, W = 2, 32, 16, 16
    C_out = 32
    K = 3

    module = DeformConv2d(C_in, C_out, K, padding=1).to('mps').to(dtype)
    input = torch.randn(B, C_in, H, W, device='mps', dtype=dtype)

    out_h, out_w = H, W  # padding=1 keeps size
    offset = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=dtype) * 0.1

    output = module(input, offset)

    ok = not torch.isnan(output).any() and not torch.isinf(output).any()
    print(f"  {name} DeformConv2d module: shape={output.shape}, ok={ok}")
    return ok

def test_modulated_module(dtype, name):
    """Test ModulatedDeformConv2d module"""
    torch.manual_seed(42)

    B, C_in, H, W = 2, 32, 16, 16
    C_out = 32
    K = 3

    module = ModulatedDeformConv2d(C_in, C_out, K, padding=1).to('mps').to(dtype)
    input = torch.randn(B, C_in, H, W, device='mps', dtype=dtype)

    output = module(input)

    ok = not torch.isnan(output).any() and not torch.isinf(output).any()
    print(f"  {name} ModulatedDeformConv2d: shape={output.shape}, ok={ok}")
    return ok

def test_stride_dilation(dtype, name):
    """Test with stride and dilation"""
    torch.manual_seed(42)

    B, C_in, H, W = 2, 32, 32, 32
    C_out = 32
    K = 3
    stride = 2
    dilation = 2
    padding = dilation

    input = torch.randn(B, C_in, H, W, device='mps', dtype=dtype)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=dtype)

    # Calculate output size with dilation
    kernel_d = dilation * (K - 1) + 1
    out_h = (H + 2 * padding - kernel_d) // stride + 1
    out_w = (W + 2 * padding - kernel_d) // stride + 1

    offset = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=dtype) * 0.1

    output = deform_conv2d(input, offset, weight, stride=(stride, stride),
                           padding=(padding, padding), dilation=(dilation, dilation))

    ok = not torch.isnan(output).any() and not torch.isinf(output).any()
    print(f"  {name} stride={stride}, dilation={dilation}: shape={output.shape}, ok={ok}")
    return ok

def test_gradient_correctness():
    """Test gradient correctness via numerical differentiation"""
    torch.manual_seed(42)

    B, C_in, H, W = 1, 4, 6, 6
    C_out = 4
    K = 3
    out_h = H - K + 1
    out_w = W - K + 1
    eps = 1e-3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32, requires_grad=True)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=torch.float32, requires_grad=True)
    offset = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=torch.float32) * 0.1
    offset.requires_grad_(True)

    def forward():
        return deform_conv2d(input, offset, weight, stride=(1, 1), padding=(0, 0))

    def check_grad(name, param):
        # Analytical gradient
        input.grad = weight.grad = offset.grad = None
        output = forward()
        output.sum().backward()
        grad_analytical = param.grad.clone()

        # Numerical gradient
        grad_numerical = torch.zeros_like(param)
        param_data = param.data.clone()

        for i in range(min(param.numel(), 50)):  # Check first 50 elements
            param.data = param_data.clone()
            param.data.view(-1)[i] += eps
            out_plus = forward().sum().item()

            param.data = param_data.clone()
            param.data.view(-1)[i] -= eps
            out_minus = forward().sum().item()

            grad_numerical.view(-1)[i] = (out_plus - out_minus) / (2 * eps)

        param.data = param_data

        # Compare (only checked elements)
        n = min(param.numel(), 50)
        diff = (grad_analytical.view(-1)[:n] - grad_numerical.view(-1)[:n]).abs()
        rel_diff = (diff / (grad_analytical.view(-1)[:n].abs() + 1e-6)).mean().item() * 100
        return rel_diff < 10.0  # 10% tolerance for numerical precision

    ok_input = check_grad('input', input)
    ok_weight = check_grad('weight', weight)
    ok_offset = check_grad('offset', offset)

    ok = ok_input and ok_weight and ok_offset
    print(f"  Gradient correctness: input={ok_input}, weight={ok_weight}, offset={ok_offset}")
    return ok

def compare_fp32_bf16():
    """Compare FP32 vs BF16 outputs"""
    torch.manual_seed(42)

    B, C_in, H, W = 1, 32, 16, 16
    C_out = 32
    K = 3

    input_fp32 = torch.randn(B, C_in, H, W, device='mps', dtype=torch.float32)
    weight_fp32 = torch.randn(C_out, C_in, K, K, device='mps', dtype=torch.float32)

    out_h = H - K + 1
    out_w = W - K + 1
    offset_fp32 = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=torch.float32) * 0.1

    output_fp32 = deform_conv2d(input_fp32, offset_fp32, weight_fp32, stride=(1, 1), padding=(0, 0))

    # BF16
    input_bf16 = input_fp32.to(torch.bfloat16)
    weight_bf16 = weight_fp32.to(torch.bfloat16)
    offset_bf16 = offset_fp32.to(torch.bfloat16)

    output_bf16 = deform_conv2d(input_bf16, offset_bf16, weight_bf16, stride=(1, 1), padding=(0, 0))

    diff = (output_fp32 - output_bf16.to(torch.float32)).abs()
    max_diff = diff.max().item()
    rel_diff = (diff / (output_fp32.abs() + 1e-6)).mean().item() * 100

    ok = rel_diff < 5.0
    print(f"  FP32 vs BF16: max_diff={max_diff:.6f}, rel_diff={rel_diff:.2f}%, ok={ok}")
    return ok

def benchmark(dtype, name, warmup=5, runs=20):
    """Benchmark forward pass"""
    torch.manual_seed(42)

    B, C_in, H, W = 4, 128, 64, 64
    C_out = 128
    K = 3

    input = torch.randn(B, C_in, H, W, device='mps', dtype=dtype)
    weight = torch.randn(C_out, C_in, K, K, device='mps', dtype=dtype)

    out_h = H - K + 1
    out_w = W - K + 1
    offset = torch.randn(B, 2 * K * K, out_h, out_w, device='mps', dtype=dtype) * 0.1

    for _ in range(warmup):
        _ = deform_conv2d(input, offset, weight, stride=(1, 1), padding=(0, 0))
    torch.mps.synchronize()

    start = time.time()
    for _ in range(runs):
        _ = deform_conv2d(input, offset, weight, stride=(1, 1), padding=(0, 0))
    torch.mps.synchronize()
    elapsed = time.time() - start

    ms = (elapsed / runs) * 1000
    print(f"  {name}: {ms:.2f} ms")
    return ms

if __name__ == "__main__":
    print("\n" + "=" * 50)
    print("Testing mps-deform-conv")
    print("=" * 50)

    all_ok = True

    print("\n0. Validation tests:")
    all_ok &= test_stride_validation()
    all_ok &= test_dilation_validation()
    all_ok &= test_negative_padding_validation()
    all_ok &= test_groups_validation()
    all_ok &= test_device_mismatch()
    all_ok &= test_invalid_output_dimensions()
    all_ok &= test_bf16_precision_acceptable()

    print("\n1. Forward pass:")
    all_ok &= test_forward(torch.float32, "FP32")
    all_ok &= test_forward(torch.float16, "FP16")
    all_ok &= test_forward(torch.bfloat16, "BF16")

    print("\n2. Forward with mask (DCNv2):")
    all_ok &= test_forward_with_mask(torch.float32, "FP32")
    all_ok &= test_forward_with_mask(torch.float16, "FP16")
    all_ok &= test_forward_with_mask(torch.bfloat16, "BF16")

    print("\n3. Backward pass:")
    all_ok &= test_backward(torch.float32, "FP32")
    all_ok &= test_backward(torch.float16, "FP16")
    all_ok &= test_backward(torch.bfloat16, "BF16")

    print("\n4. DeformConv2d module:")
    all_ok &= test_module(torch.float32, "FP32")
    all_ok &= test_module(torch.float16, "FP16")
    all_ok &= test_module(torch.bfloat16, "BF16")

    print("\n5. ModulatedDeformConv2d module:")
    all_ok &= test_modulated_module(torch.float32, "FP32")
    all_ok &= test_modulated_module(torch.float16, "FP16")
    all_ok &= test_modulated_module(torch.bfloat16, "BF16")

    print("\n6. Stride and dilation:")
    all_ok &= test_stride_dilation(torch.float32, "FP32")
    all_ok &= test_stride_dilation(torch.float16, "FP16")
    all_ok &= test_stride_dilation(torch.bfloat16, "BF16")

    print("\n7. Gradient correctness:")
    all_ok &= test_gradient_correctness()

    print("\n8. FP32 vs BF16 comparison:")
    all_ok &= compare_fp32_bf16()

    print("\n9. Benchmarks:")
    benchmark(torch.float32, "FP32")
    benchmark(torch.float16, "FP16")
    benchmark(torch.bfloat16, "BF16")

    print("\n" + "=" * 50)
    if all_ok:
        print("ALL TESTS PASSED!")
    else:
        print("SOME TESTS FAILED!")
    print("=" * 50)
