import numpy as np

class base_sampler(object):
    def __init__(self, data_source):
        pass

    def __iter__(self):
        raise NotImplementedError

    def __len__(self):
        raise NotImplementedError

class sequential_sampler(base_sampler):
    """return inidices in orginal order"""
    def __init__(self, data_source):
        self.data_source = data_source
    
    def __iter__(self):
        return iter(range(len(self.data_source)))

    def __len__(self):
        return len(self.data_source)

class random_sampler(base_sampler):
    """shuffles indicies for sampling"""
    def __init__(self, data_source):
        self.data_source = data_source
    def __iter__(self):
        return iter(np.random.permutation(len(self.data_source)))
    def __len__(self):
        return len(self.data_source)
    
class minibatch_sampler(base_sampler):
    def __init__(self, sampler, batch_size, drop_last):
        self.sampler = sampler
        self.batch_size = batch_size
        self.drop_last = drop_last

    def __iter__(self):
        batch = []
        for idx in self.sampler:
            batch.append(idx)
            if len(batch) == self.batch_size:
                yield batch
                batch = []
        if len(batch) > 0 and not self.drop_last:
            yield batch
    
    def __len__(self):
        if self.drop_last:
            return len(self.sampler) // self.batch_size
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size
        

