import numpy as np
import pytest

from pandas import DataFrame, Series
import pandas._testing as tm


class DotSharedTests:
    @pytest.fixture
    def obj(self):
        raise NotImplementedError

    @pytest.fixture
    def other(self) -> DataFrame:
        """
        other is a DataFrame that is indexed so that obj.dot(other) is valid
        """
        raise NotImplementedError

    @pytest.fixture
    def expected(self, obj, other) -> DataFrame:
        """
        The expected result of obj.dot(other)
        """
        raise NotImplementedError

    @classmethod
    def reduced_dim_assert(cls, result, expected):
        """
        Assertion about results with 1 fewer dimension that self.obj
        """
        raise NotImplementedError

    def test_dot_equiv_values_dot(self, obj, other, expected):
        # `expected` is constructed from obj.values.dot(other.values)
        result = obj.dot(other)
        tm.assert_equal(result, expected)

    def test_dot_2d_ndarray(self, obj, other, expected):
        # Check ndarray argument; in this case we get matching values,
        #  but index/columns may not match
        result = obj.dot(other.values)
        assert np.all(result == expected.values)

    def test_dot_1d_ndarray(self, obj, expected):
        # can pass correct-length array
        row = obj.iloc[0] if obj.ndim == 2 else obj

        result = obj.dot(row.values)
        expected = obj.dot(row)
        self.reduced_dim_assert(result, expected)

    def test_dot_series(self, obj, other, expected):
        # Check series argument
        result = obj.dot(other["1"])
        self.reduced_dim_assert(result, expected["1"])

    def test_dot_series_alignment(self, obj, other, expected):
        result = obj.dot(other.iloc[::-1]["1"])
        self.reduced_dim_assert(result, expected["1"])

    def test_dot_aligns(self, obj, other, expected):
        # Check index alignment
        other2 = other.iloc[::-1]
        result = obj.dot(other2)
        tm.assert_equal(result, expected)

    def test_dot_shape_mismatch(self, obj):
        msg = "Dot product shape mismatch"
        # exception raised is of type Exception
        with pytest.raises(Exception, match=msg):
            obj.dot(obj.values[:3])

    def test_dot_misaligned(self, obj, other):
        msg = "matrices are not aligned"
        with pytest.raises(ValueError, match=msg):
            obj.dot(other.T)


class TestSeriesDot(DotSharedTests):
    @pytest.fixture
    def obj(self):
        return Series(np.random.randn(4), index=["p", "q", "r", "s"])

    @pytest.fixture
    def other(self):
        return DataFrame(
            np.random.randn(3, 4), index=["1", "2", "3"], columns=["p", "q", "r", "s"]
        ).T

    @pytest.fixture
    def expected(self, obj, other):
        return Series(np.dot(obj.values, other.values), index=other.columns)

    @classmethod
    def reduced_dim_assert(cls, result, expected):
        """
        Assertion about results with 1 fewer dimension that self.obj
        """
        tm.assert_almost_equal(result, expected)


class TestDataFrameDot(DotSharedTests):
    @pytest.fixture
    def obj(self):
        return DataFrame(
            np.random.randn(3, 4), index=["a", "b", "c"], columns=["p", "q", "r", "s"]
        )

    @pytest.fixture
    def other(self):
        return DataFrame(
            np.random.randn(4, 2), index=["p", "q", "r", "s"], columns=["1", "2"]
        )

    @pytest.fixture
    def expected(self, obj, other):
        return DataFrame(
            np.dot(obj.values, other.values), index=obj.index, columns=other.columns
        )

    @classmethod
    def reduced_dim_assert(cls, result, expected):
        """
        Assertion about results with 1 fewer dimension that self.obj
        """
        tm.assert_series_equal(result, expected, check_names=False)
        assert result.name is None