from torchvision import datasets, transforms
import pandas as pd
import numpy as np

COLUMN_LABEL = '_label'
SEED = 998822
#---
np.random.seed(SEED)

#---
def get_data(train):
	data_raw = datasets.MNIST('../data/dl/', train=train, download=True,  transform=transforms.Compose([
							transforms.ToTensor(),
							lambda x: x.numpy().flatten()]), target_transform=lambda y: y.numpy())

	data_x, data_y = zip(*data_raw)
	data_x = np.array(data_x)
	data_y = np.array(data_y, dtype='int32').reshape(-1, 1)

	data = pd.DataFrame(data_x)
	data[COLUMN_LABEL] = data_y

	return data, data_x.mean(), data_x.std()

#---
data_train, avg, std = get_data(train=True)
data_test, _, _  = get_data(train=False)

#shuffle
val_idx = np.random.choice(data_train.shape[0], 10000, replace=False).tolist()

data_val   = data_train.iloc[val_idx]
data_train = data_train.drop(val_idx)

print(data_train.head())

print("Number of features:", data_train.shape[1] - 1)
print("Classes:", data_train.iloc[:, -1].unique())

print()
print("Train len:", data_train.shape[0])
print("Val len:  ", data_val.shape[0])
print("Test len: ", data_test.shape[0])

data_train.to_pickle("../data/mnist-train")
data_val.to_pickle("../data/mnist-val")
data_test.to_pickle("../data/mnist-test")

#--- prepare meta
idx = data_train.columns[:-1]
meta = pd.DataFrame(index=idx, dtype='float32')

meta['avg'] = avg	#data_train.mean()
meta['std'] = std	#data_train.std()
meta['cost'] = 1.

meta.loc[ meta['std'] == 0., 'std' ] = 1.0
meta = meta.astype('float32')

print()
print(meta)

meta.to_pickle("../data/mnist-meta")
