"""Tests for NumpyChain basic functionality."""

import numpy as np
import pytest

from crystflux.v1 import Crystallizer
from crystflux.v1.adapters import enable_dream_mode
from crystflux.v1.adapters.numpy_chain import NumpyChain, as_numpy_chain


class TestNumpyChain:
    """Test case for NumpyChain core functionality."""

    def test_initialization(self):
        """Test NumpyChain initialization."""
        array = np.array([1, 2, 3])
        context = NumpyChain(array)

        assert context.array is array
        assert context._array.shape == (3,)  # type:ignore
        assert str(context) == "[1 2 3]"

    def test_array_property(self):
        """Test access to the underlying array."""
        array = np.array([[1, 2], [3, 4]])
        context = NumpyChain(array)

        assert context.shape == (2, 2)
        assert context.dtype == np.int64  # Default dtype for integer arrays

    def test_repr_and_str(self) -> None:
        """__repr__ / __str__ include shape and dtype."""
        array = np.array([[1, 2], [3, 4]], dtype=np.int64)
        context = NumpyChain(array)

        assert repr(context) == "NumpyChain(shape=(2, 2), dtype=int64)"
        assert str(context) == "[[1 2]\n [3 4]]"

    def test_to_cryst(self):
        """Test conversion back to CrystFlux values."""
        array = np.array([1, 2, 3])
        context = NumpyChain(array)

        strict_result = context.to_cryst()
        assert type(strict_result).__name__ == "StrictValue"
        assert strict_result.as_json_value() == (1, 2, 3)

        latent_context = as_numpy_chain(Crystallizer.latent([1, 2, 3]))
        latent_result = latent_context.to_cryst()
        assert type(latent_result).__name__ == "LatentValue"
        assert latent_result.as_json_value() == (1, 2, 3)

        with enable_dream_mode():
            dream_context = as_numpy_chain(Crystallizer.dream([1, 2, 3]))
            dream_result = dream_context.to_cryst()
        assert type(dream_result).__name__ == "DreamValue"
        assert dream_result.as_json_value() == (1, 2, 3)


class TestAsNumpy:
    """Test case for the as_numpy_chain function."""

    def test_as_numpy_with_list(self):
        """Test as_numpy_chain with list data."""
        value = Crystallizer.strict([1, 2, 3, 4, 5])
        context = as_numpy_chain(value)

        assert isinstance(context, NumpyChain)
        assert np.array_equal(context.array, np.array([1, 2, 3, 4, 5]))

    def test_round_trip_conversion(self):
        """Test round-trip conversion: Value → NumpyChain → Value."""
        original_value = Crystallizer.strict([1, 2, 3])

        result = as_numpy_chain(original_value).to_cryst()
        assert result.as_json_value() == original_value.as_json_value()

    def test_as_numpy_with_dream_mode(self) -> None:
        """Test as_numpy_chain with dream mode."""
        with enable_dream_mode():
            value = Crystallizer.dream([1, 2, 3, 4, 5])
            context = as_numpy_chain(value)

            assert isinstance(context, NumpyChain)
            assert context._vesicle.mode == "dream"
            result = context.to_cryst()
        assert type(result).__name__ == "DreamValue"
        assert result.as_json_value() == (1, 2, 3, 4, 5)


class TestMethodChaining:
    """Test case for method chaining functionality."""

    def test_single_array_method(self):
        """Test a single array method."""
        value = Crystallizer.strict([1, 2, 3, 4, 5])
        result = as_numpy_chain(value).mean()

        # mean returns Python scalar
        assert isinstance(result, (np.generic, float, int))
        assert result == 3.0

    def test_single_array_method_chain(self):
        """Test a single method in the chain that returns an array."""
        value = Crystallizer.strict([1, 2, 3, 4])
        result = as_numpy_chain(value).reshape(2, 2).to_cryst()

        assert result.as_json_value() == ((1, 2), (3, 4))

    def test_multiple_method_chain(self):
        """Test multiple methods that return arrays."""
        value = Crystallizer.strict([1, 2, 3, 4])
        result = as_numpy_chain(value).reshape(2, 2).transpose().to_cryst()

        assert result.as_json_value() == ((1, 3), (2, 4))

    def test_chaining_with_scalar_result(self):
        """Test chaining that ends with a scalar operation."""
        value = Crystallizer.strict([[1, 2], [3, 4]])
        result = as_numpy_chain(value).transpose().flatten().sum()

        # sum returns Python scalar
        assert isinstance(result, (np.generic, float, int))
        assert result == 10

    def test_method_wrapper_returns_python_sequence(self):
        """Ensure non-ndarray/non-generic results are returned raw (e.g., tolist)."""
        value = Crystallizer.strict([[1, 2], [3, 4]])
        result = as_numpy_chain(value).reshape(2, 2).tolist()

        assert isinstance(result, list)
        assert result == [[1, 2], [3, 4]]

    def test_numpy_function_consumes_context(self):
        """NumPy native functions should work with NumpyChain via __array__."""
        value = Crystallizer.strict([1, 2, 3, 4])
        context = as_numpy_chain(value)
        np_mean = np.mean(context)
        assert np_mean == 2.5


class TestMapFilterIntegration:
    """Integration of strict map/filter with NumPy operations."""

    def test_map_mean_scalar_strict(self):
        """map + NumPy mean returns Python scalars (tuple after freeze)."""
        value = Crystallizer.strict([[1, 2, 3], [4, 5, 6]])

        result = value.map(lambda v: Crystallizer.strict(as_numpy_chain(v).mean()))

        assert result.as_json_value() == (2.0, 5.0)

    def test_map_reshape_transpose_strict(self):
        """map + reshape/transpose and back to cryst per element."""
        value = Crystallizer.strict([[1, 2, 3, 4], [5, 6, 7, 8]])

        result = value.map(lambda v: as_numpy_chain(v).reshape(2, 2).transpose().to_cryst())

        assert result.as_json_value() == (((1, 3), (2, 4)), ((5, 7), (6, 8)))

    def test_filter_sum_positive_strict(self):
        """filter with NumPy sum keeps positive-sum arrays only."""
        value = Crystallizer.strict([[1, 2], [0, 0], [3, 4]])

        result = value.filter(lambda v: as_numpy_chain(v).sum() > 0)

        assert result.as_json_value() == ((1, 2), (3, 4))

    def test_filter_shape_match_strict(self):
        """filter by NumPy shape to keep only 2-length arrays."""
        value = Crystallizer.strict([[1, 2], [1, 2, 3], [4, 5]])

        result = value.filter(lambda v: as_numpy_chain(v).shape == (2,))

        assert result.as_json_value() == ((1, 2), (4, 5))

    def test_complex_chain_with_mask_and_repeat(self):
        """Combine transpose/ravel/compress/repeat to ensure chaining stays coherent."""
        value = Crystallizer.strict([[1, -2, 0], [3, 4, -1]])
        ctx = as_numpy_chain(value)

        tflat = ctx.transpose().ravel().array
        mask = (tflat > 0).tolist()  # ndarray bool mask
        result = tflat.compress(mask).repeat(2).tolist()

        assert result == [1, 1, 3, 3, 4, 4]


class TestErrorHandling:
    """Test case for basic error handling in NumpyChain."""

    def test_invalid_attribute(self):
        """Test accessing non-existent attributes."""
        array = np.array([1, 2, 3])
        context = NumpyChain(array)

        with pytest.raises(AttributeError):
            _ = context.nonexistent_method

    def test_invalid_method_call(self):
        """Test calling methods with invalid arguments."""
        array = np.array([1, 2, 3])
        context = NumpyChain(array)

        with pytest.raises((ValueError, TypeError)):
            _ = context.reshape(-1, -1)  # Invalid reshape dimensions
