import numpy as np
import os
import pprint

from data_provider import DataProvider

DATA_DIR = "/data/ziz/not-backed-up/datasets-ziz-only"
import pprint

def info():
    meta_dict = {
        "name": "mnist",
        "location": os.path.join(DATA_DIR, "raw_data/mnist", "mnist.npz"),
        "available set": ['train', 'test'],
        "available mode": ['all', 'batch'],
        "data source": "",
        "remark": "",
    }
    pp = pprint.PrettyPrinter(indent=4)
    pp.pprint(meta_dict)

def load_data(which_set="train", mode="all", batch_size=100, shuffle=True):
    assert which_set in ['train', 'valid', 'test'], \
            "which_set takes values in [train, test]"
    assert mode in ['all', 'batch'], \
            "mode takes values in [all, batch]"
    path = os.path.join(DATA_DIR, "raw_data/mnist", "mnist.npz")
    if mode == 'all':
        return _load_data_all(path, which_set)
    elif mode == 'batch':
        return _load_data_batch(path, which_set, batch_size, shuffle)

def _load_data_all(path, which_set):
    data = np.load(path)
    return data["x_"+which_set], data['y_'+which_set]

def _load_data_batch(path, which_set, batch_size, shuffle):
    generator = MNISTDataProvider(which_set=which_set, batch_size=batch_size, shuffle_order=shuffle)
    return generator

class MNISTDataProvider(DataProvider):

    def __init__(self, which_set="train", batch_size=100, max_num_batches=-1, shuffle_order=True):

        inputs, targets = load_data(which_set, "all")
        super().__init__(
            inputs, targets, batch_size, max_num_batches, shuffle_order, rng=None)

    def next(self):
        """Returns next data batch or raises `StopIteration` if at end."""
        inputs_batch, targets_batch = super().next()
        return inputs_batch, targets_batch

import numpy as np
import os
import pprint

from data_provider import DataProvider

DATA_DIR = "/data/ziz/not-backed-up/datasets-ziz-only"
import pprint

def info():
    meta_dict = {
        "name": "mnist",
        "location": os.path.join(DATA_DIR, "raw_data/mnist", "mnist.npz"),
        "available set": ['train', 'test'],
        "available mode": ['all', 'batch'],
        "data source": "",
        "remark": "",
    }
    pp = pprint.PrettyPrinter(indent=4)
    pp.pprint(meta_dict)

def load_data(which_set="train", mode="all", batch_size=100, shuffle=True):
    assert which_set in ['train', 'valid', 'test'], \
            "which_set takes values in [train, test]"
    assert mode in ['all', 'batch'], \
            "mode takes values in [all, batch]"
    path = os.path.join(DATA_DIR, "raw_data/mnist", "mnist.npz")
    if mode == 'all':
        return _load_data_all(path, which_set)
    elif mode == 'batch':
        return _load_data_batch(path, which_set, batch_size, shuffle)

def _load_data_all(path, which_set):
    data = np.load(path)
    return data["x_"+which_set], data['y_'+which_set]

def _load_data_batch(path, which_set, batch_size, shuffle):
    generator = MNISTDataProvider(which_set=which_set, batch_size=batch_size, shuffle_order=shuffle)
    return generator

class MNISTDataProvider(DataProvider):

    def __init__(self, which_set="train", batch_size=100, max_num_batches=-1, shuffle_order=True):

        inputs, targets = load_data(which_set, "all")
        super().__init__(
            inputs, targets, batch_size, max_num_batches, shuffle_order, rng=None)

    def next(self):
        """Returns next data batch or raises `StopIteration` if at end."""
        inputs_batch, targets_batch = super().next()
        return inputs_batch, targets_batch

Returns next data batch or raises `StopIteration` if at end.