import numpy as np
from sklearn import linear_model
from sklearn.cross_validation import KFold
import matplotlib.pyplot as plt


n = 20
sims = 1
p = 15
beta0 = 3
beta = np.array([2, 2, 2, 2, 2, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0])

Sigma = 0.2*np.eye(p, dtype=float) + 0.8*np.ones((p, p))

sd_e = 5

	
# Calculate X with n = 20
X_train = np.random.multivariate_normal(np.zeros(p), Sigma, n)
X_test = np.random.multivariate_normal(np.zeros(p), Sigma, 1000)


# Initialize arrays
prediction_error = np.empty([sims, 1000])
RMSE = []

alphas = [1e20, 1000, 500, 100, 50, 10, 5, 1, 0.5, 0.1, 0.05, 0.01, 0.005, 0.001]

kf = KFold(n, n_folds=10)
al = 0.0
temp = 10000000

for alpha in alphas:
	prediction_MSE = []
	for train_index, test_index in kf:

		# Take the subset of X_train and X_test
		X_train1 = X_train[train_index]
		X_test1 = X_train[test_index]

		# Calculate the Y_train and Y_test based on the real beta values
		Y_train1 = beta0 + np.dot(X_train1, beta) + np.random.normal(0, sd_e, 1)
		Y_test1 = beta0 + np.dot(X_test1, beta) + np.random.normal(0, sd_e, 1)

		# Create Ridge regression object
		#regr = linear_model.LinearRegression(fit_intercept=False, normalize=False)
		#regr = linear_model.Ridge(alpha, fit_intercept=False, normalize=True)
		regr = linear_model.Lasso(alpha, fit_intercept=False, normalize=True, max_iter=2000)

		# Train the model using the training sets
		regr.fit(X_train1, Y_train1)

		# Calclulate prediction error
		prediction_error = Y_test1 - regr.predict(X_test1)

		# Calculate RMSE for each fold
		prediction_MSE.append(np.sqrt(np.sum(prediction_error**2)/len(prediction_error))) 

	#Store RMSE for each alpha and keep track of the best alpha
	pmse = np.mean(prediction_MSE)
	if (pmse < temp):
		temp = pmse
		al = alpha
	RMSE.append(pmse)

print ("Best alpha:", al)
#print np.log(1/(np.asarray(alphas) + 1e-08))
plt.plot(np.log(1/(np.asarray(alphas) + 1e-08)), RMSE, 'o')
plt.show()

