"""Provides an adapted version of the stack gunpowder node
(Fixes random seed bug)
"""
import time

from gunpowder.nodes.batch_filter import BatchFilter
from gunpowder.array import Array
from gunpowder.batch import Batch
from gunpowder.profiling import Timing
import numpy as np


class Stack(BatchFilter):
    '''Request several batches and stack them together, introducing a new
    dimension for each array. This is useful to create batches with several
    samples and only makes sense if there is a source of randomness upstream.

    This node stacks only arrays, not points. The resulting batch will have the
    same point sets as found in the first batch requested upstream.

    Args:

        num_repetitions (``int``):

            How many upstream batches to stack.
    '''

    def __init__(self, num_repetitions):

        self.num_repetitions = num_repetitions

    def provide(self, request):

        batches = []
        for _ in range(self.num_repetitions):
            request._random_seed = int(time.time() * 1e6)
            batches.append(self.get_upstream_provider().request_batch(request))

        timing = Timing(self)
        timing.start()

        batch = Batch()
        for b in batches:
            batch.profiling_stats.merge_with(b.profiling_stats)

        for key, spec in request.array_specs.items():

            data = np.stack([b[key].data for b in batches])
            batch[key] = Array(
                data,
                batches[0][key].spec.copy())

        # copy points of first batch requested
        for key, spec in request.points_specs.items():
            batch[key] = batches[0][key]

        timing.stop()
        batch.profiling_stats.add(timing)

        return batch
