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

import hypothesis.strategies as st
import numpy as np
import numpy.testing as npt
import unittest
from hypothesis import given

import caffe2.python.hypothesis_test_util as hu

from caffe2.python import (
    layer_model_instantiator,
    core,
    schema,
    workspace,
)
from caffe2.python.layers.layers import (
    InstantiationContext,
)
from caffe2.python.layers.tags import Tags
from caffe2.python.layer_test_util import (
    LayersTestCase,
    OpSpec,
)
from caffe2.python.layers.layers import (
    IdList,
    set_request_only,
    is_request_only_scalar,
    get_key,
)

import logging
logger = logging.getLogger(__name__)


class TestLayers(LayersTestCase):
    def testAddLoss(self):
        input_record_LR = self.new_record(
            schema.Struct(
                ('label', schema.Scalar((np.float64, (1, )))),
                ('logit', schema.Scalar((np.float32, (2, )))),
                ('weight', schema.Scalar((np.float64, (1, ))))
            )
        )
        loss_LR = self.model.BatchLRLoss(input_record_LR)

        self.model.add_loss(loss_LR)
        assert 'unnamed' in self.model.loss
        self.assertEqual(
            schema.Scalar((np.float32, tuple())), self.model.loss.unnamed
        )
        self.assertEqual(loss_LR, self.model.loss.unnamed)

        self.model.add_loss(loss_LR, 'addLoss')
        assert 'addLoss' in self.model.loss
        self.assertEqual(
            schema.Scalar((np.float32, tuple())), self.model.loss.addLoss
        )
        self.assertEqual(loss_LR, self.model.loss.addLoss)

        self.model.add_loss(
            schema.Scalar(
                dtype=np.float32, blob=core.BlobReference('loss_blob_1')
            ), 'addLoss'
        )
        assert 'addLoss_auto_0' in self.model.loss
        self.assertEqual(
            schema.Scalar((np.float32, tuple())), self.model.loss.addLoss_auto_0
        )
        assert core.BlobReference('loss_blob_1') in self.model.loss.field_blobs()

        self.model.add_loss(
            schema.Struct(
                (
                    'structName', schema.Scalar(
                        dtype=np.float32,
                        blob=core.BlobReference('loss_blob_2')
                    )
                )
            ), 'addLoss'
        )
        assert 'addLoss_auto_1' in self.model.loss
        self.assertEqual(
            schema.Struct(('structName', schema.Scalar((np.float32, tuple())))),
            self.model.loss.addLoss_auto_1
        )
        assert core.BlobReference('loss_blob_2') in self.model.loss.field_blobs()

        loss_in_tuple_0 = schema.Scalar(
            dtype=np.float32, blob=core.BlobReference('loss_blob_in_tuple_0')
        )

        loss_in_tuple_1 = schema.Scalar(
            dtype=np.float32, blob=core.BlobReference('loss_blob_in_tuple_1')
        )

        loss_tuple = schema.NamedTuple(
            'loss_in_tuple', * [loss_in_tuple_0, loss_in_tuple_1]
        )
        self.model.add_loss(loss_tuple, 'addLoss')
        assert 'addLoss_auto_2' in self.model.loss
        self.assertEqual(
            schema.Struct(
                ('loss_in_tuple_0', schema.Scalar((np.float32, tuple()))),
                ('loss_in_tuple_1', schema.Scalar((np.float32, tuple())))
            ), self.model.loss.addLoss_auto_2
        )
        assert core.BlobReference('loss_blob_in_tuple_0')\
         in self.model.loss.field_blobs()
        assert core.BlobReference('loss_blob_in_tuple_1')\
         in self.model.loss.field_blobs()

    def testAddOutputSchema(self):
        # add the first field
        self.model.add_output_schema('struct', schema.Struct())
        expected_output_schema = schema.Struct(('struct', schema.Struct()))
        self.assertEqual(
            self.model.output_schema,
            expected_output_schema,
        )

        # add the second field
        self.model.add_output_schema('scalar', schema.Scalar(np.float64))
        expected_output_schema = schema.Struct(
            ('struct', schema.Struct()),
            ('scalar', schema.Scalar(np.float64)),
        )
        self.assertEqual(
            self.model.output_schema,
            expected_output_schema,
        )

        # overwrite a field should raise
        with self.assertRaises(AssertionError):
            self.model.add_output_schema('scalar', schema.Struct())

    def _test_net(self, net, ops_list):
        """
        Helper function to assert the net contains some set of operations and
        then to run the net.

        Inputs:
            net -- the network to test and run
            ops_list -- the list of operation specifications to check for
                        in the net
        """
        ops_output = self.assertNetContainOps(net, ops_list)
        workspace.RunNetOnce(net)
        return ops_output

    def testFCWithoutBias(self):
        output_dims = 2
        fc_without_bias = self.model.FCWithoutBias(
            self.model.input_feature_schema.float_features, output_dims)
        self.model.output_schema = fc_without_bias

        self.assertEqual(
            schema.Scalar((np.float32, (output_dims, ))),
            fc_without_bias
        )

        train_init_net, train_net = self.get_training_nets()

        init_ops = self.assertNetContainOps(
            train_init_net,
            [
                OpSpec("UniformFill", None, None),
            ]
        )

        mat_mul_spec = OpSpec(
            "MatMul",
            [
                self.model.input_feature_schema.float_features(),
                init_ops[0].output[0],
            ],
            fc_without_bias.field_blobs()
        )

        self.assertNetContainOps(train_net, [mat_mul_spec])

        predict_net = self.get_predict_net()
        self.assertNetContainOps(predict_net, [mat_mul_spec])

    def testSparseLookupSumPooling(self):
        record = schema.NewRecord(self.model.net, schema.Struct(
            ('sparse', schema.Struct(
                ('sparse_feature_0', schema.List(
                    schema.Scalar(np.int64,
                                  metadata=schema.Metadata(categorical_limit=1000)))),
            )),
        ))
        embedding_dim = 64
        embedding_after_pooling = self.model.SparseLookup(
            record.sparse.sparse_feature_0, [embedding_dim], 'Sum')
        self.model.output_schema = schema.Struct()
        self.assertEqual(
            schema.Scalar((np.float32, (embedding_dim, ))),
            embedding_after_pooling
        )

        train_init_net, train_net = self.get_training_nets()

        init_ops = self.assertNetContainOps(
            train_init_net,
            [
                OpSpec("UniformFill", None, None),
                OpSpec("ConstantFill", None, None),
            ]
        )
        sparse_lookup_op_spec = OpSpec(
            'SparseLengthsSum',
            [
                init_ops[0].output[0],
                record.sparse.sparse_feature_0.items(),
                record.sparse.sparse_feature_0.lengths(),
            ],
            [embedding_after_pooling()]
        )
        self.assertNetContainOps(train_net, [sparse_lookup_op_spec])

        predict_net = self.get_predict_net()
        self.assertNetContainOps(predict_net, [sparse_lookup_op_spec])

    @given(
        use_hashing=st.booleans(),
        modulo=st.integers(min_value=100, max_value=200),
    )
    def testSparseFeatureHashIdList(self, use_hashing, modulo):
        record = schema.NewRecord(
            self.model.net,
            schema.List(schema.Scalar(
                np.int64,
                metadata=schema.Metadata(categorical_limit=60000)
            ))
        )
        output_schema = self.model.SparseFeatureHash(
            record,
            modulo=modulo,
            use_hashing=use_hashing)

        self.model.output_schema = output_schema

        self.assertEqual(len(self.model.layers), 1)
        self.assertEqual(output_schema._items.metadata.categorical_limit,
                modulo)
        train_init_net, train_net = self.get_training_nets()

    @given(
        use_hashing=st.booleans(),
        modulo=st.integers(min_value=100, max_value=200),
    )
    def testSparseFeatureHashIdScoreList(self, use_hashing, modulo):
        record = schema.NewRecord(self.model.net,
                schema.Map(schema.Scalar(np.int64,
                    metadata=schema.Metadata(
                        categorical_limit=60000)),
                    np.float32))

        output_schema = self.model.SparseFeatureHash(
            record,
            modulo=modulo,
            use_hashing=use_hashing)

        self.model.output_schema = output_schema

        self.assertEqual(len(self.model.layers), 1)
        self.assertEqual(output_schema._items.keys.metadata.categorical_limit,
                modulo)
        train_init_net, train_net = self.get_training_nets()

    def testSparseLookupIncorrectPositionWeightedOnIdList(self):
        '''
        Currently the implementation of SparseLookup assumed input is id_score_list
        when use PositionWeighted.
        '''
        record = schema.NewRecord(self.model.net, schema.Struct(
            ('sparse', schema.Struct(
                ('sparse_feature_0', schema.List(
                    schema.Scalar(np.int64,
                                  metadata=schema.Metadata(categorical_limit=1000)))),
            )),
        ))

        embedding_dim = 64
        with self.assertRaises(AssertionError):
            self.model.SparseLookup(
                record.sparse.sparse_feature_0, [embedding_dim], 'PositionWeighted')

    def testSparseLookupPositionWeightedOnIdList(self):
        record = schema.NewRecord(self.model.net, schema.Struct(
            ('sparse', schema.Struct(
                ('sparse_feature_0', schema.List(
                    schema.Scalar(np.int64,
                                  metadata=schema.Metadata(categorical_limit=1000)))),
            )),
        ))

        # convert id_list to id_score_list with PositionWeighted layer
        sparse_segment = record.sparse.sparse_feature_0
        pos_w_layer = self.model.PositionWeighted(sparse_segment)

        sparse_segment = schema.Map(
            keys=get_key(sparse_segment),
            values=pos_w_layer.position_weights,
            lengths_blob=sparse_segment.lengths
        )

        embedding_dim = 64
        embedding_after_pooling = self.model.SparseLookup(
            sparse_segment, [embedding_dim], 'PositionWeighted')
        self.model.output_schema = schema.Struct()
        self.assertEqual(
            schema.Scalar((np.float32, (embedding_dim, ))),
            embedding_after_pooling
        )

        train_init_net, train_net = self.get_training_nets()

        self.assertNetContainOps(
            train_init_net,
            [
                OpSpec("ConstantFill", None, None),  # position_weights/pos_w
                OpSpec("UniformFill", None, None),
                OpSpec("ConstantFill", None, None),
            ]
        )
        self.assertNetContainOps(train_net, [
            OpSpec("LengthsRangeFill", None, None),
            OpSpec("Gather", None, None),
            OpSpec("SparseLengthsWeightedSum", None, None),
        ])

        predict_net = self.get_predict_net()
        self.assertNetContainOps(predict_net, [
            OpSpec("LengthsRangeFill", None, None),
            OpSpec("Gather", None, None),
            OpSpec("SparseLengthsWeightedSum", None, None),
        ])

    def testSparseLookupPositionWeightedOnIdScoreList(self):
        record = schema.NewRecord(self.model.net, schema.Struct(
            ('sparse', schema.Struct(
                ('id_score_list_0', schema.Map(
                    schema.Scalar(
                        np.int64,
                        metadata=schema.Metadata(
                            categorical_limit=1000
                        ),
                    ),
                    np.float32
                )),
            )),
        ))

        embedding_dim = 64
        embedding_after_pooling = self.model.SparseLookup(
            record.sparse.id_score_list_0, [embedding_dim], 'PositionWeighted')
        self.model.output_schema = schema.Struct()
        self.assertEqual(
            schema.Scalar((np.float32, (embedding_dim, ))),
            embedding_after_pooling
        )

        train_init_net, train_net = self.get_training_nets()

        init_ops = self.assertNetContainOps(
            train_init_net,
            [
                OpSpec("UniformFill", None, None),
                OpSpec("ConstantFill", None, None),
            ]
        )
        sparse_lookup_op_spec = OpSpec(
            'SparseLengthsWeightedSum',
            [
                init_ops[0].output[0],
                record.sparse.id_score_list_0.values(),
                record.sparse.id_score_list_0.keys(),
                record.sparse.id_score_list_0.lengths(),
            ],
            [embedding_after_pooling()]
        )
        self.assertNetContainOps(train_net, [sparse_lookup_op_spec])

        predict_net = self.get_predict_net()
        self.assertNetContainOps(predict_net, [sparse_lookup_op_spec])

    def testPairwiseDotProductWithAllEmbeddings(self):
        embedding_dim = 64
        N = 5
        record = schema.NewRecord(self.model.net, schema.Struct(
            ('all_embeddings', schema.Scalar(
                ((np.float32, (N, embedding_dim)))
            )),
        ))
        current = self.model.PairwiseDotProduct(
            record, N * N)

        self.assertEqual(
            schema.Scalar((np.float32, (N * N, ))),
            current
        )

        train_init_net, train_net = self.get_training_nets()
        self.assertNetContainOps(train_init_net, [])
        self.assertNetContainOps(train_net, [
            OpSpec("BatchMatMul", None, None),
            OpSpec("Flatten", None, None),
        ])

    def testPairwiseDotProductWithXandYEmbeddings(self):
        embedding_dim = 64
        record = schema.NewRecord(self.model.net, schema.Struct(
            ('x_embeddings', schema.Scalar(
                ((np.float32, (5, embedding_dim)))
            )),
            ('y_embeddings', schema.Scalar(
                ((np.float32, (6, embedding_dim)))
            )),
        ))
        current = self.model.PairwiseDotProduct(
            record, 5 * 6)

        self.assertEqual(
            schema.Scalar((np.float32, (5 * 6, ))),
            current
        )

        train_init_net, train_net = self.get_training_nets()
        self.assertNetContainOps(train_init_net, [])
        self.assertNetContainOps(train_net, [
            OpSpec("BatchMatMul", None, None),
            OpSpec("Flatten", None, None),
        ])

    def testPairwiseDotProductWithXandYEmbeddingsAndGather(self):
        embedding_dim = 64

        output_idx = [1, 3, 5]
        output_idx_blob = self.model.add_global_constant(
            str(self.model.net.NextScopedBlob('pairwise_dot_product_gather')),
            output_idx,
            dtype=np.int32,
        )
        indices_to_gather = schema.Scalar(
            (np.int32, len(output_idx)),
            output_idx_blob,
        )

        record = schema.NewRecord(self.model.net, schema.Struct(
            ('x_embeddings', schema.Scalar(
                ((np.float32, (5, embedding_dim)))
            )),
            ('y_embeddings', schema.Scalar(
                ((np.float32, (6, embedding_dim)))
            )),
            ('indices_to_gather', indices_to_gather),
        ))
        current = self.model.PairwiseDotProduct(
            record, len(output_idx))

        # This assert is not necessary,
        # output size is passed into PairwiseDotProduct
        self.assertEqual(
            schema.Scalar((np.float32, (len(output_idx), ))),
            current
        )

        train_init_net, train_net = self.get_training_nets()
        self.assertNetContainOps(train_init_net, [])
        self.assertNetContainOps(train_net, [
            OpSpec("BatchMatMul", None, None),
            OpSpec("Flatten", None, None),
            OpSpec("BatchGather", None, None),
        ])

    def testPairwiseDotProductIncorrectInput(self):
        embedding_dim = 64
        record = schema.NewRecord(self.model.net, schema.Struct(
            ('x_embeddings', schema.Scalar(
                ((np.float32, (5, embedding_dim)))
            )),
        ))
        with self.assertRaises(AssertionError):
            self.model.PairwiseDotProduct(
                record, 25)

        record = schema.NewRecord(self.model.net, schema.Struct(
            ('all_embeddings', schema.List(np.float32))
        ))
        with self.assertRaises(AssertionError):
            self.model.PairwiseDotProduct(
                record, 25)

    def testConcat(self):
        embedding_dim = 64
        input_record = self.new_record(schema.Struct(
            ('input1', schema.Scalar((np.float32, (embedding_dim, )))),
            ('input2', schema.Scalar((np.float32, (embedding_dim, )))),
            ('input3', schema.Scalar((np.float32, (embedding_dim, )))),
        ))

        output = self.model.Concat(input_record)
        self.assertEqual(
            schema.Scalar((np.float32, ((len(input_record.fields) * embedding_dim, )))),
            output
        )

        # Note that in Concat layer we assume first dimension is batch.
        # so input is B * embedding_dim
        # add_axis=1 make it B * 1 * embedding_dim
        # concat on axis=1 make it B * N * embedding_dim
        output = self.model.Concat(input_record, axis=1, add_axis=1)
        self.assertEqual(
            schema.Scalar((np.float32, ((len(input_record.fields), embedding_dim)))),
            output
        )

    def testSamplingTrain(self):
        output_dims = 1000

        indices = self.new_record(schema.Scalar((np.int32, (10,))))
        sampling_prob = self.new_record(schema.Scalar((np.float32, (10, ))))

        sampled_fc = self.model.SamplingTrain(
            schema.Struct(
                ('input', self.model.input_feature_schema.float_features),
                ('indices', indices),
                ('sampling_prob', sampling_prob),
            ),
            "FC",
            output_dims,
        )
        self.model.output_schema = sampled_fc

        # Check that we don't add prediction layer into the model
        self.assertEqual(1, len(self.model.layers))

        self.assertEqual(
            schema.Scalar((np.float32, (output_dims, ))),
            sampled_fc
        )

        train_init_net, train_net = self.get_training_nets()

        init_ops = self.assertNetContainOps(
            train_init_net,
            [
                OpSpec("UniformFill", None, None),
                OpSpec("UniformFill", None, None),
            ]
        )

        sampled_fc_layer = self.model.layers[0]

        gather_w_spec = OpSpec(
            "Gather",
            [
                init_ops[0].output[0],
                indices(),
            ],
            [
                sampled_fc_layer._prediction_layer.train_param_blobs[0]
            ]
        )
        gather_b_spec = OpSpec(
            "Gather",
            [
                init_ops[1].output[0],
                indices(),
            ],
            [
                sampled_fc_layer._prediction_layer.train_param_blobs[1]
            ]
        )
        train_fc_spec = OpSpec(
            "FC",
            [
                self.model.input_feature_schema.float_features(),
            ] + sampled_fc_layer._prediction_layer.train_param_blobs,
            sampled_fc.field_blobs()
        )
        log_spec = OpSpec("Log", [sampling_prob()], [None])
        sub_spec = OpSpec(
            "Sub",
            [sampled_fc.field_blobs()[0], None],
            sampled_fc.field_blobs()
        )

        train_ops = self.assertNetContainOps(
            train_net,
            [gather_w_spec, gather_b_spec, train_fc_spec, log_spec, sub_spec])

        self.assertEqual(train_ops[3].output[0], train_ops[4].input[1])

        predict_net = self.get_predict_net()
        self.assertNetContainOps(
            predict_net,
            [
                OpSpec(
                    "FC",
                    [
                        self.model.input_feature_schema.float_features(),
                        init_ops[0].output[0],
                        init_ops[1].output[0],
                    ],
                    sampled_fc.field_blobs()
                )
            ]
        )

    def testBatchLRLoss(self):
        input_record = self.new_record(schema.Struct(
            ('label', schema.Scalar((np.float64, (1,)))),
            ('logit', schema.Scalar((np.float32, (2,)))),
            ('weight', schema.Scalar((np.float64, (1,))))
        ))
        loss = self.model.BatchLRLoss(input_record)
        self.assertEqual(schema.Scalar((np.float32, tuple())), loss)

    def testMarginRankLoss(self):
        input_record = self.new_record(schema.Struct(
            ('pos_prediction', schema.Scalar((np.float32, (1,)))),
            ('neg_prediction', schema.List(np.float32)),
        ))
        pos_items = np.array([0.1, 0.2, 0.3], dtype=np.float32)
        neg_lengths = np.array([1, 2, 3], dtype=np.int32)
        neg_items = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6], dtype=np.float32)
        schema.FeedRecord(
            input_record,
            [pos_items, neg_lengths, neg_items]
        )
        loss = self.model.MarginRankLoss(input_record)
        self.run_train_net_forward_only()
        self.assertEqual(schema.Scalar((np.float32, tuple())), loss)

    def testBatchMSELoss(self):
        input_record = self.new_record(schema.Struct(
            ('label', schema.Scalar((np.float64, (1,)))),
            ('prediction', schema.Scalar((np.float32, (2,)))),
        ))
        loss = self.model.BatchMSELoss(input_record)
        self.assertEqual(schema.Scalar((np.float32, tuple())), loss)

    def testBatchSigmoidCrossEntropyLoss(self):
        input_record = self.new_record(schema.Struct(
            ('label', schema.Scalar((np.float32, (32,)))),
            ('prediction', schema.Scalar((np.float32, (32,))))
        ))
        loss = self.model.BatchSigmoidCrossEntropyLoss(input_record)
        self.assertEqual(schema.Scalar((np.float32, tuple())), loss)

    def testBatchSoftmaxLoss(self):
        input_record = self.new_record(schema.Struct(
            ('label', schema.Scalar((np.float32, tuple()))),
            ('prediction', schema.Scalar((np.float32, (32,))))
        ))
        loss = self.model.BatchSoftmaxLoss(input_record)
        self.assertEqual(schema.Struct(
            ('softmax', schema.Scalar((np.float32, (32,)))),
            ('loss', schema.Scalar(np.float32)),
        ), loss)

    def testBatchSoftmaxLossWeight(self):
        input_record = self.new_record(schema.Struct(
            ('label', schema.Scalar((np.float32, tuple()))),
            ('prediction', schema.Scalar((np.float32, (32,)))),
            ('weight', schema.Scalar((np.float64, (1,))))
        ))
        loss = self.model.BatchSoftmaxLoss(input_record)
        self.assertEqual(schema.Struct(
            ('softmax', schema.Scalar((np.float32, (32,)))),
            ('loss', schema.Scalar(np.float32)),
        ), loss)

    @given(
        X=hu.arrays(dims=[2, 5]),
    )
    def testBatchNormalization(self, X):
        input_record = self.new_record(schema.Scalar((np.float32, (5,))))
        schema.FeedRecord(input_record, [X])
        bn_output = self.model.BatchNormalization(input_record)
        self.assertEqual(schema.Scalar((np.float32, (5,))), bn_output)
        self.model.output_schema = schema.Struct()

        train_init_net, train_net = self.get_training_nets()

        init_ops = self.assertNetContainOps(
            train_init_net,
            [
                OpSpec("ConstantFill", None, None),
                OpSpec("ConstantFill", None, None),
                OpSpec("ConstantFill", None, None),
                OpSpec("ConstantFill", None, None),
            ]
        )

        input_blob = input_record.field_blobs()[0]
        output_blob = bn_output.field_blobs()[0]

        expand_dims_spec = OpSpec(
            "ExpandDims",
            [input_blob],
            None,
        )

        train_bn_spec = OpSpec(
            "SpatialBN",
            [None, init_ops[0].output[0], init_ops[1].output[0],
                init_ops[2].output[0], init_ops[3].output[0]],
            [output_blob, init_ops[2].output[0], init_ops[3].output[0], None, None],
            {'is_test': 0, 'order': 'NCHW', 'momentum': 0.9},
        )

        test_bn_spec = OpSpec(
            "SpatialBN",
            [None, init_ops[0].output[0], init_ops[1].output[0],
                init_ops[2].output[0], init_ops[3].output[0]],
            [output_blob],
            {'is_test': 1, 'order': 'NCHW', 'momentum': 0.9},
        )

        squeeze_spec = OpSpec(
            "Squeeze",
            [output_blob],
            [output_blob],
        )

        self.assertNetContainOps(
            train_net,
            [expand_dims_spec, train_bn_spec, squeeze_spec]
        )

        eval_net = self.get_eval_net()

        self.assertNetContainOps(
            eval_net,
            [expand_dims_spec, test_bn_spec, squeeze_spec]
        )

        predict_net = self.get_predict_net()

        self.assertNetContainOps(
            predict_net,
            [expand_dims_spec, test_bn_spec, squeeze_spec]
        )

        workspace.RunNetOnce(train_init_net)
        workspace.RunNetOnce(train_net)

        schema.FeedRecord(input_record, [X])
        workspace.RunNetOnce(eval_net)

        schema.FeedRecord(input_record, [X])
        workspace.RunNetOnce(predict_net)

    @given(
        X=hu.arrays(dims=[5, 2]),
        num_to_collect=st.integers(min_value=1, max_value=10),
    )
    def testLastNWindowCollector(self, X, num_to_collect):
        input_record = self.new_record(schema.Scalar(np.float32))
        schema.FeedRecord(input_record, [X])
        last_n = self.model.LastNWindowCollector(input_record, num_to_collect)
        self.run_train_net_forward_only()
        output_record = schema.FetchRecord(last_n.last_n)
        start = max(0, 5 - num_to_collect)
        npt.assert_array_equal(X[start:], output_record())
        num_visited = schema.FetchRecord(last_n.num_visited)
        npt.assert_array_equal([5], num_visited())

    def testUniformSampling(self):
        input_record = self.new_record(schema.Scalar(np.int32))
        input_array = np.array([3, 10, 11, 15, 20, 99], dtype=np.int32)
        schema.FeedRecord(input_record, [input_array])
        num_samples = 20
        num_elements = 100
        uniform_sampling_output = self.model.UniformSampling(
            input_record, num_samples, num_elements)
        self.model.loss = uniform_sampling_output
        self.run_train_net()
        samples = workspace.FetchBlob(uniform_sampling_output.samples())
        sampling_prob = workspace.FetchBlob(
            uniform_sampling_output.sampling_prob())
        self.assertEqual(num_samples, len(samples))
        np.testing.assert_array_equal(input_array, samples[:len(input_array)])
        np.testing.assert_almost_equal(
            np.array([float(num_samples) / num_elements] * num_samples,
                     dtype=np.float32),
            sampling_prob
        )

    def testUniformSamplingWithIncorrectSampleSize(self):
        input_record = self.new_record(schema.Scalar(np.int32))
        num_samples = 200
        num_elements = 100
        with self.assertRaises(AssertionError):
            self.model.UniformSampling(input_record, num_samples, num_elements)

    def testGatherRecord(self):
        indices = np.array([1, 3, 4], dtype=np.int32)
        dense = np.array(list(range(20)), dtype=np.float32).reshape(10, 2)
        lengths = np.array(list(range(10)), dtype=np.int32)
        items = np.array(list(range(lengths.sum())), dtype=np.int64)
        items_lengths = np.array(list(range(lengths.sum())), dtype=np.int32)
        items_items = np.array(list(range(items_lengths.sum())), dtype=np.int64)
        record = self.new_record(schema.Struct(
            ('dense', schema.Scalar(np.float32)),
            ('sparse', schema.Struct(
                ('list', schema.List(np.int64)),
                ('list_of_list', schema.List(schema.List(np.int64))),
            )),
            ('empty_struct', schema.Struct())
        ))
        indices_record = self.new_record(schema.Scalar(np.int32))
        input_record = schema.Struct(
            ('indices', indices_record),
            ('record', record),
        )
        schema.FeedRecord(
            input_record,
            [indices, dense, lengths, items, lengths, items_lengths,
             items_items])
        gathered_record = self.model.GatherRecord(input_record)
        self.assertTrue(schema.equal_schemas(gathered_record, record))

        self.run_train_net_forward_only()
        gathered_dense = workspace.FetchBlob(gathered_record.dense())
        np.testing.assert_array_equal(