# Copyright (c) 2017 Sony Corporation. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from six import exec_

import pytest
import numpy as np
import nnabla as nn
import nnabla.functions as F


@pytest.mark.parametrize("seed", [313])
@pytest.mark.parametrize("op", ["+", "-", "*", "/", "**"])
@pytest.mark.parametrize("x_var, y_var", [(False, False), (False, True), (True, False)])
def test_ndarray_arithmetic_ops2(seed, op, x_var, y_var):
    rng = np.random.RandomState(seed)
    shape = [2, 3, 4]
    vx_data = rng.randn(*shape).astype(np.float32)
    vy_data = rng.randn(*shape).astype(np.float32)
    if op == "**":
        vx_data += - vx_data.min() + 1.0
    if x_var:
        vx = nn.Variable.from_numpy_array(vx_data)
    else:
        vx = nn.NdArray.from_numpy_array(vx_data)
    if y_var:
        vy = nn.Variable.from_numpy_array(vy_data)
    else:
        vy = nn.NdArray.from_numpy_array(vy_data)
    vz = eval("vx {0} vy".format(op))
    ref_z = eval("vx_data {0} vy_data".format(op))
    assert np.allclose(ref_z, vz.data)

    if x_var:
        return

    # Inplace test
    vx_bak = vx
    exec_("vx {0}= vy".format(op))
    assert np.allclose(vx.data, vz.data)
    assert vx is vx_bak


@pytest.mark.parametrize("seed", [313, 314])
@pytest.mark.parametrize("op", ["+", "-", "*", "/", "**"])
def test_ndarray_arithmetic_scalar_ops(seed, op):
    rng = np.random.RandomState(seed)
    vx = nn.NdArray.from_numpy_array(rng.randn(2, 3, 4).astype(np.float32))
    a = rng.randn()
    if op == "**":
        vx.data += - vx.data.min() + 1.0
    vz = eval("vx {0} a".format(op))
    ref_z = eval("vx.data {0} a".format(op))
    assert np.allclose(ref_z, vz.data)

    # Inplace test
    vx_bak = vx
    exec_("vx {0}= a".format(op))
    assert np.allclose(vx.data, vz.data)
    assert vx is vx_bak


@pytest.mark.parametrize("seed", [313, 314])
@pytest.mark.parametrize("op", ["+", "-", "*", "/", "**"])
def test_ndarray_arithmetic_scalar_rops(seed, op):
    rng = np.random.RandomState(seed)
    vx = nn.NdArray.from_numpy_array(rng.randn(2, 3, 4).astype(np.float32))
    a = rng.randn()
    if op == "**":
        a = np.abs(a)
    vz = eval("a {0} vx".format(op))
    ref_z = eval("a {0} vx.data".format(op))
    assert np.allclose(ref_z, vz.data)


@pytest.mark.parametrize("seed", [313, 314])
@pytest.mark.parametrize("op", ["+", "-"])
def test_ndarray_arithmetic_unary_ops(seed, op):
    rng = np.random.RandomState(seed)
    vx = nn.NdArray.from_numpy_array(rng.randn(2, 3, 4).astype(np.float32))
    vz = eval("{0} vx".format(op))
    ref_z = eval("{0} vx.data".format(op))
    assert np.allclose(ref_z, vz.data)

@@ -895,16 +895,69 @@
 alse)%5D)%0A
+@pytest.mark.parametrize(%22shape%22, %5B(2, 3, 4), (0,)%5D)%0A
 def test
@@ -1003,16 +1003,23 @@
 r, y_var
+, shape
 ):%0A    r
@@ -1055,30 +1055,8 @@
 ed)%0A
-    shape = %5B2, 3, 4%5D%0A
     
@@ -1162,32 +1162,53 @@
    if op == %22**%22
+ and vx_data.size %3E 0
 :%0A        vx_dat
@@ -1876,32 +1876,85 @@
 *%22, %22/%22, %22**%22%5D)%0A
+@pytest.mark.parametrize(%22shape%22, %5B(2, 3, 4), (0,)%5D)%0A
 def test_ndarray
@@ -1976,32 +1976,39 @@
 lar_ops(seed, op
+, shape
 ):%0A    rng = np.
@@ -2071,39 +2071,38 @@
 array(rng.randn(
-2, 3, 4
+*shape
 ).astype(np.floa
@@ -2136,32 +2136,48 @@
    if op == %22**%22
+ and vx.size %3E 0
 :%0A        vx.dat
@@ -2561,24 +2561,77 @@
 %22/%22, %22**%22%5D)%0A
+@pytest.mark.parametrize(%22shape%22, %5B(2, 3, 4), (0,)%5D)%0A
 def test_nda
@@ -2658,32 +2658,39 @@
 ar_rops(seed, op
+, shape
 ):%0A    rng = np.
@@ -2753,39 +2753,38 @@
 array(rng.randn(
-2, 3, 4
+*shape
 ).astype(np.floa
@@ -3061,16 +3061,69 @@
 , %22-%22%5D)%0A
+@pytest.mark.parametrize(%22shape%22, %5B(2, 3, 4), (0,)%5D)%0A
 def test
@@ -3160,16 +3160,23 @@
 seed, op
+, shape
 ):%0A    r
@@ -3255,23 +3255,22 @@
 g.randn(
-2, 3, 4
+*shape
 ).astype

Test basic ops for NdArray with size 0.