import os
import numpy as np
import pickle


def load_CIFAR_batch(filename):
    """ load single batch of cifar """
    with open(filename, 'rb') as f:
        datadict = pickle.load(f, encoding='latin1')
        X = datadict['data']
        Y = datadict['labels']
        X = X.reshape(10000, 3, 32, 32).transpose(0, 2, 3, 1).astype("float")
        Y = np.array(Y)

    return X, Y


def load_CIFAR10(ROOT):
    """ load all of cifar """
    xs = []
    ys = []
    for b in range(1, 6):
        f = os.path.join(ROOT, 'data_batch_%d' % (b,))
        X, Y = load_CIFAR_batch(f)
        xs.append(X)
        ys.append(Y)
        Xtr = np.concatenate(xs)  # 使变成行向量
        Ytr = np.concatenate(ys)
        del X, Y
        Xte, Yte = load_CIFAR_batch(os.path.join(ROOT, 'test_batch'))

    return Xtr, Ytr, Xte, Yte


def next_batch(img, label, batch_size, step):
    img_batch = img[step * batch_size:step * batch_size + batch_size]
    lab_batch = label[step * batch_size:step * batch_size + batch_size]

    return img_batch, lab_batch


def shuffle_data(imgData, labData):
    index = np.random.permutation(len(imgData))
    shuffled_image = imgData[index]
    shuffled_label = labData[index]

    return shuffled_image, shuffled_label
