
import numpy as np 
import cv2
from glob import glob
from tensorflow.keras.datasets import mnist
def load_mnist(SIZE=128,num_train=2000, num_test=500):
	'''
	SIZE: points image size
	'''

	(x_train, _), (x_test, _) = mnist.load_data()
	x_train = x_train.astype('float32') / 255.
	x_test = x_test.astype('float32') / 255.
	x_train = np.reshape(x_train, (len(x_train), 28, 28, 1))  
	x_test = np.reshape(x_test, (len(x_test), 28, 28, 1)) 

	X_train, X_test = x_train[0][...,-1], x_test[0][...,-1]

	X_train, X_test = cv2.resize(X_train, (SIZE,SIZE)), cv2.resize(X_test, (SIZE,SIZE))
	X_train, X_test = X_train.reshape(1,SIZE,SIZE,1), X_test.reshape(1,SIZE,SIZE,1)
	for i in x_train[:num_train]:
		x_i = cv2.resize(i[...,-1], (SIZE, SIZE))
		X_train = np.concatenate((X_train, x_i.reshape(1,SIZE,SIZE,1)), axis=0)
		
	for j in x_test[-num_train:]:
		x_j = cv2.resize(j[...,-1], (SIZE, SIZE))
		X_test = np.concatenate((X_test, x_j.reshape(1,SIZE,SIZE,1)), axis=0)
	return X_train, X_test


def load_path(img_path = '/home/nikki_intern/GAN/DATA/img_align_celeba/*', SIZE = 256, 
              num_train=2000, num_test=1000):

    img_name = glob(img_path)
    train_name = img_name[:2000]
    test_name = img_name[-500:]

    X_train = cv2.resize(cv2.imread(train_name[0])[:,:,::-1], (SIZE,SIZE)) 
    X_test = cv2.resize(cv2.imread(test_name[0])[:,:,::-1], (SIZE,SIZE))
    H,W,C = X_train.shape
    X_train, X_test = X_train.reshape(1,SIZE,SIZE,C), X_test.reshape(1,SIZE,SIZE,C)
    for i in train_name[1:]:
        x_i = cv2.resize(cv2.imread(i)[:,:,::-1], (SIZE,SIZE)) 
        X_train = np.concatenate((X_train, x_i.reshape(1,SIZE,SIZE,C)), axis=0)
        
    for j in test_name[1:]:
        x_j = cv2.resize(cv2.imread(j)[:,:,::-1], (SIZE,SIZE)) 
        X_test = np.concatenate((X_test, x_j.reshape(1,SIZE,SIZE,C)), axis=0)
        
    return X_train.astype('float32') / 255., X_test.astype('float32') / 255.