#!/usr/bin/env python3
"""
Test script to verify both Flash Attention v2 and v3 can perform inference
"""

import torch


def test_fa2_fa3():
    """Test both FA2 and FA3 inference capabilities"""

    print('Testing FA2 and FA3 availability...')
    print('=' * 60)

    # Test FA2
    fa2_available = False
    try:
        import flash_attn
        print('✓ FA2 (flash_attn) imported successfully')
        print(f'  Version: {flash_attn.__version__}')
        print(f'  Location: {flash_attn.__file__}')

        # Try to import the main function
        from flash_attn import flash_attn_func
        print('  ✓ flash_attn_func available')
        fa2_available = True
    except Exception as e:
        print(f'✗ FA2 import failed: {e}')

    print()

    # Test FA3
    fa3_available = False
    try:
        import flash_attn_interface
        print('✓ FA3 (flash_attn_interface) imported successfully')
        print(f'  Location: {flash_attn_interface.__file__}')

        # Check what functions are available
        funcs = [attr for attr in dir(flash_attn_interface) if not attr.startswith('_')]
        print(f'  Available functions: {funcs[:5]}...')  # Show first 5
        fa3_available = True
    except Exception as e:
        print(f'✗ FA3 import failed: {e}')

    print()
    print('=' * 60)
    print('Testing actual inference...')
    print('=' * 60)

    # Test parameters
    batch_size = 2
    seqlen = 128
    nheads = 8
    headdim = 64

    # Create test tensors
    q = torch.randn(batch_size, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)
    k = torch.randn(batch_size, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)
    v = torch.randn(batch_size, seqlen, nheads, headdim, device='cuda', dtype=torch.float16)

    # Test FA2
    if fa2_available:
        try:
            from flash_attn import flash_attn_func

            out_fa2 = flash_attn_func(q, k, v)
            print(f'✓ FA2 inference working: output shape {out_fa2.shape}')
        except Exception as e:
            print(f'✗ FA2 inference failed: {e}')
    else:
        print('⚠ FA2 not available, skipping inference test')

    # Test FA3
    if fa3_available:
        try:
            import flash_attn_interface

            # FA3 typically has the same interface
            if hasattr(flash_attn_interface, 'flash_attn_func'):
                out_fa3 = flash_attn_interface.flash_attn_func(q, k, v)
                print(f'✓ FA3 inference working: output shape {out_fa3.shape}')
            else:
                print('⚠ FA3 available but API differs from FA2')
                print(f'  Available in flash_attn_interface: {[x for x in dir(flash_attn_interface) if "flash" in x.lower()]}')
        except Exception as e:
            print(f'✗ FA3 inference test failed: {e}')
    else:
        print('⚠ FA3 not available, skipping inference test')

    print('=' * 60)


if __name__ == '__main__':
    test_fa2_fa3()