# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

"""DataIter for negative sampling.
"""
import mxnet as mx
import numpy as np

class NegativeSamplingDataIter(mx.io.DataIter):
    """Wraps an existing DataIter to produce a new DataIter with negative samples.
    Assumes that all the relevant inputs are in data, not labels.
    Drops (replaces) any labels in the original DataIter.

    It only shuffles one of the input data columns, specified in the
    constructor as shuffle_data_idx.  So if the original input data
    has three columns, ('item_ids', 'item_words', 'users') and you want
    to keep the two "item_*" together, then set `shuffle_data_idx=2`
    and `users` will be shuffled for the negative samples.

    Output batches will be larger than input batches by a factor
    of (1+sample_ratio)

    Negative samples are always drawn from the same minibatch.
    So they're effectively sampled according to the frequency at
    which they appear in the training data.  (Other reasonable
    samling strategies are not implemented here.)
    The shuffling is checked to ensure that a true positive sample
    isn't returned as a negative sample.
    """

    def __init__(self, source_dataiter, sample_ratio=1, shuffle_data_idx=1,
                 positive_label=1, negative_label=0):
        self._sourcedata = source_dataiter
        source_dataiter.reset()
        self.positive_label = positive_label  # output shapes = input shapes
        self.negative_label = negative_label
        self.shuffle_data_idx = shuffle_data_idx
        if sample_ratio == int(sample_ratio):
            self.sample_ratio = int(sample_ratio)
        else:
            raise ValueError("sample_ratio must be an integer, not %s" % sample_ratio)
        self._clear_queue()

        self.provide_data = source_dataiter.provide_data
        self.provide_label = source_dataiter.provide_label
        self.batch_size = source_dataiter.batch_size

    def _clear_queue(self):
        self._sampled_queue = []

    def _push_queue(self, data_list, labels):
        """Takes a list of numpy arrays for data,
        and a numpy array for labels.
        Converts to minibatches and puts it on the queue.
        """
        num_minibatches = 1+self.sample_ratio
        total_size = len(labels)
        slice_size = total_size / num_minibatches
        def slicer(x, s):
            idx = range(s*slice_size, (s+1)*slice_size)
            return np.take(x,idx,0)

        for i in range(1+self.sample_ratio):
            nddata = [mx.nd.array(slicer(x,i)) for x in data_list]
            ndlabels = mx.nd.array(slicer(labels,i))
            batch = mx.io.DataBatch(nddata, [ndlabels], provide_data=self.provide_data,
                                    provide_label=self.provide_label)
            self._sampled_queue.append(batch)

    def next(self):
        if not self._sampled_queue:
            self._refill_queue()
        batch = self._sampled_queue.pop()
        return batch

    def reset(self):
        self._sourcedata.reset()
        self._clear_queue()

    def _shuffle_batch(self, data):
        # Takes a list of NDArrays.  Returns a shuffled version as numpy
        a = data[self.shuffle_data_idx].asnumpy()

        # Come up with a shuffled index
        batch_size = data[0].shape[0]
        si = np.arange(batch_size)
        np.random.shuffle(si)
        matches = (si == np.arange(batch_size)) # everywhere it didn't shuffle
        si -= matches  # Shifts down by 1 when true, ensuring it differs
        # Note shifting down by 1 works in python because -1 is a valid index.
        #Q: Is this shifting introducing bias?

        # Shuffle the data with the shuffle index
        shuf_a = np.take(a,si,0)  # like a[si,:] but general for ndarray's

        # Return similar datastructure to what we got.  Convert all to numpy
        out = [d.asnumpy() for d in data]
        out[self.shuffle_data_idx] = shuf_a
        return out

    def _refill_queue(self):
        """Fetch another batch from the source, and shuffle it to make
        negative samples.
        """
        original = self._sourcedata.next().data  # List of NDArrays: one per input
        batch_size = original[0].shape[0]
        num_inputs = len(original)

        #Start with positive examples, copied straight
        outdata = [[o.asnumpy()] for o in original] # list of lists of numpy arrays
        outlabels = [np.ones(batch_size) * self.positive_label] # list of numpy arrays
        # The inner list of both is the set of samples.  We'll recombine later.

        # Construct negative samples.
        for _ in range(self.sample_ratio):
            shuffled = self._shuffle_batch(original)
            for i,d in enumerate(shuffled):
                outdata[i].append(d)
            outlabels.append(np.ones(batch_size) * self.negative_label)
        def stacker(x):
            if len(x[0].shape)==1:
                return np.hstack(x)
            else:
                return np.vstack(x)
        outdata = [stacker(x) for x in outdata] # Single tall vectors
        outlabels = stacker(outlabels)

        # Record-level shuffle so the negatives are mixed in.
        def shuffler(x, idx):
            return np.take(x,idx,0)
        shuf_idx = np.arange(len(outlabels))
        np.random.shuffle(shuf_idx)
        outdata = [shuffler(o,shuf_idx) for o in outdata]
        outlabels = shuffler(outlabels,shuf_idx)
        self._push_queue(outdata,outlabels)


if __name__ == "__main__":
    print("Simple test of NegativeSamplingDataIter")
    np.random.seed(123)
    A = np.random.randint(-20,150,size=(100,5))
    B = np.random.randint(-2,15,size=(100,2))
    R = np.random.randint(1,5,size=(100,))
    batch_size=3
    oridi = mx.io.NDArrayIter(data={'a':A,'b':B},label=R, batch_size=batch_size)
    oribatch = oridi.next()
    oridi.reset()
    for ratio in range(0,5):
        nsdi = NegativeSamplingDataIter(oridi, sample_ratio=ratio)

        # Check sizes of output
        bat = nsdi.next()
        for i in range(len(bat.data)):
            assert bat.data[i].shape[0] == batch_size
            assert bat.data[i].shape[1] == oribatch.data[i].shape[1]
        assert bat.label.shape[0] == batch_size

        # Check that we get more minibatches
        oridi.reset()
        ori_cnt = len(list(oridi))
        nsdi.reset()
        ns_cnt = len(list(nsdi))
        assert ns_cnt == ori_cnt * (1+ratio)

    print("Tests done")

