from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from caffe2.proto import caffe2_pb2
from caffe2.python import core, workspace
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.serialized_test.serialized_test_util as serial

from hypothesis import given
import hypothesis.strategies as st
import numpy as np


def _fill_diagonal(shape, value):
    result = np.zeros(shape)
    np.fill_diagonal(result, value)
    return (result,)


class TestFillerOperator(serial.SerializedTestCase):

    @given(**hu.gcs)
    def test_shape_error(self, gc, dc):
        op = core.CreateOperator(
            'GaussianFill',
            [],
            'out',
            shape=32,  # illegal parameter
            mean=0.0,
            std=1.0,
        )
        exception = False
        try:
            workspace.RunOperatorOnce(op)
        except Exception:
            exception = True
        self.assertTrue(exception, "Did not throw exception on illegal shape")

        op = core.CreateOperator(
            'ConstantFill',
            [],
            'out',
            shape=[],  # scalar
            value=2.0,
        )
        exception = False
        self.assertTrue(workspace.RunOperatorOnce(op))
        self.assertEqual(workspace.FetchBlob('out'), [2.0])

    @given(**hu.gcs)
    def test_int64_shape(self, gc, dc):
        large_dim = 2 ** 31 + 1
        net = core.Net("test_shape_net")
        net.UniformFill(
            [],
            'out',
            shape=[0, large_dim],
            min=0.0,
            max=1.0,
        )
        self.assertTrue(workspace.CreateNet(net))
        self.assertTrue(workspace.RunNet(net.Name()))
        self.assertEqual(workspace.blobs['out'].shape, (0, large_dim))

    @given(
        shape=hu.dims().flatmap(
            lambda dims: hu.arrays(
                [dims], dtype=np.int64,
                elements=st.integers(min_value=0, max_value=20)
            )
        ),
        a=st.integers(min_value=0, max_value=100),
        b=st.integers(min_value=0, max_value=100),
        **hu.gcs
    )
    def test_uniform_int_fill_op_blob_input(self, shape, a, b, gc, dc):
        net = core.Net('test_net')

        with core.DeviceScope(core.DeviceOption(caffe2_pb2.CPU)):
            shape_blob = net.Const(shape, dtype=np.int64)
        a_blob = net.Const(a, dtype=np.int32)
        b_blob = net.Const(b, dtype=np.int32)
        uniform_fill = net.UniformIntFill([shape_blob, a_blob, b_blob],
                                          1, input_as_shape=1)

        workspace.RunNetOnce(net)

        blob_out = workspace.FetchBlob(uniform_fill)
        if b < a:
            new_shape = shape[:]
            new_shape[0] = 0
            np.testing.assert_array_equal(new_shape, blob_out.shape)
        else: