import pytest
import numpy as np
from sorix import Tensor, tensor, no_grad

def test_tensor_initialization():
    """Test standard tensor creation and properties."""
    data = [1, 2, 3]
    t = tensor(data)
    assert np.array_equal(t.data, np.array(data))
    assert t.device == 'cpu'
    assert t.requires_grad == False
    assert t.grad is None

    t_grad = tensor(data, requires_grad=True)
    assert t_grad.requires_grad == True
    assert t_grad.grad is not None
    assert np.all(t_grad.grad == 0)

def test_basic_arithmetic_autograd():
    """Test +, -, *, / and their gradients."""
    a = tensor([2.0, 3.0], requires_grad=True)
    b = tensor([4.0, 5.0], requires_grad=True)
    
    # Addition
    c = a + b
    c.backward()
    assert np.array_equal(c.data, [6.0, 8.0])
    assert np.array_equal(a.grad, [1.0, 1.0])
    assert np.array_equal(b.grad, [1.0, 1.0])
    
    a.grad = None; b.grad = None
    
    # Subtraction
    d = a - b
    d.backward()
    assert np.array_equal(d.data, [-2.0, -2.0])
    assert np.array_equal(a.grad, [1.0, 1.0])
    assert np.array_equal(b.grad, [-1.0, -1.0])
    
    a.grad = None; b.grad = None

    # Multiplication
    e = a * b
    e.backward()
    assert np.array_equal(e.data, [8.0, 15.0])
    assert np.array_equal(a.grad, [4.0, 5.0]) # de/da = b
    assert np.array_equal(b.grad, [2.0, 3.0]) # de/db = a

    a.grad = None; b.grad = None
    
    # Division
    f = a / b
    f.backward()
    assert np.allclose(f.data, [0.5, 0.6])
    assert np.allclose(a.grad, [1/4.0, 1/5.0]) # df/da = 1/b
    assert np.allclose(b.grad, [-2.0/16.0, -3.0/25.0]) # df/db = -a/b^2

def test_matmul_autograd():
    """Test matrix multiplication and its gradients."""
    # (2, 3) @ (3, 2) -> (2, 2)
    a_data = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
    b_data = np.array([[7.0, 8.0], [9.0, 10.0], [11.0, 12.0]])
    
    a = tensor(a_data, requires_grad=True)
    b = tensor(b_data, requires_grad=True)
    
    c = a @ b
    c.backward()
    
    expected_c = a_data @ b_data
    assert np.array_equal(c.data, expected_c)
    
    # dc/da = grad_out @ b.T. Since grad_out is ones(2,2)
    # grad_a = ones(2,2) @ b.T
    expected_grad_a = np.ones((2, 2)) @ b_data.T
    expected_grad_b = a_data.T @ np.ones((2, 2))
    
    assert np.allclose(a.grad, expected_grad_a)
    assert np.allclose(b.grad, expected_grad_b)

def test_nonlinear_functions():
    """Test tanh and sigmoid and their gradients."""
    x = tensor([0.0, 1.0, -1.0], requires_grad=True)
    
    # Tanh
    y = x.tanh()
    y.backward()
    expected_y = np.tanh(x.data)
    assert np.allclose(y.data, expected_y)
    # dy/dx = 1 - tanh^2
    expected_grad = 1 - expected_y**2
    assert np.allclose(x.grad, expected_grad)

    x.grad = None
    # Sigmoid
    z = x.sigmoid()
    z.backward()
    expected_z = 1 / (1 + np.exp(-x.data))
    assert np.allclose(z.data, expected_z)
    # dz/dx = s * (1 - s)
    expected_grad_z = expected_z * (1 - expected_z)
    assert np.allclose(x.grad, expected_grad_z)

def test_reduction_ops():
    """Test sum and mean."""
    x_data = np.array([[1.0, 2.0], [3.0, 4.0]])
    x = tensor(x_data, requires_grad=True)
    
    # Sum
    s = x.sum()
    s.backward()
    assert s.data == 10.0
    assert np.array_equal(x.grad, np.ones((2, 2)))
    
    x.grad = None
    # Mean
    m = x.mean()
    m.backward()
    assert m.data == 2.5
    assert np.allclose(x.grad, np.full((2, 2), 1.0/4.0))

def test_broadcasting_autograd():
    """Test if gradients are correctly matched during broadcasting."""
    a = tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True) # (2, 2)
    b = tensor([10.0, 20.0], requires_grad=True)            # (2,) -> broadcast to (2, 2)
    
    c = a + b
    c.backward()
    
    assert np.array_equal(a.grad, np.ones((2, 2)))
    # b is broadcasted, so its gradient should be the sum over the broadcasted dimension
    assert np.array_equal(b.grad, [2.0, 2.0])

def test_no_grad():
    """Test the no_grad context manager."""
    a = tensor([1.0, 2.0], requires_grad=True)
    b = tensor([3.0, 4.0], requires_grad=True)
    
    with no_grad():
        c = a + b
    
    assert c.requires_grad == False
    assert len(c._prev) == 0

def test_complex_expression():
    """A more complex expression to verify chain rule."""
    x = tensor(2.0, requires_grad=True)
    y = tensor(3.0, requires_grad=True)
    
    # f(x, y) = x^2 * y + y^2
    f = (x**2) * y + (y**2)
    f.backward()
    
    # df/dx = 2xy = 2*2*3 = 12
    # df/dy = x^2 + 2y = 2^2 + 2*3 = 4 + 6 = 10
    assert x.grad == 12.0
    assert y.grad == 10.0

def test_tensor_inplace():
    a = tensor([1.0, 2.0])
    
    a.add_(1.0)
    assert np.array_equal(a.data, [2.0, 3.0])
    
    a.sub_(1.0)
    assert np.array_equal(a.data, [1.0, 2.0])
    
    a.mul_(2.0)
    assert np.array_equal(a.data, [2.0, 4.0])
    
    a.div_(2.0)
    assert np.array_equal(a.data, [1.0, 2.0])

def test_tensor_reshape_transpose():
    a = tensor([[1, 2, 3], [4, 5, 6]], requires_grad=True)
    
    # Reshape
    b = a.reshape(3, 2)
    assert b.shape == (3, 2)
    b.backward()
    assert a.grad.shape == (2, 3)
    
    a.grad = None
    # Flatten
    c = a.flatten()
    assert c.shape == (6,)
    c.backward()
    assert a.grad.shape == (2, 3)
    
    a.grad = None
    # Transpose
    d = a.transpose(1, 0)
    assert d.shape == (3, 2)
    d.backward()
    assert a.grad.shape == (2, 3)
    
    # T property
    assert a.T.shape == (3, 2)

def test_tensor_abs():
    a = tensor([-1.0, 2.0, -3.0])
    b = a.abs()
    assert np.array_equal(b.data, [1.0, 2.0, 3.0])

def test_tensor_comparisons():
    a = tensor([1, 2, 3])
    b = tensor([2, 2, 2])
    
    assert np.array_equal((a > b).data, [False, False, True])
    assert np.array_equal((a < b).data, [True, False, False])
    assert np.array_equal((a >= b).data, [False, True, True])
    assert np.array_equal((a <= b).data, [True, True, False])
    assert np.array_equal((a == b).data, [False, True, False])
    assert np.array_equal((a != b).data, [True, False, True])
    
    # Comparison with scalar
    assert np.array_equal((a > 2).data, [False, False, True])

def test_tensor_utility_methods():
    a = tensor([1.5, 2.5], requires_grad=True)
    
    # astype
    b = a.astype(int)
    assert b.dtype == int
    
    # item
    c = tensor([5.0])
    assert c.item() == 5.0
    
    # to_numpy
    assert isinstance(a.numpy(), np.ndarray)
    
    # __array__
    assert isinstance(np.array(a), np.ndarray)
    
    # __iter__
    items = [x for x in a]
    assert len(items) == 2
    
    # __str__, __repr__
    assert "tensor" in str(a).lower()
    assert "tensor" in repr(a).lower()

def test_tensor_dtypes():
    """Test explicit dtype setting and aliases."""
    import sorix
    
    # Using aliases
    t1 = tensor([1, 2], dtype=sorix.float32)
    assert t1.dtype == 'float32'
    assert t1.data.dtype == np.float32
    
    t2 = tensor([1, 2], dtype=sorix.int64)
    assert t2.dtype == 'int64'
    assert t2.data.dtype == np.int64
    
    t3 = tensor([1.5, 2.5], dtype=sorix.int32)
    assert t3.dtype == 'int32'
    assert np.array_equal(t3.data, [1, 2])
    
    # Verify backward respects dtype
    x = tensor([2], dtype=sorix.int64, requires_grad=True)
    y = x * x # 4
    y.backward()
    assert x.grad.dtype == x.dtype
    assert x.grad[0] == 4

def test_tensor_getitem_autograd():
    a = tensor([1.0, 2.0, 3.0], requires_grad=True)
    b = a[1]
    assert b.data == 2.0
    b.backward()
    assert np.array_equal(a.grad, [0.0, 1.0, 0.0])

def test_tensor_radd_rsub_rmul_rtruediv_rmatmul():
    a = tensor([2.0], requires_grad=True)
    
    # radd
    b = 1.0 + a
    assert b.data == 3.0
    
    # rsub
    c = 5.0 - a
    assert c.data == 3.0
    c.backward()
    assert a.grad == [-1.0]
    
    a.grad = None
    # rmul
    d = 3.0 * a
    assert d.data == 6.0
    d.backward()
    assert a.grad == [3.0]
    
    a.grad = None
    # rtruediv
    e = 10.0 / a
    assert e.data == 5.0
    e.backward()
    # d/da (10/a) = -10/a^2 = -10/4 = -2.5
    assert a.grad == [-2.5]
    
    # rmatmul
    x = tensor([[1.0, 2.0]])
    y = np.array([[3.0], [4.0]])
    z = y @ x # This uses __rmatmul__ of x
    assert z.shape == (2, 2)

def test_tensor_backward_no_grad_root():
    a = tensor([1.0, 2.0], requires_grad=False)
    a.backward() # Should initialize a.grad to ones
    assert np.array_equal(a.grad, [1.0, 1.0])

def test_tensor_to_gpu_error():
    from sorix.cupy.cupy import _cupy_available
    if not _cupy_available:
        a = tensor([1, 2])
        with pytest.raises(RuntimeError, match="CuPy is not installed"):
            a.to('cuda')
        with pytest.raises(ValueError, match="Invalid device type"):
            a.to("invalid")

def test_tensor_broadcasting_complex():
    a = tensor(np.ones((2, 1, 4)), requires_grad=True)
    b = tensor(np.ones((1, 3, 4)), requires_grad=True)
    c = a + b # (2, 3, 4)
    c.backward()
    assert a.grad.shape == (2, 1, 4)
    assert b.grad.shape == (1, 3, 4)
    assert np.all(a.grad == 3.0)
    assert np.all(b.grad == 2.0)

def test_tensor_gpu_ops_if_available():
    from sorix.cupy.cupy import _cupy_available
    if _cupy_available:
        import cupy as cp
        a = tensor([1.0, 2.0], device='cuda', requires_grad=True)
        b = tensor([3.0, 4.0], device='cuda', requires_grad=True)
        
        # Arithmetic
        c = a + b
        assert c.device == 'cuda'
        assert isinstance(c.data, cp.ndarray)
        
        # Non-linear with backward
        d = a.tanh()
        assert d.device == 'cuda'
        d.sum().backward()
        
        a.grad = cp.zeros_like(a.data)
        e = a.sigmoid()
        assert e.device == 'cuda'
        e.sum().backward()
        
        a.grad = cp.zeros_like(a.data)
        f = a.softmax()
        assert f.device == 'cuda'
        f.sum().backward()
        
        # Power
        p = a**2
        p.sum().backward()

def test_tensor_shape_ops_autograd():
    """Test reshape, transpose, and flatten with autograd."""
    a = tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
    
    # Reshape
    b = a.reshape(4, 1)
    loss = (b * tensor([[10.0], [20.0], [30.0], [40.0]])).sum()
    loss.backward()
    assert np.allclose(a.grad, [[10, 20], [30, 40]])
    a.grad = None
    
    # Transpose
    c = a.transpose(1, 0)
    loss_c = (c * tensor([[1.0, 2.0], [3.0, 4.0]])).sum()
    loss_c.backward()
    assert np.allclose(a.grad, [[1, 3], [2, 4]])
    a.grad = None
    
    # .T property
    d = a.T
    loss_d = (d * tensor([[1.0, 3.0], [2.0, 4.0]])).sum()
    loss_d.backward()
    assert np.allclose(a.grad, [[1, 2], [3, 4]])
    a.grad = None
    
    # Flatten
    e = a.flatten()
    loss_e = (e * tensor([1.0, 2.0, 3.0, 4.0])).sum()
    loss_e.backward()
    assert np.allclose(a.grad, [[1, 2], [3, 4]])

def test_tensor_reductions_autograd():
    """Test sum and mean with autograd and axes."""
    a = tensor([[1.0, 2.0], [3.0, 4.0]], requires_grad=True)
    
    # Sum over axis 0
    b = a.sum(axis=0) # [4, 6]
    loss = (b * tensor([1.0, 2.0])).sum()
    loss.backward()
    assert np.allclose(a.grad, [[1, 2], [1, 2]])
    a.grad = None
    
    # Mean with keepdims
    c = a.mean(axis=1, keepdims=True) # [[1.5], [3.5]]
    loss_c = (c * tensor([[10.0], [20.0]])).sum()
    loss_c.backward()
    # d_loss/d_c = [10, 20] -> d_loss/d_a = [5, 5], [10, 10]
    assert np.allclose(a.grad, [[5, 5], [10, 10]])

def test_tensor_advanced_math():
    """Test sqrt, log, exp with autograd via sorix.utils.math."""
    from sorix.utils import math
    x = tensor([1.0, 4.0, 0.5], requires_grad=True)
    
    # Sqrt
    y = math.sqrt(x)
    y.sum().backward()
    # d/dx(sqrt(x)) = 1/(2*sqrt(x))
    # 1/2, 1/4, 1/(2*sqrt(0.5))
    assert np.allclose(x.grad[:2], [0.5, 0.25])
    x.grad = None
    
    # Log (x+1 to avoid log(0))
    z = math.log(x + 1)
    z.sum().backward()
    # d/dx(log(x+1)) = 1/(x+1)
    assert np.allclose(x.grad, [1/2, 1/5, 1/1.5])
    x.grad = None
    
    # Exp
    w = math.exp(x)
    w.sum().backward()
    # d/dx(exp(x)) = exp(x)
    assert np.allclose(x.grad, np.exp([1.0, 4.0, 0.5]))

def test_no_grad_context():
    """Test no_grad disables tracking."""
    a = tensor([1.0], requires_grad=True)
    with no_grad():
        b = a * 2
        c = a + 2
        d = a - 2
        e = a / 2
        f = a.sum()
        g = a.mean()
        assert not any([b.requires_grad, c.requires_grad, d.requires_grad, e.requires_grad, f.requires_grad, g.requires_grad])
    
    assert (a * 2).requires_grad == True

def test_tensor_init_edge_cases():
    """Test initializing from Tensor and dtypes."""
    a = tensor([1.2, 3.4])
    b = tensor(a) # Copy from tensor
    assert np.array_equal(a.data, b.data)
    assert a is not b # Should be a copy
    
    # Repr with multiline
    long_t = tensor(np.arange(100))
    r = repr(long_t)
    assert "\n" in r or "..." in r # Formatting check

def test_tensor_indexing_advanced():
    """Test __getitem__ with more types."""
    t = tensor([[1, 2], [3, 4]], requires_grad=True)
    
    # Tuple indexing
    assert t[0, 1].item() == 2
    
    # Mask indexing (if supported, sorix likely supports it via numpy)
    mask = t > 2
    res = t[mask]
    assert np.array_equal(res.numpy(), [3, 4])

def test_tensor_reshape_transpose_extra():
    # Empty tensor reshape if allowed
    # (Checking if it hits some branches)
    z = tensor([])
    assert z.size == 0
    with pytest.raises(Exception):
        # Using a tuple as a single argument for shape
        z.reshape((2, 2))

def test_in_place_ops():
    """Test add_, sub_, mul_, div_."""
    a = tensor([10.0, 20.0])
    a.add_(5.0)
    assert np.array_equal(a.data, [15.0, 25.0])
    a.sub_(tensor([5.0, 5.0]))
    assert np.array_equal(a.data, [10.0, 20.0])
    a.mul_(2.0)
    assert np.array_equal(a.data, [20.0, 40.0])
    a.div_(10.0)
    assert np.array_equal(a.data, [2.0, 4.0])

def test_reverse_math():
    """Test __radd__, __rsub__, __rmul__, __rtruediv__."""
    a = tensor([2.0, 4.0], requires_grad=True)
    
    # 10 + a
    b = 10 + a
    assert np.array_equal(b.data, [12.0, 14.0])
    
    # 10 - a
    c = 10 - a
    c.sum().backward()
    assert np.array_equal(c.data, [8.0, 6.0])
    assert np.array_equal(a.grad, [-1.0, -1.0])
    a.grad = None
    
    # 10 * a
    d = 10 * a
    assert np.array_equal(d.data, [20.0, 40.0])
    
    # 10 / a
    e = 10 / a
    e.sum().backward()
    # -10/a^2 = -10/4, -10/16
    assert np.allclose(a.grad, [-2.5, -10/16.0])

def test_tensor_device_movement():
    """Test cpu() and to() movement."""
    a = tensor([1, 2])
    # to same device should be a no-op but returns self
    b = a.to('cpu')
    assert a is b
    
    # test dummy cuda move if not available
    c = a.cpu()
    assert c is a
    # Wait, let's check how sorix handles initialization of grad in backward.
