import random

class Dataloader():
    """
    Community dataset from "GraphRNN: Generating Realistic Graphs with Deep Auto-regressive Models"
    Arxiv link: https://arxiv.org/pdf/1802.08773.pdf

     500 two-community graphs with 60 ≤ |V | ≤ 160. Each community is generated by the Erdos-R ˝ enyi ´model (E-R)
     (Erdos & R ˝ enyi ´ , 1959) with n = |V |/2 nodes and p = 0.3. We then add 0.05|V | inter-community edgeswith
     uniform probability.


    The paper description and paper implementation of this dataset differs
    """
    def __init__(self, dataset, batch_size=1, shuffle=True, n_nodes=0):
        self.dataset = dataset
        self.batch_size = batch_size
        self.idx = 0
        self.shuffle = shuffle
        if self.shuffle:
            random.shuffle(self.dataset.graphs)

    def __iter__(self):
        return self

    def __next__(self):
        if self.idx > len(self.dataset) - self.batch_size:
            self.idx = 0
            random.shuffle(self.dataset.graphs)
            raise StopIteration  # Done iterating.
        else:
            samples = self.dataset[self.idx:self.idx + self.batch_size]
            self.idx += self.batch_size
            return samples

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


if __name__ == "__main__":
    from ae_datasets import d_selector

    dataset = d_selector.retrieve_dataset("erdosrenyi_100_16_20")
    dataloader = Dataloader(dataset, batch_size=1)
    for sample in dataloader:
        print(sample)
