import unittest
from pythran.tests import TestEnv
import numpy
from pythran.typing import NDArray


class TestNumpyLinalg(TestEnv):

    def test_linalg_norm0(self):
        self.run_test("def linalg_norm0(x): from numpy.linalg import norm ; return norm(x)", numpy.arange(6.), linalg_norm0=[NDArray[float,:]])

    def test_linalg_norm1(self):
        self.run_test("def linalg_norm1(x): from numpy.linalg import norm ; return norm(x)", numpy.arange(6.).reshape(2,3), linalg_norm1=[NDArray[float,:,:]])

    def test_linalg_norm2(self):
        self.run_test("def linalg_norm2(x): from numpy.linalg import norm ; from numpy import inf ; return norm(x, inf)", numpy.arange(6.), linalg_norm2=[NDArray[float,:]])

    def test_linalg_norm3(self):
        self.run_test("def linalg_norm3(x): from numpy.linalg import norm ; from numpy import inf ; return norm(x, -inf)", numpy.arange(6.), linalg_norm3=[NDArray[float,:]])

    def test_linalg_norm4(self):
        self.run_test("def linalg_norm4(x): from numpy.linalg import norm ; from numpy import inf ; return norm(x, 0)", numpy.arange(6.), linalg_norm4=[NDArray[float,:]])

    def test_linalg_norm5(self):
        self.run_test("def linalg_norm5(x): from numpy.linalg import norm ; from numpy import inf ; return norm(x, ord=inf, axis=1)", (numpy.arange(9) - 4).reshape((3,3)), linalg_norm5=[NDArray[int,:,:]])

    def test_linalg_norm6(self):
        self.run_test("def linalg_norm6(x): from numpy.linalg import norm ; from numpy import inf ; return norm(x, ord=5, axis=(0,))", (numpy.arange(9) - 4).reshape((3,3)), linalg_norm6=[NDArray[int,:,:]])

    def test_linalg_norm7(self):
        self.run_test("def linalg_norm7(x): from numpy.linalg import norm ; return norm(x)", numpy.arange(6).reshape(2,3), linalg_norm7=[NDArray[int,:,:]])

    def test_linalg_norm8(self):
        self.run_test("def linalg_norm8(x): from numpy.linalg import norm ; return norm(x)",
                      numpy.arange(6).reshape(2,3) * 1j + 1,
                      linalg_norm8=[NDArray[complex,:,:]])

    def test_linalg_norm_pydoc(self):
        self.run_test('''
            def linalg_norm_pydoc(x):
                      import numpy as np
                      from numpy import linalg as LA
                      a = np.arange(9) - x
                      b = a.reshape((3, 3))
                      c = np.array([[ 1, 2, 3], [-1, 1, x]])
                      return (LA.norm(a), LA.norm(b),
                              LA.norm(a, np.inf), #LA.norm(b, np.inf),
                              LA.norm(a, -np.inf), #LA.norm(b, -np.inf),
                              LA.norm(a, 1), #LA.norm(b, 1),
                              LA.norm(a, -1), #LA.norm(b, -1),
                              LA.norm(a, 2), #LA.norm(b, 2),
                              LA.norm(a, -2), #LA.norm(b, -2),
                              LA.norm(a, 3), LA.norm(a, -3),
                              LA.norm(c, axis=0),
                              LA.norm(c, axis=1),
                              LA.norm(c, ord=1, axis=1),
                     )''',
                      10, linalg_norm_pydoc=[int])
