import scipy.io
import matplotlib.pyplot as plt
import numpy as np

def load_2D_dataset():
    data = scipy.io.loadmat('datasets/data.mat')
    train_X = data['X'].T
    train_Y = data['y'].T
    test_X = data['Xval'].T
    test_Y = data['yval'].T

    plt.figure(1)
    z = np.reshape(data['y'],data['y'].shape[0])
    plt.scatter(train_X[0, :], train_X[1, :], c=z, s=40, cmap=plt.cm.Spectral);
    
    return train_X, train_Y, test_X, test_Y