//! Additional coverage tests for layers/mod.rs - Part 10
//!
//! Targets remaining uncovered code paths in:
//! - QuantizedLinear: 1D output shape construction (lines 710-712)
//! - Linear: weight_mut, bias_mut accessors
//! - LayerNorm: normalized_shape(), eps() getters
//! - FusedLayerNormLinear: feature_dim(), out_features() getters
//! - FeedForward: fc1_mut, fc2_mut with actual weight loading
//! - Attention submodules: additional edge cases
//! - Position embeddings: additional RoPE/ALiBi edge cases

use crate::error::RealizarError;
use crate::layers::*;
use crate::tensor::Tensor;

// =============================================================================
// QuantizedLinear: 1D Input/Output Shape Handling (lines 710-712)
// =============================================================================

#[test]
fn test_quantized_linear_forward_produces_correct_output_size() {
    // Create valid Q4_K layer: 256 in_features, 2 out_features
    // 256 values = 1 super-block = 144 bytes per row
    // 2 output rows = 2 * 144 = 288 bytes
    let weight_bytes = vec![0u8; 288];
    let bias = vec![0.5f32, -0.5f32];

    let layer = QuantizedLinear::new(256, 2, weight_bytes, bias).expect("create layer");

    // 2D input: [batch=3, in_features=256]
    let input = Tensor::from_vec(vec![3, 256], vec![0.1f32; 768]).expect("input");
    let output = layer.forward(&input).expect("forward");

    // Output shape should be [3, 2]
    assert_eq!(output.shape(), &[3, 2]);
    assert_eq!(output.data().len(), 6);
}

#[test]
fn test_quantized_linear_multi_batch_forward() {
    // Q4_K: 256 in_features, 4 out_features
    // 256 values = 1 super-block = 144 bytes per row
    // 4 output rows = 4 * 144 = 576 bytes
    let weight_bytes = vec![0u8; 576];
    let bias = vec![0.1f32, 0.2f32, 0.3f32, 0.4f32];

    let layer = QuantizedLinear::new(256, 4, weight_bytes, bias).expect("create layer");

    // 3D input: [2, 3, 256] - batch of 2, seq of 3
    let input = Tensor::from_vec(vec![2, 3, 256], vec![0.05f32; 1536]).expect("input");
    let output = layer.forward(&input).expect("forward");

    // Output shape should be [2, 3, 4]
    assert_eq!(output.shape(), &[2, 3, 4]);
    assert_eq!(output.data().len(), 24);
}

// =============================================================================
// Linear: Weight and Bias Accessor Tests
// =============================================================================

#[test]
fn test_linear_weight_mut_accessor() {
    let mut linear = Linear::new(4, 8).expect("create layer");

    // Access and modify weights
    let weights = linear.weight_mut();
    assert_eq!(weights.len(), 4 * 8, "Weights should have in*out elements");

    // Set specific pattern
    for (i, w) in weights.iter_mut().enumerate() {
        *w = (i as f32) * 0.1;
    }

    // Verify modification persists
    assert!((linear.weight_mut()[0] - 0.0).abs() < 1e-6);
    assert!((linear.weight_mut()[1] - 0.1).abs() < 1e-6);
    assert!((linear.weight_mut()[31] - 3.1).abs() < 1e-6);
}

#[test]
fn test_linear_bias_mut_accessor() {
    let mut linear = Linear::new(4, 8).expect("create layer");

    // Access and modify bias
    let bias = linear.bias_mut();
    assert_eq!(bias.len(), 8, "Bias should have out_features elements");

    // Set specific values
    for (i, b) in bias.iter_mut().enumerate() {
        *b = (i as f32) * -0.5;
    }

    // Verify modification persists
    assert!((linear.bias_mut()[0] - 0.0).abs() < 1e-6);
    assert!((linear.bias_mut()[1] - (-0.5)).abs() < 1e-6);
    assert!((linear.bias_mut()[7] - (-3.5)).abs() < 1e-6);
}

#[test]
fn test_linear_forward_with_custom_weights() {
    let mut linear = Linear::new(2, 3).expect("create layer");

    // Set identity-like weights: weight[i][j] = 1 if i==j else 0
    // For 2x3: column-major (out_features = 3)
    // weight[0,0]=1, weight[1,1]=1, rest=0
    for w in linear.weight_mut().iter_mut() {
        *w = 0.0;
    }
    linear.weight_mut()[0] = 1.0; // in=0, out=0
    linear.weight_mut()[4] = 1.0; // in=1, out=1

    // Zero bias
    for b in linear.bias_mut().iter_mut() {
        *b = 0.0;
    }

    let input = Tensor::from_vec(vec![2], vec![5.0, 3.0]).expect("input");
    let output = linear.forward(&input).expect("forward");

    // output[0] = input[0]*1 + input[1]*0 = 5.0
    // output[1] = input[0]*0 + input[1]*1 = 3.0
    // output[2] = input[0]*0 + input[1]*0 = 0.0
    assert_eq!(output.shape(), &[3]);
    assert!((output.data()[0] - 5.0).abs() < 1e-5);
    assert!((output.data()[1] - 3.0).abs() < 1e-5);
    assert!((output.data()[2] - 0.0).abs() < 1e-5);
}

#[test]
fn test_linear_forward_with_bias() {
    let mut linear = Linear::new(2, 2).expect("create layer");

    // Set weights to identity
    for w in linear.weight_mut().iter_mut() {
        *w = 0.0;
    }
    linear.weight_mut()[0] = 1.0; // in=0, out=0
    linear.weight_mut()[3] = 1.0; // in=1, out=1

    // Set non-zero bias
    linear.bias_mut()[0] = 10.0;
    linear.bias_mut()[1] = 20.0;

    let input = Tensor::from_vec(vec![2], vec![1.0, 2.0]).expect("input");
    let output = linear.forward(&input).expect("forward");

    // output = input * W + bias
    // output[0] = 1.0*1 + 2.0*0 + 10.0 = 11.0
    // output[1] = 1.0*0 + 2.0*1 + 20.0 = 22.0
    assert!((output.data()[0] - 11.0).abs() < 1e-5);
    assert!((output.data()[1] - 22.0).abs() < 1e-5);
}

// =============================================================================
// LayerNorm: Getter Methods Coverage
// =============================================================================

#[test]
fn test_layer_norm_normalized_shape_getter() {
    for size in [1, 64, 128, 512, 4096] {
        let layer_norm = LayerNorm::new(size, 1e-5).expect("create");
        assert_eq!(
            layer_norm.normalized_shape(),
            size,
            "normalized_shape() should return construction parameter"
        );
    }
}

#[test]
fn test_layer_norm_eps_getter() {
    let test_cases = [1e-12, 1e-6, 1e-5, 1e-3, 0.1, 1.0];

    for eps in test_cases {
        let layer_norm = LayerNorm::new(64, eps).expect("create");
        assert!(
            (layer_norm.eps() - eps).abs() < 1e-12,
            "eps() should return construction parameter, expected {} got {}",
            eps,
            layer_norm.eps()
        );
    }
}

#[test]
fn test_layer_norm_zero_variance_handling() {
    let layer_norm = LayerNorm::new(4, 1e-5).expect("create");

    // All same values -> variance = 0, but eps prevents division by zero
    let input = Tensor::from_vec(vec![4], vec![5.0, 5.0, 5.0, 5.0]).expect("input");
    let output = layer_norm.forward(&input).expect("forward");

    // With zero variance and eps=1e-5:
    // normalized = (x - mean) / sqrt(0 + eps) = 0 / sqrt(eps) = 0
    // With gamma=1 and beta=0: output = 0
    for &val in output.data() {
        assert!(val.is_finite(), "Should handle zero variance");
        assert!(
            val.abs() < 1e-2,
            "Zero variance should produce near-zero output"
        );
    }
}

#[test]
fn test_layer_norm_negative_values() {
    let layer_norm = LayerNorm::new(4, 1e-5).expect("create");

    let input = Tensor::from_vec(vec![4], vec![-10.0, -5.0, 5.0, 10.0]).expect("input");
    let output = layer_norm.forward(&input).expect("forward");

    assert_eq!(output.shape(), &[4]);

    // Mean should be 0, so outputs should be symmetric around 0
    let mean: f32 = output.data().iter().sum::<f32>() / 4.0;
    assert!(mean.abs() < 1e-5, "LayerNorm output mean should be ~0");
}

// =============================================================================
// FusedLayerNormLinear: Getter Methods Coverage
// =============================================================================

#[test]
fn test_fused_layer_norm_linear_feature_dim_getter() {
    for dim in [4, 64, 256, 512] {
        let fused = FusedLayerNormLinear::new(dim, 128, 1e-5).expect("create");
        assert_eq!(
            fused.feature_dim(),
            dim,
            "feature_dim() should return construction parameter"
        );
    }
}

#[test]
fn test_fused_layer_norm_linear_out_features_getter() {
    for out in [4, 64, 256, 512] {
        let fused = FusedLayerNormLinear::new(128, out, 1e-5).expect("create");
        assert_eq!(
            fused.out_features(),
            out,
            "out_features() should return construction parameter"
        );
    }
}

#[test]
fn test_fused_layer_norm_linear_combined_operation() {
    let mut fused = FusedLayerNormLinear::new(4, 2, 1e-5).expect("create");

    // Set norm weights to 2.0 (scale by 2)
    for w in fused.norm_weight_mut().iter_mut() {
        *w = 2.0;
    }

    // Set linear weights (identity-like for first 2 dims)
    for w in fused.linear_weight_mut().iter_mut() {
        *w = 0.0;
    }
    fused.linear_weight_mut()[0] = 0.5; // in=0, out=0
    fused.linear_weight_mut()[3] = 0.5; // in=1, out=1

    let input = Tensor::from_vec(vec![4], vec![1.0, 2.0, 3.0, 4.0]).expect("input");
    let output = fused.forward(&input).expect("forward");

    assert_eq!(output.shape(), &[2]);
    // Output should be finite
    for &val in output.data() {
        assert!(val.is_finite());
    }
}

#[test]
fn test_fused_layer_norm_linear_parallel_matches_serial_different_sizes() {
    for (feature_dim, out_features, batch_size) in [(8, 4, 10), (16, 8, 50), (32, 16, 100)] {
        let fused = FusedLayerNormLinear::new(feature_dim, out_features, 1e-5).expect("create");

        let input = Tensor::from_vec(
            vec![batch_size, feature_dim],
            vec![0.1f32; batch_size * feature_dim],
        )
        .expect("input");

        let serial = fused.forward(&input).expect("serial");
        let parallel = fused.forward_parallel(&input).expect("parallel");

        assert_eq!(serial.shape(), parallel.shape());

        for i in 0..serial.data().len() {
            assert!(
                (serial.data()[i] - parallel.data()[i]).abs() < 1e-4,
                "Mismatch at {} for dims ({}, {}, {}): {} vs {}",
                i,
                feature_dim,
                out_features,
                batch_size,
                serial.data()[i],
                parallel.data()[i]
            );
        }
    }
}

// =============================================================================
// FeedForward: fc1_mut and fc2_mut Accessor Tests
// =============================================================================

#[test]
fn test_ffn_fc1_mut_accessor() {
    let mut ffn = FeedForward::new(4, 16).expect("create");

    let fc1 = ffn.fc1_mut();
    assert_eq!(fc1.in_features(), 4);
    assert_eq!(fc1.out_features(), 16);

    // Modify weights
    fc1.weight_mut()[0] = 42.0;
    assert!((ffn.fc1_mut().weight_mut()[0] - 42.0).abs() < 1e-6);
}

#[test]
fn test_ffn_fc2_mut_accessor() {
    let mut ffn = FeedForward::new(4, 16).expect("create");

    let fc2 = ffn.fc2_mut();
    assert_eq!(fc2.in_features(), 16);
    assert_eq!(fc2.out_features(), 4);

    // Modify weights
    fc2.weight_mut()[0] = 99.0;
    assert!((ffn.fc2_mut().weight_mut()[0] - 99.0).abs() < 1e-6);
}

#[test]
fn test_ffn_with_loaded_weights_integration() {
    let mut ffn = FeedForward::new(2, 4).expect("create");

    // Set fc1: 2->4 expansion
    // Simple pattern: each output is sum of inputs
    for w in ffn.fc1_mut().weight_mut().iter_mut() {
        *w = 1.0;
    }
    for b in ffn.fc1_mut().bias_mut().iter_mut() {
        *b = 0.0;
    }

    // Set fc2: 4->2 projection
    // Average of intermediate outputs
    for w in ffn.fc2_mut().weight_mut().iter_mut() {
        *w = 0.25;
    }
    for b in ffn.fc2_mut().bias_mut().iter_mut() {
        *b = 0.0;
    }

    let input = Tensor::from_vec(vec![2], vec![1.0, 2.0]).expect("input");
    let output = ffn.forward(&input).expect("forward");

    assert_eq!(output.shape(), &[2]);
    // Output should be finite and reasonable
    for &val in output.data() {
        assert!(val.is_finite());
    }
}

// =============================================================================
// softmax: Additional Edge Cases
// =============================================================================

#[test]
fn test_softmax_2d_multiple_rows() {
    // Test 2D tensor with multiple rows, verifying each row sums to 1
    let input = Tensor::from_vec(
        vec![3, 4],
        vec![
            1.0, 2.0, 3.0, 4.0, // row 0
            0.0, 0.0, 0.0, 0.0, // row 1 (uniform)
            -1.0, -2.0, -3.0, -4.0, // row 2 (negative)
        ],
    )
    .expect("input");

    let output = softmax(&input).expect("softmax");

    assert_eq!(output.shape(), &[3, 4]);

    // Each row should sum to 1.0
    for row in 0..3 {
        let row_sum: f32 = (0..4).map(|col| output.data()[row * 4 + col]).sum();
        assert!(
            (row_sum - 1.0).abs() < 1e-5,
            "Row {} should sum to 1.0, got {}",
            row,
            row_sum
        );
    }

    // Row 1 (all zeros) should be uniform
    for col in 0..4 {
        assert!(
            (output.data()[4 + col] - 0.25).abs() < 1e-5,
            "Uniform row element should be 0.25"
        );
    }
}

#[test]
fn test_softmax_numerical_stability_extreme_range() {
    // Mix of very small and very large values
    let input = Tensor::from_vec(vec![4], vec![-500.0, 0.0, 500.0, 1000.0]).expect("input");

    let output = softmax(&input).expect("softmax");

    // Should produce finite values
    for &val in output.data() {
        assert!(val.is_finite(), "Should handle extreme range");
    }

    // Sum should be 1.0
    let sum: f32 = output.data().iter().sum();
    assert!((sum - 1.0).abs() < 1e-5);

    // Largest input should dominate
    assert!(output.data()[3] > 0.9, "Largest input should dominate");
}

// =============================================================================
// gelu: Additional Edge Cases
// =============================================================================

#[test]
fn test_gelu_preserves_shape_multi_dimensional() {
    // Test that GELU preserves shape for various dimensions
    let shapes = [vec![4], vec![2, 4], vec![2, 3, 4], vec![2, 3, 4, 5]];

    for shape in shapes {
        let size: usize = shape.iter().product();
        let input = Tensor::from_vec(shape.clone(), vec![0.5f32; size]).expect("input");
        let output = gelu(&input).expect("gelu");

        assert_eq!(output.shape(), &shape[..], "GELU should preserve shape");
    }
}

#[test]
fn test_gelu_monotonicity_positive_range() {
    // GELU should be monotonically increasing for x > 0
    let values: Vec<f32> = (0..10).map(|i| i as f32 * 0.5).collect();
    let input = Tensor::from_vec(vec![10], values).expect("input");
    let output = gelu(&input).expect("gelu");

    for i in 1..10 {
        assert!(
            output.data()[i] >= output.data()[i - 1],
            "GELU should be monotonic for positive x"
        );
    }
}

#[test]
fn test_gelu_approximation_accuracy() {
    // Test that GELU approximation is reasonably accurate
    let test_points = [0.0, 0.5, 1.0, 2.0, -0.5, -1.0];

    for &x in &test_points {
        let input = Tensor::from_vec(vec![1], vec![x]).expect("input");
        let output = gelu(&input).expect("gelu");

        // GELU(0) = 0
        if x == 0.0 {
            assert!(output.data()[0].abs() < 1e-6);
        }
        // GELU(x) > 0 for x > 0
        if x > 0.0 {
            assert!(output.data()[0] > 0.0);
        }
        // GELU should be finite
        assert!(output.data()[0].is_finite());
    }
}

// =============================================================================
// Error Path Coverage: Construction Errors
// =============================================================================

#[test]
fn test_linear_zero_in_features_error() {
    let result = Linear::new(0, 8);
    assert!(result.is_err());

    if let Err(RealizarError::InvalidShape { reason }) = result {
        assert!(reason.contains("> 0") || reason.contains("in_features"));
    }
}

#[test]
fn test_linear_zero_out_features_error() {
    let result = Linear::new(8, 0);
    assert!(result.is_err());

    if let Err(RealizarError::InvalidShape { reason }) = result {
        assert!(reason.contains("> 0") || reason.contains("out_features"));
    }
}

#[test]
fn test_layer_norm_zero_normalized_shape_error() {
    let result = LayerNorm::new(0, 1e-5);
    assert!(result.is_err());

    if let Err(RealizarError::InvalidShape { reason }) = result {
        assert!(reason.contains("> 0") || reason.contains("normalized_shape"));
    }
}

#[test]
fn test_fused_layer_norm_linear_zero_feature_dim_error() {
    let result = FusedLayerNormLinear::new(0, 8, 1e-5);
    assert!(result.is_err());

    if let Err(RealizarError::InvalidShape { reason }) = result {
        assert!(reason.contains("> 0") || reason.contains("feature_dim"));
    }
}

#[test]
fn test_fused_layer_norm_linear_zero_out_features_error() {
    let result = FusedLayerNormLinear::new(8, 0, 1e-5);
    assert!(result.is_err());

    if let Err(RealizarError::InvalidShape { reason }) = result {
        assert!(reason.contains("> 0") || reason.contains("out_features"));
    }
}

#[test]
fn test_ffn_zero_hidden_dim_error() {
    let result = FeedForward::new(0, 16);
    assert!(result.is_err());
}

#[test]
fn test_ffn_zero_intermediate_dim_error() {
    let result = FeedForward::new(4, 0);
    assert!(result.is_err());
}

// =============================================================================
// Batched Input Edge Cases
// =============================================================================

#[test]
fn test_linear_single_batch_dimension() {
    let linear = Linear::new(4, 8).expect("create");

    // Single item batch: [1, 4]
    let input = Tensor::from_vec(vec![1, 4], vec![1.0, 2.0, 3.0, 4.0]).expect("input");
    let output = linear.forward(&input).expect("forward");

    assert_eq!(output.shape(), &[1, 8]);
}

#[test]
fn test_layer_norm_single_batch() {
    let layer_norm = LayerNorm::new(4, 1e-5).expect("create");

    // Single item batch: [1, 4]
    let input = Tensor::from_vec(vec![1, 4], vec![1.0, 2.0, 3.0, 4.0]).expect("input");
    let output = layer_norm.forward(&input).expect("forward");

    assert_eq!(output.shape(), &[1, 4]);
}

#[test]
fn test_fused_layer_norm_linear_single_batch() {
    let fused = FusedLayerNormLinear::new(4, 8, 1e-5).expect("create");

    // Single item batch
    let input = Tensor::from_vec(vec![1, 4], vec![1.0, 2.0, 3.0, 4.0]).expect("input");

    let serial = fused.forward(&input).expect("serial");
    let parallel = fused.forward_parallel(&input).expect("parallel");

    assert_eq!(serial.shape(), &[1, 8]);
    assert_eq!(parallel.shape(), &[1, 8]);
}

// =============================================================================
// Property-Based Style Tests
// =============================================================================

#[test]
fn test_softmax_output_bounds() {
    // Softmax output should be in [0, 1] for all inputs
    let test_values: Vec<f32> = (-10..10).map(|i| i as f32 * 0.5).collect();
    let input = Tensor::from_vec(vec![20], test_values).expect("input");
    let output = softmax(&input).expect("softmax");

    for &val in output.data() {
        assert!(val >= 0.0, "Softmax output should be >= 0");
        assert!(val <= 1.0, "Softmax output should be <= 1");
    }
}

#[test]
fn test_layer_norm_output_mean_zero() {
    // LayerNorm with default gamma=1, beta=0 should produce mean ~= 0
    let layer_norm = LayerNorm::new(8, 1e-5).expect("create");

    // Various input distributions
    let test_inputs = [
        vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
        vec![-4.0, -3.0, -2.0, -1.0, 1.0, 2.0, 3.0, 4.0],
        vec![100.0, 100.1, 100.2, 100.3, 100.4, 100.5, 100.6, 100.7],
    ];

    for values in test_inputs {
        let input = Tensor::from_vec(vec![8], values).expect("input");
        let output = layer_norm.forward(&input).expect("forward");

        let mean: f32 = output.data().iter().sum::<f32>() / 8.0;
        assert!(
            mean.abs() < 1e-4,
            "LayerNorm output mean should be ~0, got {}",
            mean
        );
    }
}

#[test]
fn test_linear_zero_input_produces_bias() {
    let mut linear = Linear::new(4, 3).expect("create");

    // Set all weights to 0, set specific bias
    for w in linear.weight_mut().iter_mut() {
        *w = 0.0;
    }
    linear.bias_mut()[0] = 1.0;
    linear.bias_mut()[1] = 2.0;
    linear.bias_mut()[2] = 3.0;

    // Zero input should produce bias
    let input = Tensor::from_vec(vec![4], vec![0.0, 0.0, 0.0, 0.0]).expect("input");
    let output = linear.forward(&input).expect("forward");

    assert!((output.data()[0] - 1.0).abs() < 1e-6);
    assert!((output.data()[1] - 2.0).abs() < 1e-6);
    assert!((output.data()[2] - 3.0).abs() < 1e-6);
}

// =============================================================================
// RoPE and Position Embedding Coverage
// =============================================================================

#[test]
fn test_rope_creation_various_dimensions() {
    for dim in [4, 8, 16, 32, 64, 128] {
        let rope = RoPE::new(dim, 512).expect("create");
        assert_eq!(rope.head_dim(), dim);
        assert_eq!(rope.max_seq_len(), 512);
    }
}

#[test]
fn test_scaled_rope_scaling_types() {
    // Test all scaling types can be created
    let linear = ScaledRoPE::new(64, 2048, RopeScalingType::Linear { scale: 2.0 }).expect("linear");
    assert_eq!(linear.head_dim(), 64);

    let dynamic =
        ScaledRoPE::new(64, 2048, RopeScalingType::Dynamic { scale: 2.0 }).expect("dynamic");
    assert_eq!(dynamic.head_dim(), 64);

    let ntk = ScaledRoPE::new(
        64,
        2048,
        RopeScalingType::NTKAware {
            scale: 2.0,
            alpha: 1.0,
        },
    )
    .expect("ntk");
    assert_eq!(ntk.head_dim(), 64);

    let yarn = ScaledRoPE::new(
        64,
        2048,
        RopeScalingType::YaRN {
            scale: 2.0,
            original_max_seq_len: 4096,
            attention_factor: 1.0,
            beta_fast: 32.0,
            beta_slow: 1.0,
        },
    )
    .expect("yarn");
    assert_eq!(yarn.head_dim(), 64);
}

#[test]
fn test_alibi_creation_various_heads() {
    for num_heads in [1, 2, 4, 8, 16, 32] {
        let alibi = ALiBi::new(num_heads).expect("create");
        assert_eq!(alibi.num_heads(), num_heads);
    }
}

#[test]
fn test_rope_forward_shape_preservation() {
    let rope = RoPE::new(8, 512).expect("create");

    // Test various input shapes
    let shapes = [vec![4, 8], vec![2, 4, 8], vec![1, 1, 8]];

    for shape in shapes {
        let size: usize = shape.iter().product();
        let input = Tensor::from_vec(shape.clone(), vec![0.1f32; size]).expect("input");
        let output = rope.forward(&input, 0).expect("forward");

        assert_eq!(output.shape(), &shape[..], "RoPE should preserve shape");
    }
}

#[test]
fn test_alibi_get_bias_various_seq_lengths() {
    let alibi = ALiBi::new(4).expect("create");

    for seq_len in [1, 4, 16, 64, 128] {
        let bias = alibi.get_bias(seq_len).expect("get_bias");

        // Bias shape should be [num_heads, seq_len, seq_len]
        assert_eq!(bias.shape(), &[4, seq_len, seq_len]);
    }
}

// =============================================================================
// Attention Submodule Additional Coverage
// =============================================================================

#[test]
fn test_attention_zero_head_dim_error() {
    let result = Attention::new(0);
    assert!(result.is_err(), "Should error on zero head_dim");
}

#[test]
fn test_attention_large_head_dim() {
    let attn = Attention::new(256).expect("create");

    let q = Tensor::from_vec(vec![2, 256], vec![0.1f32; 512]).expect("q");
    let k = Tensor::from_vec(vec![2, 256], vec![0.1f32; 512]).expect("k");
    let v = Tensor::from_vec(vec![2, 256], vec![0.1f32; 512]).expect("v");

    let output = attn.forward(&q, &k, &v).expect("forward");
    assert_eq!(output.shape(), &[2, 256]);
}

#[test]
fn test_sliding_window_attention_window_larger_than_seq() {
    let swa = SlidingWindowAttention::new(8, 100).expect("create");

    // seq_len=4 < window_size=100
    let q = Tensor::from_vec(vec![4, 8], vec![0.1f32; 32]).expect("q");
    let k = Tensor::from_vec(vec![4, 8], vec![0.1f32; 32]).expect("k");
    let v = Tensor::from_vec(vec![4, 8], vec![0.1f32; 32]).expect("v");

    // Should work - effectively full attention
    let output = swa.forward(&q, &k, &v).expect("forward");
    assert_eq!(output.shape(), &[4, 8]);
}

#[test]
fn test_multi_head_attention_weight_accessors() {
    let mut mha = MultiHeadAttention::mha(64, 8).expect("create");

    // Test all weight accessors are accessible and writable
    let w_q = mha.w_q_mut();
    assert!(!w_q.is_empty());
    w_q[0] = 1.0;

    let w_k = mha.w_k_mut();
    assert!(!w_k.is_empty());
    w_k[0] = 2.0;

    let w_v = mha.w_v_mut();
    assert!(!w_v.is_empty());
    w_v[0] = 3.0;

    let w_o = mha.w_o_mut();
    assert!(!w_o.is_empty());
    w_o[0] = 4.0;
}
