import pytest
import tensorflow as tf
import numpy as np
import os
import sys

# Add the src directory to the Python path to allow imports
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))

try:
    from gru_model_handler import build_gru_model_v3
    MODEL_HANDLER_AVAILABLE = True
except ImportError as e:
    print(f"Failed to import gru_model_handler: {e}. Skipping tests that require it.")
    MODEL_HANDLER_AVAILABLE = False
    # Define a dummy function to avoid errors if import fails
    def build_gru_model_v3(*args, **kwargs):
        raise ImportError("gru_model_handler could not be imported.")

@pytest.mark.skipif(not MODEL_HANDLER_AVAILABLE, reason="GRUModelHandler module not found or TensorFlow/Keras missing")
def test_causal_mask_lambda_layer():
    """Tests if the Lambda layer correctly generates a lower-triangular causal mask."""
    lookback = 60
    n_features = 96
    attention_units = 16 # Must be > 0 to include the attention layer and mask

    # Build a minimal model instance to access the layer
    try:
        model = build_gru_model_v3(
            lookback=lookback, 
            n_features=n_features, 
            attention_units=attention_units
            # Use default values for other params for this test
        )
    except Exception as build_err:
        pytest.fail(f"Failed to build GRU model for testing: {build_err}")


    # Find the Lambda layer by name
    try:
        mask_lambda_layer = model.get_layer('causal_mask_lambda')
    except ValueError:
        pytest.fail("Could not find the 'causal_mask_lambda' layer in the model.")

    # Create a dummy input tensor (batch size 1, sequence length lookback, feature dim n_features)
    # The content doesn't matter, only the shape for the mask generation
    dummy_input_tensor = tf.zeros((1, lookback, n_features))

    # Get the mask output from the Lambda layer
    # Note: Calling the layer directly works outside tf.function context
    mask_output = mask_lambda_layer(dummy_input_tensor)

    # Expected mask is lower triangular matrix of shape (lookback, lookback)
    expected_mask = np.tril(np.ones((lookback, lookback)))

    # Compare the Lambda layer output with the expected mask
    # The mask_output might have a batch dimension added by Keras depending on context,
    # but the core logic in create_mask returns (seq_len, seq_len).
    # Let's check the shape and squeeze if necessary, although direct call might not add batch dim.
    
    # Convert tensor to numpy for comparison
    mask_numpy = mask_output.numpy()
    
    # Assert shape is as expected (seq_len, seq_len)
    assert mask_numpy.shape == (lookback, lookback), f"Expected mask shape {(lookback, lookback)}, but got {mask_numpy.shape}"

    # Assert the values are close to the lower triangular matrix
    assert np.allclose(mask_numpy, expected_mask), "Mask generated by Lambda layer is not lower-triangular."

    print("Causal mask test passed.")

# Add a main block to run the test if the file is executed directly (optional)
if __name__ == "__main__":
    # Basic check if dependencies are met
    if not MODEL_HANDLER_AVAILABLE:
         print("Skipping test execution: GRUModelHandler or TensorFlow/Keras not available.")
    else:
        print("Running causal mask test...")
        # Simple manual run without pytest fixtures/discovery
        try:
            test_causal_mask_lambda_layer()
        except Exception as e:
            print(f"Test failed with error: {e}")
        