import numpy as num

#I couldn't resist implementing these two functions
def gcd_bin(a,b):
    while b != 0:
        a, b = b, a%b
    return a

def gcd(*args):
    if len(args) == 1:
        return args[0]
    if len(args) == 2:
        return gcd_bin(*args)
    return gcd(gcd_bin(args[0], args[1]), *args[2:])

class HFDataSource(object):
    def __init__(self, data, mbsz, gradBatchSize, curvatureBatchSize, lineSearchBatchSize, gradBatchIsTrainingSet):
        """
        mbsz is the computational minibatch size (as opposed to the
        conceptual minibatch sizes used in the optimizer). mbsz is the
        actual number of cases processed in a single matrix operation
        on the gpu at one time. The larger the mbsz, the more
        parallelism the GPU can exploit and the more GPU memory will
        be needed.
        
        We always assume cgBacktrackBatchSize == curvatureBatchSize.
        If gradBatchIsTrainingSet then we assume lineSearchBatchSize equals gradBatchSize.
        
        data is either an infinite iterator over the smallest
        minibatch or the dataset itself (as in a finite list of
        batches of size mbsz).
        """
        assert(gcd(gradBatchSize, lineSearchBatchSize, curvatureBatchSize) == curvatureBatchSize)
        assert(gcd(gradBatchSize, lineSearchBatchSize, curvatureBatchSize, mbsz) == mbsz)
        
        self.gradBatchSize = gradBatchSize
        self.curvatureBatchSize = curvatureBatchSize
        self.lineSearchBatchSize = lineSearchBatchSize
        self.mbsz = mbsz
        self.data = data
        self.gradBatchIsTrainingSet = gradBatchIsTrainingSet
        if self.gradBatchIsTrainingSet:
            assert(lineSearchBatchSize == gradBatchSize)
            assert(mbsz == curvatureBatchSize)
            self.batchList = [b for b in data]
            self.gradBatch = self.batchList
            self.batchOffset = 0
        
    def obtainBatch(self, batchSize):
        assert(batchSize % self.mbsz == 0)
        reqs = batchSize/self.mbsz
        if self.gradBatchIsTrainingSet:
            if batchSize == self.gradBatchSize:
                return self.gradBatch
            #we start at 1 because it doesn't matter and because it helps match Ilya's code during debugging
            self.batchOffset += 1
            batch = [self.batchList[(self.batchOffset+i) % len(self.batchList)] for i in range(reqs)]
            self.batchOffset += reqs-1
            return batch
        return [self.data.next() for i in range(reqs)]
