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

import torch
import time

# 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()}")

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 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("\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. FP32 vs BF16 comparison:")
    all_ok &= compare_fp32_bf16()

    print("\n8. 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)
