// Copyright (c) 2024 RustyTorch++ Team
// Licensed under the Apache License, Version 2.0

//! Comprehensive TDD tests for SparseMatrix operations.
//! Tests all new implementations: bandwidth(), scalar multiplication, and matrix addition.

use nalgebra::DVector;
use rtx_fea::assembly::sparse_matrix::SparseMatrix;

#[test]
fn test_bandwidth_empty_matrix() {
    // RED: Test empty matrix has bandwidth 0
    let matrix = SparseMatrix::new(5, 5);
    assert_eq!(
        matrix.bandwidth(),
        0,
        "Empty matrix should have bandwidth 0"
    );
}

#[test]
fn test_bandwidth_diagonal_matrix() {
    // GREEN: Test diagonal matrix has bandwidth 0
    let mut matrix = SparseMatrix::new(4, 4);

    // Add diagonal elements
    matrix.add_entry(0, 0, 1.0).unwrap();
    matrix.add_entry(1, 1, 2.0).unwrap();
    matrix.add_entry(2, 2, 3.0).unwrap();
    matrix.add_entry(3, 3, 4.0).unwrap();
    matrix.finalize().unwrap();

    assert_eq!(
        matrix.bandwidth(),
        0,
        "Diagonal matrix should have bandwidth 0"
    );
}

#[test]
fn test_bandwidth_tridiagonal_matrix() {
    // Test tridiagonal matrix has bandwidth 1
    let mut matrix = SparseMatrix::new(4, 4);

    // Main diagonal
    matrix.add_entry(0, 0, 2.0).unwrap();
    matrix.add_entry(1, 1, 2.0).unwrap();
    matrix.add_entry(2, 2, 2.0).unwrap();
    matrix.add_entry(3, 3, 2.0).unwrap();

    // Upper diagonal
    matrix.add_entry(0, 1, -1.0).unwrap();
    matrix.add_entry(1, 2, -1.0).unwrap();
    matrix.add_entry(2, 3, -1.0).unwrap();

    // Lower diagonal
    matrix.add_entry(1, 0, -1.0).unwrap();
    matrix.add_entry(2, 1, -1.0).unwrap();
    matrix.add_entry(3, 2, -1.0).unwrap();

    matrix.finalize().unwrap();

    assert_eq!(
        matrix.bandwidth(),
        1,
        "Tridiagonal matrix should have bandwidth 1"
    );
}

#[test]
fn test_bandwidth_full_matrix() {
    // Test full dense matrix has maximum bandwidth
    let mut matrix = SparseMatrix::new(3, 3);

    // Fill all entries
    for i in 0..3 {
        for j in 0..3 {
            matrix.add_entry(i, j, (i * 3 + j) as f64 + 1.0).unwrap();
        }
    }
    matrix.finalize().unwrap();

    assert_eq!(
        matrix.bandwidth(),
        2,
        "3x3 full matrix should have bandwidth 2"
    );
}

#[test]
fn test_bandwidth_sparse_random_pattern() {
    // Test sparse matrix with specific pattern
    let mut matrix = SparseMatrix::new(5, 5);

    // Create a specific sparse pattern with known bandwidth
    matrix.add_entry(0, 0, 1.0).unwrap();
    matrix.add_entry(0, 3, 2.0).unwrap(); // bandwidth = 3
    matrix.add_entry(1, 1, 3.0).unwrap();
    matrix.add_entry(2, 2, 4.0).unwrap();
    matrix.add_entry(3, 3, 5.0).unwrap();
    matrix.add_entry(4, 1, 6.0).unwrap(); // bandwidth = 3
    matrix.finalize().unwrap();

    assert_eq!(
        matrix.bandwidth(),
        3,
        "Sparse matrix bandwidth calculation incorrect"
    );
}

#[test]
fn test_bandwidth_single_element() {
    // Test matrix with single element
    let mut matrix = SparseMatrix::new(3, 3);
    matrix.add_entry(1, 1, 5.0).unwrap();
    matrix.finalize().unwrap();

    assert_eq!(
        matrix.bandwidth(),
        0,
        "Single element at diagonal should have bandwidth 0"
    );
}

#[test]
fn test_bandwidth_single_off_diagonal() {
    // Test matrix with single off-diagonal element
    let mut matrix = SparseMatrix::new(5, 5);
    matrix.add_entry(1, 4, 7.0).unwrap(); // bandwidth = 3
    matrix.finalize().unwrap();

    assert_eq!(
        matrix.bandwidth(),
        3,
        "Single off-diagonal element bandwidth incorrect"
    );
}

// Scalar multiplication tests

#[test]
fn test_scalar_multiply_by_zero() {
    // Multiply by 0 should give zero matrix
    let mut matrix = SparseMatrix::new(3, 3);
    matrix.add_entry(0, 0, 5.0).unwrap();
    matrix.add_entry(1, 2, 3.0).unwrap();
    matrix.add_entry(2, 1, -2.0).unwrap();
    matrix.finalize().unwrap();

    let result = &matrix * 0.0;

    // Check all values are zero
    let test_vec = DVector::from_vec(vec![1.0, 1.0, 1.0]);
    let product = result.multiply_vector(&test_vec).unwrap();

    assert_eq!(product[0], 0.0, "Multiply by zero should zero all elements");
    assert_eq!(product[1], 0.0, "Multiply by zero should zero all elements");
    assert_eq!(product[2], 0.0, "Multiply by zero should zero all elements");
}

#[test]
fn test_scalar_multiply_by_one() {
    // Multiply by 1 should be identity operation
    let mut matrix = SparseMatrix::new(3, 3);
    matrix.add_entry(0, 0, 5.0).unwrap();
    matrix.add_entry(1, 2, 3.0).unwrap();
    matrix.add_entry(2, 1, -2.0).unwrap();
    matrix.finalize().unwrap();

    let original_values = matrix.values().to_vec();
    let result = &matrix * 1.0;

    assert_eq!(
        result.values(),
        original_values.as_slice(),
        "Multiply by 1 should preserve values"
    );
    // Check structure is preserved by testing with multiplication
    let test_vec = DVector::from_vec(vec![1.0, 1.0, 1.0]);
    let original_product = matrix.multiply_vector(&test_vec).unwrap();
    let result_product = result.multiply_vector(&test_vec).unwrap();
    assert_eq!(
        result_product, original_product,
        "Multiply by 1 should preserve structure"
    );
}

#[test]
fn test_scalar_multiply_by_negative_one() {
    // Multiply by -1 should negate all values
    let mut matrix = SparseMatrix::new(3, 3);
    matrix.add_entry(0, 0, 5.0).unwrap();
    matrix.add_entry(1, 2, 3.0).unwrap();
    matrix.add_entry(2, 1, -2.0).unwrap();
    matrix.finalize().unwrap();

    let result = &matrix * -1.0;

    let result_values = result.values();
    assert_eq!(result_values[0], -5.0, "Values should be negated");
    assert_eq!(result_values[1], -3.0, "Values should be negated");
    assert_eq!(result_values[2], 2.0, "Values should be negated");
}

#[test]
fn test_scalar_multiply_by_fraction() {
    // Multiply by 0.5 should halve all values
    let mut matrix = SparseMatrix::new(2, 2);
    matrix.add_entry(0, 0, 10.0).unwrap();
    matrix.add_entry(1, 1, 20.0).unwrap();
    matrix.finalize().unwrap();

    let result = &matrix * 0.5;

    let result_values = result.values();
    assert_eq!(result_values[0], 5.0, "Values should be halved");
    assert_eq!(result_values[1], 10.0, "Values should be halved");
}

#[test]
fn test_scalar_multiply_preserves_sparsity() {
    // Multiplication should preserve sparsity pattern
    let mut matrix = SparseMatrix::new(4, 4);
    matrix.add_entry(0, 0, 1.0).unwrap();
    matrix.add_entry(1, 2, 2.0).unwrap();
    matrix.add_entry(3, 1, 3.0).unwrap();
    matrix.finalize().unwrap();

    let nnz_before = matrix.nnz();
    let result = &matrix * 2.5;
    let nnz_after = result.nnz();

    assert_eq!(
        nnz_after, nnz_before,
        "Scalar multiplication should preserve sparsity"
    );
}

#[test]
fn test_scalar_multiply_unfinalized_matrix() {
    // Should work with unfinalized matrix
    let mut matrix = SparseMatrix::new(2, 2);
    matrix.add_entry(0, 0, 4.0).unwrap();
    matrix.add_entry(1, 1, 6.0).unwrap();
    // Don't finalize

    let mut result = &matrix * 3.0;

    // Result should be unfinalized like the original
    // Finalize and check values
    result.finalize().unwrap();
    let result_values = result.values();
    assert_eq!(
        result_values[0], 12.0,
        "Unfinalized multiplication should work"
    );
    assert_eq!(
        result_values[1], 18.0,
        "Unfinalized multiplication should work"
    );
}

// Matrix addition tests

#[test]
fn test_matrix_add_zero_matrix() {
    // Adding zero matrix should be identity operation
    let mut matrix = SparseMatrix::new(3, 3);
    matrix.add_entry(0, 0, 5.0).unwrap();
    matrix.add_entry(1, 1, 3.0).unwrap();
    matrix.add_entry(2, 2, 7.0).unwrap();
    matrix.finalize().unwrap();

    let zero_matrix = SparseMatrix::new(3, 3);
    let result = matrix.clone() + &zero_matrix;

    // Result should equal original matrix
    assert_eq!(
        result.nnz(),
        matrix.nnz(),
        "Adding zero should preserve sparsity"
    );
    let result_values = result.values();
    let matrix_values = matrix.values();
    for i in 0..result_values.len() {
        assert_eq!(
            result_values[i], matrix_values[i],
            "Adding zero should preserve values"
        );
    }
}

#[test]
fn test_matrix_add_to_itself() {
    // Adding matrix to itself should double all values
    let mut matrix = SparseMatrix::new(2, 2);
    matrix.add_entry(0, 0, 2.0).unwrap();
    matrix.add_entry(0, 1, 3.0).unwrap();
    matrix.add_entry(1, 0, 4.0).unwrap();
    matrix.add_entry(1, 1, 5.0).unwrap();
    matrix.finalize().unwrap();

    let mut result = matrix.clone() + &matrix;
    result.finalize().unwrap();

    // Check doubled values
    let test_vec = DVector::from_vec(vec![1.0, 1.0]);
    let original_product = matrix.multiply_vector(&test_vec).unwrap();
    let result_product = result.multiply_vector(&test_vec).unwrap();

    assert_eq!(
        result_product[0],
        2.0 * original_product[0],
        "Values should be doubled"
    );
    assert_eq!(
        result_product[1],
        2.0 * original_product[1],
        "Values should be doubled"
    );
}

#[test]
fn test_matrix_add_different_patterns() {
    // Add matrices with different sparsity patterns
    let mut matrix1 = SparseMatrix::new(3, 3);
    matrix1.add_entry(0, 0, 1.0).unwrap();
    matrix1.add_entry(1, 1, 2.0).unwrap();
    matrix1.finalize().unwrap();

    let mut matrix2 = SparseMatrix::new(3, 3);
    matrix2.add_entry(1, 1, 3.0).unwrap(); // Overlapping
    matrix2.add_entry(2, 2, 4.0).unwrap(); // Non-overlapping
    matrix2.finalize().unwrap();

    let mut result = matrix1.clone() + &matrix2;
    result.finalize().unwrap();

    // Test with vector multiplication to verify correctness
    let test_vec = DVector::from_vec(vec![1.0, 1.0, 1.0]);
    let product = result.multiply_vector(&test_vec).unwrap();

    assert_eq!(
        product[0], 1.0,
        "Non-overlapping entries should be preserved"
    );
    assert_eq!(
        product[1], 5.0,
        "Overlapping entries should be summed (2+3)"
    );
    assert_eq!(
        product[2], 4.0,
        "Non-overlapping entries should be preserved"
    );
}

#[test]
fn test_matrix_add_dimension_mismatch() {
    // Adding matrices with different dimensions should panic or return error
    let matrix1 = SparseMatrix::new(3, 3);
    let matrix2 = SparseMatrix::new(2, 2);

    // This should panic due to dimension mismatch
    let result = std::panic::catch_unwind(|| matrix1 + &matrix2);

    assert!(
        result.is_err(),
        "Adding matrices with different dimensions should panic"
    );
}

#[test]
fn test_matrix_add_maintains_symmetry() {
    // Test that addition is commutative
    let mut matrix1 = SparseMatrix::new(3, 3);
    matrix1.add_entry(0, 1, 2.0).unwrap();
    matrix1.add_entry(1, 2, 3.0).unwrap();
    matrix1.finalize().unwrap();

    let mut matrix2 = SparseMatrix::new(3, 3);
    matrix2.add_entry(0, 0, 1.0).unwrap();
    matrix2.add_entry(1, 2, 1.0).unwrap();
    matrix2.finalize().unwrap();

    let mut result1 = matrix1.clone() + &matrix2;
    let mut result2 = matrix2.clone() + &matrix1;

    result1.finalize().unwrap();
    result2.finalize().unwrap();

    // Both results should be identical
    assert_eq!(
        result1.nnz(),
        result2.nnz(),
        "Addition should be commutative"
    );

    // Sort values for comparison (order might differ due to assembly)
    let mut vals1 = result1.values().to_vec();
    let mut vals2 = result2.values().to_vec();
    vals1.sort_by(|a, b| a.total_cmp(b));
    vals2.sort_by(|a, b| a.total_cmp(b));

    assert_eq!(vals1, vals2, "Addition should be commutative");
}

#[test]
fn test_matrix_add_large_sparse() {
    // Test addition with larger sparse matrices
    let size = 100;
    let mut matrix1 = SparseMatrix::new(size, size);
    let mut matrix2 = SparseMatrix::new(size, size);

    // Add some sparse entries
    for i in 0..size {
        matrix1.add_entry(i, i, i as f64).unwrap();
        if i > 0 {
            matrix2.add_entry(i, i - 1, 0.5).unwrap();
        }
    }

    matrix1.finalize().unwrap();
    matrix2.finalize().unwrap();

    let mut result = matrix1.clone() + &matrix2;
    result.finalize().unwrap();

    // Check that result has elements from both matrices
    assert!(
        result.nnz() >= matrix1.nnz(),
        "Result should contain all entries"
    );
    assert!(
        result.nnz() >= matrix2.nnz(),
        "Result should contain all entries"
    );
}

// Combined operations tests

#[test]
fn test_scalar_multiply_then_add() {
    // Test: 2*A + B
    let mut matrix_a = SparseMatrix::new(2, 2);
    matrix_a.add_entry(0, 0, 3.0).unwrap();
    matrix_a.add_entry(1, 1, 4.0).unwrap();
    matrix_a.finalize().unwrap();

    let mut matrix_b = SparseMatrix::new(2, 2);
    matrix_b.add_entry(0, 0, 1.0).unwrap();
    matrix_b.add_entry(1, 1, 2.0).unwrap();
    matrix_b.finalize().unwrap();

    let scaled_a = &matrix_a * 2.0;
    let mut result = scaled_a + &matrix_b;
    result.finalize().unwrap();

    // Verify: result[0,0] = 2*3 + 1 = 7, result[1,1] = 2*4 + 2 = 10
    let test_vec = DVector::from_vec(vec![1.0, 0.0]);
    let product1 = result.multiply_vector(&test_vec).unwrap();
    assert_eq!(product1[0], 7.0, "Combined operation incorrect");

    let test_vec2 = DVector::from_vec(vec![0.0, 1.0]);
    let product2 = result.multiply_vector(&test_vec2).unwrap();
    assert_eq!(product2[1], 10.0, "Combined operation incorrect");
}

#[test]
fn test_associativity_of_addition() {
    // Test: (A + B) + C == A + (B + C)
    let mut matrix_a = SparseMatrix::new(2, 2);
    matrix_a.add_entry(0, 0, 1.0).unwrap();
    matrix_a.finalize().unwrap();

    let mut matrix_b = SparseMatrix::new(2, 2);
    matrix_b.add_entry(1, 1, 2.0).unwrap();
    matrix_b.finalize().unwrap();

    let mut matrix_c = SparseMatrix::new(2, 2);
    matrix_c.add_entry(0, 1, 3.0).unwrap();
    matrix_c.finalize().unwrap();

    // (A + B) + C
    let ab = &matrix_a + &matrix_b;
    let mut abc1 = ab + &matrix_c;
    abc1.finalize().unwrap();

    // A + (B + C)
    let bc = &matrix_b + &matrix_c;
    let mut abc2 = matrix_a.clone() + &bc;
    abc2.finalize().unwrap();

    // Compare results
    assert_eq!(abc1.nnz(), abc2.nnz(), "Associativity should hold");

    // Test with vector multiplication
    let test_vec = DVector::from_vec(vec![1.0, 1.0]);
    let product1 = abc1.multiply_vector(&test_vec).unwrap();
    let product2 = abc2.multiply_vector(&test_vec).unwrap();

    for i in 0..2 {
        assert!(
            (product1[i] - product2[i]).abs() < 1e-10,
            "Associativity should hold"
        );
    }
}

#[test]
fn test_distributivity() {
    // Test: a*(B + C) == a*B + a*C
    let mut matrix_b = SparseMatrix::new(2, 2);
    matrix_b.add_entry(0, 0, 2.0).unwrap();
    matrix_b.add_entry(1, 1, 3.0).unwrap();
    matrix_b.finalize().unwrap();

    let mut matrix_c = SparseMatrix::new(2, 2);
    matrix_c.add_entry(0, 0, 4.0).unwrap();
    matrix_c.add_entry(1, 1, 5.0).unwrap();
    matrix_c.finalize().unwrap();

    let scalar = 2.5;

    // a*(B + C)
    let bc = &matrix_b + &matrix_c;
    let mut result1 = &bc * scalar;
    result1.finalize().unwrap();

    // a*B + a*C
    let ab = &matrix_b * scalar;
    let ac = &matrix_c * scalar;
    let mut result2 = ab + &ac;
    result2.finalize().unwrap();

    // Compare results
    let test_vec = DVector::from_vec(vec![1.0, 1.0]);
    let product1 = result1.multiply_vector(&test_vec).unwrap();
    let product2 = result2.multiply_vector(&test_vec).unwrap();

    for i in 0..2 {
        assert!(
            (product1[i] - product2[i]).abs() < 1e-10,
            "Distributivity should hold: {} != {}",
            product1[i],
            product2[i]
        );
    }
}

// Edge cases and stress tests

#[test]
fn test_operations_on_empty_matrix() {
    let empty = SparseMatrix::new(3, 3);

    // Scalar multiply empty
    let scaled = &empty * 5.0;
    assert_eq!(scaled.nnz(), 0, "Scaled empty matrix should be empty");

    // Add empty to empty
    let sum = empty.clone() + &empty;
    assert_eq!(sum.nnz(), 0, "Sum of empty matrices should be empty");
}

#[test]
fn test_bandwidth_rectangular_matrix() {
    // Test bandwidth for non-square matrix
    let mut matrix = SparseMatrix::new(3, 5);
    matrix.add_entry(0, 4, 1.0).unwrap();
    matrix.add_entry(2, 0, 2.0).unwrap();
    matrix.finalize().unwrap();

    let bandwidth = matrix.bandwidth();
    assert_eq!(bandwidth, 4, "Rectangular matrix bandwidth calculation");
}

#[test]
fn test_numerical_stability() {
    // Test with very small and very large numbers
    let mut matrix = SparseMatrix::new(2, 2);
    matrix.add_entry(0, 0, 1e-15).unwrap();
    matrix.add_entry(1, 1, 1e15).unwrap();
    matrix.finalize().unwrap();

    // Scalar multiplication
    let scaled = &matrix * 1e10;
    let scaled_values = scaled.values();
    assert!(
        scaled_values[0] > 0.0,
        "Small number scaling should maintain sign"
    );
    assert!(
        scaled_values[1] > 0.0,
        "Large number scaling should not overflow"
    );

    // Addition
    let result = matrix.clone() + &matrix;
    let result_values = result.values();
    assert_eq!(result_values[0], 2e-15, "Small number addition");
    assert_eq!(result_values[1], 2e15, "Large number addition");
}