import numpy as np
import pytest

from pyvelora import Matrix, Vector
from pyvelora.utils import ShapeError


def test_matrix_init_valid_2d_data():
	m = Matrix([[1, 2], [3, 4]])
	assert isinstance(m, Matrix)
	assert m.shape == (2, 2)
	assert m.ndim == 2
	assert m.size == 4


def test_matrix_init_raises_for_non_2d_data():
	with pytest.raises(ValueError, match="Matrix must be 2D"):
		Matrix([1, 2, 3])


def test_matrix_getitem_scalar_and_slice():
	m = Matrix([[1, 2], [3, 4]])
	assert m[0, 1] == 2

	column = m[:, :1]
	assert isinstance(column, Matrix)
	assert np.allclose(column.data, np.array([[1], [3]]))


def test_matrix_setitem_updates_value():
	m = Matrix([[1, 2], [3, 4]])
	m[1, 0] = 9
	assert m[1, 0] == 9


def test_matrix_len_iter_and_contains():
	m = Matrix([[1, 2], [3, 4]])
	assert len(m) == 2

	rows = list(m)
	assert len(rows) == 2
	assert np.allclose(rows[0], np.array([1.0, 2.0]))
	assert 2 in m
	assert 99 not in m
	assert bool(m) is True


def test_matrix_add_and_subtract():
	m1 = Matrix([[1, 2], [3, 4]])
	m2 = Matrix([[5, 6], [7, 8]])

	added = m1 + m2
	subtracted = m2 - m1

	assert isinstance(added, Matrix)
	assert isinstance(subtracted, Matrix)
	assert np.allclose(added.data, np.array([[6, 8], [10, 12]]))
	assert np.allclose(subtracted.data, np.array([[4, 4], [4, 4]]))


def test_matrix_add_raises_for_shape_mismatch():
	m1 = Matrix([[1, 2], [3, 4]])
	m2 = Matrix([[1, 2, 3], [4, 5, 6]])

	with pytest.raises(ShapeError, match="same shape"):
		_ = m1 + m2


def test_matrix_mul_and_div_elementwise():
	m1 = Matrix([[2, 4], [6, 8]])
	m2 = Matrix([[1, 2], [3, 4]])
	multiplied = m1 * m2
	divided = m1 / m2

	assert np.allclose(multiplied.data, np.array([[2, 8], [18, 32]]))
	assert np.allclose(divided.data, np.array([[2.0, 2.0], [2.0, 2.0]]))


def test_matrix_divide_by_zero_raises():
	m = Matrix([[1, 2], [3, 4]])
	zero_matrix = Matrix([[1, 0], [1, 1]])
	with pytest.raises(ZeroDivisionError, match="Cannot divide by zero"):
		_ = m / zero_matrix





def test_matrix_matmul_matrix_returns_matrix():
	m1 = Matrix([[1, 2], [3, 4]])
	m2 = Matrix([[5, 6], [7, 8]])
	result = m1 @ m2

	assert isinstance(result, Matrix)
	assert np.allclose(result.data, np.array([[19, 22], [43, 50]]))


def test_matrix_matmul_raises_for_misaligned_shapes():
	m1 = Matrix([[1, 2], [3, 4]])
	m2 = Matrix([[1, 2], [3, 4], [5, 6]])

	with pytest.raises(ValueError, match="not aligned for matrix multiplication"):
		_ = m1 @ m2


def test_matrix_comparison_operators_return_matrix():
	m1 = Matrix([[1, 5], [3, 4]])
	m2 = Matrix([[2, 5], [1, 6]])

	less_than = m1 < m2
	equal_to = m1 == m1.copy()

	assert isinstance(less_than, Matrix)
	assert isinstance(equal_to, Matrix)
	assert np.array_equal(less_than.data, np.array([[True, False], [False, True]]))
	assert np.array_equal(equal_to.data, np.array([[True, True], [True, True]]))


def test_matrix_unary_and_copy_methods():
	m = Matrix([[1, -2], [-3, 4]])

	assert np.allclose((-m).data, np.array([[-1, 2], [3, -4]]))
	assert np.allclose((+m).data, np.array([[1, -2], [-3, 4]]))
	assert np.allclose(abs(m).data, np.array([[1, 2], [3, 4]]))

	cloned = m.copy()
	assert isinstance(cloned, Matrix)
	assert cloned is not m
	assert np.allclose(cloned.data, m.data)


def test_matrix_str_omits_trailing_decimal_for_integer_values():
	m = Matrix([[1, 2], [3, 4]])
	assert str(m) == "[[1 2]\n [3 4]]"


def test_matrix_transpose_property():
	m = Matrix([[1, 2], [3, 4]])
	transposed = m.T

	assert isinstance(transposed, Matrix)
	assert np.allclose(transposed.data, np.array([[1, 3], [2, 4]]))


def test_matrix_determinant_trace_and_inverse():
	m = Matrix([[4, 7], [2, 6]])

	assert np.isclose(m.determinant(), 10.0)
	assert np.isclose(m.trace(), 10.0)
	assert np.allclose(m.inverse().data, np.linalg.inv(np.array([[4, 7], [2, 6]])))


def test_matrix_square_only_operations_raise_for_non_square():
	m = Matrix([[1, 2, 3], [4, 5, 6]])

	with pytest.raises(ValueError, match="square matrices"):
		m.determinant()

	with pytest.raises(ValueError, match="square matrices"):
		m.inverse()

	with pytest.raises(ValueError, match="square matrices"):
		m.trace()

	with pytest.raises(ValueError, match="square matrices"):
		m.eigenvalues()

	with pytest.raises(ValueError, match="square matrices"):
		m.eigenvectors()


def test_matrix_eigenvalues_and_eigenvectors_shapes():
	m = Matrix([[2, 0], [0, 3]])

	values = m.eigenvalues()
	eigenvalues, eigenvectors = m.eigenvectors()

	assert np.allclose(np.sort(values), np.array([2.0, 3.0]))
	assert isinstance(eigenvalues, Vector)
	assert isinstance(eigenvectors, Matrix)
	assert eigenvalues.shape == (2,)
	assert eigenvectors.shape == (2, 2)


def test_matrix_radd_rsub_rmul():
	m1 = Matrix([[1, 2], [3, 4]])
	m2 = Matrix([[5, 6], [7, 8]])

	radd_result = m1.__radd__(m2)
	rsub_result = m1.__rsub__(m2)
	rmul_result = m1.__rmul__(m2)

	assert isinstance(radd_result, Matrix)
	assert isinstance(rsub_result, Matrix)
	assert isinstance(rmul_result, Matrix)
	assert np.allclose(radd_result.data, np.array([[6, 8], [10, 12]]))
	assert np.allclose(rsub_result.data, np.array([[4, 4], [4, 4]]))
	assert np.allclose(rmul_result.data, np.array([[5, 12], [21, 32]]))


def test_matrix_rtruediv_rmatmul():
	m1 = Matrix([[2, 4], [6, 8]])
	m2 = Matrix([[4, 8], [12, 16]])

	rtruediv_result = m1.__rtruediv__(m2)
	assert isinstance(rtruediv_result, Matrix)
	assert np.allclose(rtruediv_result.data, np.array([[2.0, 2.0], [2.0, 2.0]]))

	m3 = Matrix([[1, 2], [3, 4]])
	m4 = Matrix([[5, 6], [7, 8]])
	rmatmul_result = m3.__rmatmul__(m4)
	assert isinstance(rmatmul_result, Matrix)
	# m4 @ m3 = [[5,6],[7,8]] @ [[1,2],[3,4]] = [[5*1+6*3, 5*2+6*4], [7*1+8*3, 7*2+8*4]] = [[23,34], [31,46]]
	assert np.allclose(rmatmul_result.data, np.array([[23, 34], [31, 46]]))


def test_matrix_iadd_isub_imul():
	m1 = Matrix([[1, 2], [3, 4]])
	m2 = Matrix([[5, 6], [7, 8]])

	result = m1.__iadd__(m2)
	assert result is m1
	assert np.allclose(m1.data, np.array([[6, 8], [10, 12]]))

	m1 = Matrix([[10, 15], [20, 25]])
	m2 = Matrix([[2, 3], [4, 5]])
	result = m1.__isub__(m2)
	assert result is m1
	assert np.allclose(m1.data, np.array([[8, 12], [16, 20]]))

	m1 = Matrix([[2, 3], [4, 5]])
	m2 = Matrix([[6, 7], [8, 9]])
	result = m1.__imul__(m2)
	assert result is m1
	assert np.allclose(m1.data, np.array([[12, 21], [32, 45]]))


def test_matrix_itruediv():
	m1 = Matrix([[10, 20], [30, 40]])
	m2 = Matrix([[2, 4], [5, 8]])

	result = m1.__itruediv__(m2)
	assert result is m1
	assert np.allclose(m1.data, np.array([[5.0, 5.0], [6.0, 5.0]]))


def test_matrix_imatmul():
	m1 = Matrix([[1, 2], [3, 4]])
	m2 = Matrix([[5, 6], [7, 8]])

	result = m1.__imatmul__(m2)
	assert result is m1
	assert np.allclose(m1.data, np.array([[19, 22], [43, 50]]))
