# bunch of imports 
import numpy as np
import matplotlib.pyplot as plt
import h5py
import sklearn
import sklearn.datasets
import sklearn.linear_model
import scipy.io

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.scatter(train_X[0, :], train_X[1, :], c=np.squeeze(train_Y), s=40, cmap=plt.cm.Spectral);
    
    return train_X, train_Y, test_X, test_Y
    
# def load_planar_dataset(randomness, seed):
    
#     np.random.seed(seed)
    
#     m = 50
#     N = int(m/2) # number of points per class
#     D = 2 # dimensionality
#     X = np.zeros((m,D)) # data matrix where each row is a single example
#     Y = np.zeros((m,1), dtype='uint8') # labels vector (0 for red, 1 for blue)
#     a = 2 # maximum ray of the flower

#     for j in range(2):
        
#         ix = range(N*j,N*(j+1))
#         if j == 0:
#             t = np.linspace(j, 4*3.1415*(j+1),N) #+ np.random.randn(N)*randomness # theta
#             r = 0.3*np.square(t) + np.random.randn(N)*randomness # radius
#         if j == 1:
#             t = np.linspace(j, 2*3.1415*(j+1),N) #+ np.random.randn(N)*randomness # theta
#             r = 0.2*np.square(t) + np.random.randn(N)*randomness # radius
            
#         X[ix] = np.c_[r*np.cos(t), r*np.sin(t)]
#         Y[ix] = j
        
#     X = X.T
#     Y = Y.T

#     return X, Y
    
def load_planar_dataset(seed):
    
    np.random.seed(seed)
    
    m = 1000 # number of examples
    N = int(m/2) # number of points per class
    D = 2 # dimensionality
    X = np.zeros((m,D)) # data matrix where each row is a single example
    Y = np.zeros((m,1), dtype='uint8') # labels vector (0 for red, 1 for blue)
    a = 4 # maximum ray of the flower

    for j in range(2):
        ix = range(N*j,N*(j+1))
        t = np.linspace(j*3.12,(j+1)*3.12,N) + np.random.randn(N)*0.2 # theta
        r = a*np.sin(4*t) + np.random.randn(N)*0.2 # radius
        X[ix] = np.c_[r*np.sin(t), r*np.cos(t)]
        Y[ix] = j
        
    X = X.T
    Y = Y.T

    return X, Y


def plot_decision_boundary(model, X, y):
    # Set min and max values and give it some padding
    x_min, x_max = X[0, :].min() - 1, X[0, :].max() + 1
    y_min, y_max = X[1, :].min() - 1, X[1, :].max() + 1
    h = 0.01
    # Generate a grid of points with distance h between them
    xx, yy = np.meshgrid(np.arange(x_min, x_max, h), np.arange(y_min, y_max, h))
    # Predict the function value for the whole grid
    Z = model(np.c_[xx.ravel(), yy.ravel()])
    Z = Z.reshape(xx.shape)
    # Plot the contour and training examples
    plt.contourf(xx, yy, Z, cmap=plt.cm.Spectral)
    plt.ylabel('x2')
    plt.xlabel('x1')
    plt.scatter(X[0, :], X[1, :], cmap=plt.cm.Spectral, c=np.squeeze(y)) 
    plt.show()