import numpy as np
from numpy.testing import (assert_allclose,
                           assert_equal)
import scipy.linalg.cython_blas as blas


class TestBlasInt:
    """Tests for blas_int type used in cython_blas/cython_lapack."""

    def test_blas_int_size(self):
        """Verify blas_int size matches the build configuration."""
        from scipy.__config__ import CONFIG
        size = blas._blas_int_size()
        cython_blas_ilp64 = CONFIG['Build Dependencies']['blas']['cython blas ilp64']
        if cython_blas_ilp64:
            assert size == 8, f"ILP64 build but blas_int is {size} bytes"
        else:
            assert size == 4, f"LP64 build but blas_int is {size} bytes"

        # also verify blas_bint size
        bint_sizes = blas._blas_bint_bint_sizes()  # (sizeof(blas_bint), sizeof(bint))
        if cython_blas_ilp64:
            assert bint_sizes == (8, 4)
        else:
            assert bint_sizes == (4, 4)

    def test_dgemm_with_blas_int_dimensions(self):
        """Test dgemm works correctly - exercises blas_int for dimensions."""
        a = np.eye(3, dtype=np.float64)
        b = np.arange(9, dtype=np.float64).reshape((3, 3))
        c = np.empty((3, 3), dtype=np.float64)
        blas._test_dgemm(1., a, b, 0., c)
        assert_allclose(c, b)


class TestDGEMM:

    def test_transposes(self):

        a = np.arange(12, dtype=np.float64).reshape((3, 4))[:2,:2]
        b = np.arange(1, 13, dtype=np.float64).reshape((4, 3))[:2,:2]
        c = np.empty((2, 4), dtype=np.float64)[:2,:2]

        blas._test_dgemm(1., a, b, 0., c)
        assert_allclose(c, a.dot(b))

        blas._test_dgemm(1., a.T, b, 0., c)
        assert_allclose(c, a.T.dot(b))

        blas._test_dgemm(1., a, b.T, 0., c)
        assert_allclose(c, a.dot(b.T))

        blas._test_dgemm(1., a.T, b.T, 0., c)
        assert_allclose(c, a.T.dot(b.T))

        blas._test_dgemm(1., a, b, 0., c.T)
        assert_allclose(c, a.dot(b).T)

        blas._test_dgemm(1., a.T, b, 0., c.T)
        assert_allclose(c, a.T.dot(b).T)

        blas._test_dgemm(1., a, b.T, 0., c.T)
        assert_allclose(c, a.dot(b.T).T)

        blas._test_dgemm(1., a.T, b.T, 0., c.T)
        assert_allclose(c, a.T.dot(b.T).T)

    def test_shapes(self):
        a = np.arange(6, dtype=np.float64).reshape((3, 2))
        b = np.arange(-6, 2, dtype=np.float64).reshape((2, 4))
        c = np.empty((3, 4), dtype=np.float64)

        blas._test_dgemm(1., a, b, 0., c)
        assert_allclose(c, a.dot(b))

        blas._test_dgemm(1., b.T, a.T, 0., c.T)
        assert_allclose(c, b.T.dot(a.T).T)

class TestWfuncPointers:
    """ Test the function pointers that are expected to fail on
    Mac OS X without the additional entry statement in their definitions
    in fblas_l1.pyf.src. """

    def test_complex_args(self):

        cx = np.array([.5 + 1.j, .25 - .375j, 12.5 - 4.j], np.complex64)
        cy = np.array([.8 + 2.j, .875 - .625j, -1. + 2.j], np.complex64)

        assert_allclose(blas._test_cdotc(cx, cy),
                        -17.6468753815+21.3718757629j)
        assert_allclose(blas._test_cdotu(cx, cy),
                        -6.11562538147+30.3156242371j)

        assert_equal(blas._test_icamax(cx), 3)

        assert_allclose(blas._test_scasum(cx), 18.625)
        assert_allclose(blas._test_scnrm2(cx), 13.1796483994)

        assert_allclose(blas._test_cdotc(cx[::2], cy[::2]),
                        -18.1000003815+21.2000007629j)
        assert_allclose(blas._test_cdotu(cx[::2], cy[::2]),
                        -6.10000038147+30.7999992371j)
        assert_allclose(blas._test_scasum(cx[::2]), 18.)
        assert_allclose(blas._test_scnrm2(cx[::2]), 13.1719398499)

    def test_double_args(self):

        x = np.array([5., -3, -.5], np.float64)
        y = np.array([2, 1, .5], np.float64)

        assert_allclose(blas._test_dasum(x), 8.5)
        assert_allclose(blas._test_ddot(x, y), 6.75)
        assert_allclose(blas._test_dnrm2(x), 5.85234975815)

        assert_allclose(blas._test_dasum(x[::2]), 5.5)
        assert_allclose(blas._test_ddot(x[::2], y[::2]), 9.75)
        assert_allclose(blas._test_dnrm2(x[::2]), 5.0249376297)

        assert_equal(blas._test_idamax(x), 1)

    def test_float_args(self):

        x = np.array([5., -3, -.5], np.float32)
        y = np.array([2, 1, .5], np.float32)

        assert_equal(blas._test_isamax(x), 1)

        assert_allclose(blas._test_sasum(x), 8.5)
        assert_allclose(blas._test_sdot(x, y), 6.75)
        assert_allclose(blas._test_snrm2(x), 5.85234975815)

        assert_allclose(blas._test_sasum(x[::2]), 5.5)
        assert_allclose(blas._test_sdot(x[::2], y[::2]), 9.75)
        assert_allclose(blas._test_snrm2(x[::2]), 5.0249376297)

    def test_double_complex_args(self):

        cx = np.array([.5 + 1.j, .25 - .375j, 13. - 4.j], np.complex128)
        cy = np.array([.875 + 2.j, .875 - .625j, -1. + 2.j], np.complex128)

        assert_equal(blas._test_izamax(cx), 3)

        assert_allclose(blas._test_zdotc(cx, cy), -18.109375+22.296875j)
        assert_allclose(blas._test_zdotu(cx, cy), -6.578125+31.390625j)

        assert_allclose(blas._test_zdotc(cx[::2], cy[::2]), -18.5625+22.125j)
        assert_allclose(blas._test_zdotu(cx[::2], cy[::2]), -6.5625+31.875j)