#!/usr/bin/env python3
"""Test script for FlashAttentionForwardCombine cute-dsl implementation."""

import cutlass
try:
    from flash_fwd_combine import FlashAttentionForwardCombine
except ImportError:
    from flash_attn.cute.flash_fwd_combine import FlashAttentionForwardCombine


def test_can_implement():
    """Test the can_implement static method."""
    # Valid configuration
    assert FlashAttentionForwardCombine.can_implement(
        dtype=cutlass.Float16,
        dtype_partial=cutlass.Float32, 
        head_dim=64,
        m_block_size=128,
        k_block_size=64,
        log_max_splits=4,
        num_threads=256,
        is_even_k=True,
        varlen=False,
        stages=4
    )
    
    # Invalid head dimension
    assert not FlashAttentionForwardCombine.can_implement(
        dtype=cutlass.Float16,
        dtype_partial=cutlass.Float32,
        head_dim=65,  # Not multiple of 8
        m_block_size=128,
        k_block_size=64,
        log_max_splits=4,
        num_threads=256,
        is_even_k=True,
        varlen=False,
        stages=4
    )
    
    # Invalid dtype
    assert not FlashAttentionForwardCombine.can_implement(
        dtype=cutlass.Float32,  # Not supported
        dtype_partial=cutlass.Float32,
        head_dim=64,
        m_block_size=128,
        k_block_size=64,
        log_max_splits=4,
        num_threads=256,
        is_even_k=True,
        varlen=False,
        stages=4
    )
    
    print("✓ can_implement tests passed")


def test_initialization():
    """Test basic initialization of the class."""
    combine = FlashAttentionForwardCombine(
        dtype=cutlass.Float16,
        dtype_partial=cutlass.Float32,
        head_dim=64,
        m_block_size=128,
        k_block_size=64,
        log_max_splits=4,
        num_threads=256,
        is_even_k=True,
        stages=4
    )
    
    # Check attributes are set correctly
    assert combine.dtype == cutlass.Float16
    assert combine.dtype_partial == cutlass.Float32
    assert combine.head_dim == 64
    assert combine.m_block_size == 128
    assert combine.k_block_size == 64
    assert combine.max_splits == 16  # 2^4
    assert combine.num_threads == 256
    assert combine.is_even_k == True
    assert combine.stages == 4
    
    print("✓ initialization tests passed")


def test_setup_attributes():
    """Test the _setup_attributes method."""
    combine = FlashAttentionForwardCombine(
        dtype=cutlass.Float16,
        dtype_partial=cutlass.Float32,
        head_dim=64,
        m_block_size=128,
        k_block_size=64,
        log_max_splits=4,
        num_threads=256,
        is_even_k=True,
        stages=4
    )
    
    # This should not raise an exception
    combine._setup_attributes()
    
    # Check that tiled copies were created
    assert hasattr(combine, 'gmem_tiled_copy_O_partial')
    assert hasattr(combine, 'gmem_tiled_copy_O')
    assert hasattr(combine, 'gmem_tiled_copy_LSE')
    assert hasattr(combine, 's2r_tiled_copy_LSE')
    
    print("✓ _setup_attributes tests passed")


def main():
    """Run all tests."""
    print("Testing FlashAttentionForwardCombine...")
    
    test_can_implement()
    test_initialization() 
    test_setup_attributes()
    
    print("\n✓ All tests passed!")
    print("FlashAttentionForwardCombine implementation successfully ported from C++ to Cute-DSL")


if __name__ == "__main__":
    main()