"""Tests for Linen implementation."""

import pytest
import numpy as np


def test_linen_import():
    """Test that Linen module can be imported."""
    try:
        from nmn.linen import nmn
        assert True
    except ImportError as e:
        pytest.skip(f"Linen/JAX dependencies not available: {e}")


def test_linen_basic_functionality():
    """Test basic Linen NMN functionality.""" 
    try:
        import jax
        import jax.numpy as jnp
        from flax import linen as nn
        from nmn.linen.nmn import YatNMN
        
        # Create layer
        layer = YatNMN(features=10)
        
        # Initialize parameters
        key = jax.random.PRNGKey(0)
        dummy_input = jnp.ones((4, 8))
        params = layer.init(key, dummy_input)
        
        # Test forward pass
        output = layer.apply(params, dummy_input)
        
        assert output.shape == (4, 10)
        
    except ImportError:
        pytest.skip("JAX/Flax dependencies not available")