"""Unit tests for Tensor class."""

import pytest
import numpy as np
from synor_compute import Tensor
from synor_compute.types import Precision


class TestTensorCreation:
    """Tests for tensor creation methods."""

    def test_from_numpy_array(self):
        """Should create tensor from numpy array."""
        arr = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
        tensor = Tensor.from_numpy(arr)

        assert tensor.shape == (2, 3)
        assert tensor.size == 6
        assert tensor.ndim == 2
        assert tensor.dtype == Precision.FP32

    def test_from_numpy_with_dtype(self):
        """Should create tensor with specified dtype."""
        arr = np.array([1, 2, 3])
        tensor = Tensor.from_numpy(arr, dtype=Precision.FP16)

        assert tensor.dtype == Precision.FP16

    def test_infer_dtype_fp64(self):
        """Should infer FP64 from float64 array."""
        arr = np.array([1.0, 2.0], dtype=np.float64)
        tensor = Tensor.from_numpy(arr)

        assert tensor.dtype == Precision.FP64

    def test_infer_dtype_int8(self):
        """Should infer INT8 from int8 array."""
        arr = np.array([1, 2, 3], dtype=np.int8)
        tensor = Tensor.from_numpy(arr)

        assert tensor.dtype == Precision.INT8


class TestTensorFactoryMethods:
    """Tests for tensor factory methods."""

    def test_zeros(self):
        """Should create tensor of zeros."""
        tensor = Tensor.zeros((2, 3))

        assert tensor.shape == (2, 3)
        assert np.allclose(tensor.data, 0)

    def test_zeros_with_dtype(self):
        """Should create zeros tensor with specified dtype."""
        tensor = Tensor.zeros((2, 2), dtype=Precision.FP64)

        assert tensor.dtype == Precision.FP64
        assert tensor.data.dtype == np.float64

    def test_ones(self):
        """Should create tensor of ones."""
        tensor = Tensor.ones((3, 2))

        assert tensor.shape == (3, 2)
        assert np.allclose(tensor.data, 1)

    def test_random(self):
        """Should create tensor with random values in [0, 1)."""
        tensor = Tensor.random((10, 10))

        assert tensor.shape == (10, 10)
        assert tensor.data.min() >= 0
        assert tensor.data.max() < 1

    def test_randn(self):
        """Should create tensor with normal distribution."""
        tensor = Tensor.randn((1000,))

        # Check approximate normal distribution properties
        assert abs(tensor.data.mean()) < 0.2  # Mean ~= 0
        assert 0.8 < tensor.data.std() < 1.2  # Std ~= 1


class TestTensorOperations:
    """Tests for tensor operations."""

    def test_reshape(self):
        """Should reshape tensor."""
        tensor = Tensor.from_numpy(np.arange(6))
        reshaped = tensor.reshape((2, 3))

        assert reshaped.shape == (2, 3)
        assert reshaped.size == 6

    def test_reshape_preserves_data(self):
        """Should preserve data when reshaping."""
        arr = np.array([1, 2, 3, 4, 5, 6])
        tensor = Tensor.from_numpy(arr)
        reshaped = tensor.reshape((2, 3))

        assert reshaped.data[0, 0] == 1
        assert reshaped.data[1, 2] == 6

    def test_to_different_dtype(self):
        """Should convert to different dtype."""
        tensor = Tensor.from_numpy(np.array([1.5, 2.5], dtype=np.float32))
        converted = tensor.to(Precision.FP64)

        assert converted.dtype == Precision.FP64
        assert converted.data.dtype == np.float64

    def test_to_same_dtype_returns_self(self):
        """Should return self when converting to same dtype."""
        tensor = Tensor.from_numpy(np.array([1.0, 2.0], dtype=np.float32))
        same = tensor.to(Precision.FP32)

        assert same is tensor

    def test_numpy_returns_underlying_array(self):
        """Should return underlying numpy array."""
        arr = np.array([1, 2, 3], dtype=np.float32)
        tensor = Tensor.from_numpy(arr)

        assert np.array_equal(tensor.numpy(), arr)


class TestTensorProperties:
    """Tests for tensor properties."""

    def test_shape(self):
        """Should return correct shape."""
        tensor = Tensor.zeros((3, 4, 5))
        assert tensor.shape == (3, 4, 5)

    def test_size(self):
        """Should return correct size."""
        tensor = Tensor.zeros((3, 4, 5))
        assert tensor.size == 60

    def test_ndim(self):
        """Should return correct number of dimensions."""
        tensor = Tensor.zeros((3, 4, 5))
        assert tensor.ndim == 3

    def test_nbytes(self):
        """Should return correct byte size."""
        tensor = Tensor.zeros((10,), dtype=Precision.FP32)
        assert tensor.nbytes == 40  # 10 * 4 bytes


class TestTensorSerialization:
    """Tests for tensor serialization."""

    def test_serialize(self):
        """Should serialize tensor to dict."""
        tensor = Tensor.from_numpy(np.array([[1, 2], [3, 4]], dtype=np.float32))
        serialized = tensor.serialize()

        assert "data" in serialized
        assert "shape" in serialized
        assert "dtype" in serialized
        assert serialized["shape"] == [2, 2]
        assert serialized["dtype"] == "fp32"

    def test_deserialize(self):
        """Should deserialize tensor from dict."""
        original = Tensor.from_numpy(np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32))
        serialized = original.serialize()
        restored = Tensor.deserialize(serialized)

        assert restored.shape == original.shape
        assert restored.dtype == original.dtype
        assert np.array_equal(restored.data, original.data)

    def test_serialize_deserialize_fp64(self):
        """Should preserve FP64 precision."""
        original = Tensor.from_numpy(np.array([1.234567890123456], dtype=np.float64))
        serialized = original.serialize()
        restored = Tensor.deserialize(serialized)

        assert restored.dtype == Precision.FP64
        assert np.allclose(restored.data, original.data)

    def test_serialize_deserialize_large_tensor(self):
        """Should handle large tensors."""
        original = Tensor.random((100, 100))
        serialized = original.serialize()
        restored = Tensor.deserialize(serialized)

        assert restored.shape == original.shape
        assert np.allclose(restored.data, original.data)


class TestTensorRepr:
    """Tests for tensor string representations."""

    def test_repr(self):
        """Should return informative repr."""
        tensor = Tensor.zeros((2, 3), dtype=Precision.FP32)
        rep = repr(tensor)

        assert "Tensor" in rep
        assert "(2, 3)" in rep
        assert "fp32" in rep

    def test_str(self):
        """Should return string with data."""
        tensor = Tensor.from_numpy(np.array([1, 2, 3], dtype=np.float32))
        s = str(tensor)

        assert "Tensor" in s
        assert "fp32" in s


class TestTensorEdgeCases:
    """Tests for edge cases."""

    def test_scalar_like_tensor(self):
        """Should handle 1-element tensor."""
        tensor = Tensor.from_numpy(np.array([42], dtype=np.float32))

        assert tensor.shape == (1,)
        assert tensor.size == 1
        assert tensor.ndim == 1

    def test_empty_shape_dimension(self):
        """Should handle tensors with zero-size dimensions."""
        tensor = Tensor.zeros((0, 3))

        assert tensor.shape == (0, 3)
        assert tensor.size == 0

    def test_high_dimensional_tensor(self):
        """Should handle high-dimensional tensors."""
        tensor = Tensor.zeros((2, 3, 4, 5, 6))

        assert tensor.ndim == 5
        assert tensor.size == 720