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

from linearRegCostFunction import linearRegCostFunction
from trainLinearReg import trainLinearReg
from learningCurve import learningCurve
from polyFeatures import polyFeatures
from featureNormalize import featureNormalize
from plotFit import plotFit
from validationCurve import validationCurve

#  Machine Learning Online Class
#  Exercise 5 | Regularized Linear Regression and Bias-Variance
#
#  Instructions
#  ------------
#
#  This file contains code that helps you get started on the
#  exercise. You will need to complete the following functions:
#
#     linearRegCostFunction.m
#     learningCurve.m
#     validationCurve.m
#
#  For this exercise, you will not need to change any code in this file,
#  or any other files other than those mentioned above.

#  =========== Part 1: Loading and Visualizing Data =============
#  We start the exercise by first loading and visualizing the dataset.
#  The following code will load the dataset into your environment and plot
#  the data.

# Load Training Data
print('Loading and Visualizing Data ...')

# Load from ex5data1:
# You will have X, y, Xval, yval, Xtest, ytest in your environment
data = scipy.io.loadmat('ex5data1.mat')

# m = Number of examples
X = data['X'][:, 0]
y = data['y'][:, 0]
Xval = data['Xval'][:, 0]
yval = data['yval'][:, 0]
Xtest = data['Xtest'][:, 0]
ytest = data['ytest'][:, 0]
m = X.size

# Plot training data
plt.scatter(X, y, marker='x', s=60, edgecolor='r', color='r', lw=1.5)
plt.ylabel('Water flowing out of the dam (y)')  # Set the y-axis label
plt.xlabel('Change in water level (x)')  # Set the x-axis label
show()

input('Program paused. Press Enter to continue...')

#  =========== Part 2: Regularized Linear Regression Cost =============
#  You should now implement the cost function for regularized linear
#  regression.

theta = np.array([1, 1])
J, _ = linearRegCostFunction(np.column_stack((np.ones(m), X)), y, theta, Lambda=1)

print('Cost at theta = [1  1]: %f \n(this value should be about 303.993192)\n' % J)

input('Program paused. Press Enter to continue...')

#  =========== Part 3: Regularized Linear Regression Gradient =============
#  You should now implement the gradient for regularized linear
#  regression.

theta = np.array([1, 1])
J, grad = linearRegCostFunction(np.column_stack((np.ones(m), X)), y, theta, Lambda=1)

print('Gradient at theta = [1  1]:  [%f %f] \n'
      '(this value should be about [-15.303016 598.250744])\n' % (grad[0], grad[1]))

input('Program paused. Press Enter to continue...')

#  =========== Part 4: Train Linear Regression =============
#  Once you have implemented the cost and gradient correctly, the
#  trainLinearReg function will use your cost function to train
#  regularized linear regression.
#
#  Write Up Note: The data is non-linear, so this will not give a great
#                 fit.

#  Train linear regression with Lambda = 0
Lambda = 0
theta = trainLinearReg(np.column_stack((np.ones(m), X)), y, Lambda=Lambda)

#  Plot fit over the data
plt.scatter(X, y, marker='x', s=60, edgecolor='r', color='r', lw=1.5)
plt.ylabel('Water flowing out of the dam (y)')  # Set the y-axis label
plt.xlabel('Change in water level (x)')  # Set the x-axis label
plt.plot(X, np.column_stack((np.ones(m), X)).dot(theta), '--', lw=2.0)
show()

input('Program paused. Press Enter to continue...')

#  =========== Part 5: Learning Curve for Linear Regression =============
#  Next, you should implement the learningCurve function.
#
#  Write Up Note: Since the model is underfitting the data, we expect to
#                 see a graph with "high bias" -- slide 8 in ML-advice.pdf
#

Lambda = 0
error_train, error_val = learningCurve(np.column_stack((np.ones(m), X)), y,
                                       np.column_stack((np.ones(Xval.shape[0]), Xval)),
                                       yval,
                                       Lambda=Lambda)
plt.figure()
plt.plot(range(1,m + 1), error_train, color='b', lw=2, label='Train')
plt.plot(range(1,m + 1), error_val, color='g', lw=2, label='Cross Validation')
plt.title('Learning Curve for Linear Regression')
plt.legend()
plt.xlabel('Number of training examples')
plt.ylabel('Error')

plt.xlim(0, 13)
plt.legend(loc='upper right', shadow=True, fontsize='x-large', numpoints=1)
show()

print('Training Examples\tTrain Error\tCross Validation Error')
plt.ylim(0, max(np.concatenate((error_train, error_val))) + 10)
for i in range(m):
    print('  \t%d\t\t%f\t%f' % (i, error_train[i], error_val[i]))

input('Program paused. Press Enter to continue...')"""

#  =========== Part 6: Feature Mapping for Polynomial Regression =============
#  One solution to this is to use polynomial regression. You should now
#  complete polyFeatures to map each example into its powers

p = 8

# Map X onto Polynomial Features and Normalize
X_poly = polyFeatures(X, p)
X_poly, mu, sigma = featureNormalize(X_poly)  # Normalize
X_poly = np.column_stack((np.ones(m), X_poly))  # Add Ones

# Map X_poly_test and normalize (using mu and sigma)
X_poly_test = polyFeatures(Xtest, p)
X_poly_test = X_poly_test - mu
X_poly_test = X_poly_test / sigma
X_poly_test = np.column_stack((np.ones(X_poly_test.shape[0]), X_poly_test))  # Add Ones

# Map X_poly_val and normalize (using mu and sigma)
X_poly_val = polyFeatures(Xval, p)
X_poly_val = X_poly_val - mu
X_poly_val = X_poly_val / sigma
X_poly_val = np.column_stack((np.ones(X_poly_test.shape[0]), X_poly_val))  # Add Ones
"""
print('Normalized Training Example 1:')
print(X_poly[0, :])

input('Program paused. Press Enter to continue...')

#  =========== Part 7: Learning Curve for Polynomial Regression =============
#  Now, you will get to experiment with polynomial regression with multiple
#  values of Lambda. The code below runs polynomial regression with
#  Lambda = 0. You should try running the code with different values of
#  Lambda to see how the fit and learning curve change.

Lambda = 0.1
theta = trainLinearReg(X_poly, y, Lambda, method='BFGS', maxiter=10)

# Plot training data and fit
plt.figure()
plt.scatter(X, y, marker='x', s=60, edgecolor='r', color='r', lw=1.5)

plotFit(min(X), max(X), mu, sigma, theta, p)

plt.xlabel('Change in water level (x)')  # Set the y-axis label
plt.ylabel('Water flowing out of the dam (y)')  # Set the x-axis label
plt.title('Polynomial Regression Fit (Lambda = %f)' % Lambda)
show()

input('Program paused. Press Enter to continue...')

error_train, error_val = learningCurve(X_poly, y, X_poly_val, yval, Lambda)
plt.figure()
plt.plot(range(1, m + 1), error_train, color='b', lw=2, label='Train')
plt.plot(range(1, m + 1), error_val, color='g', lw=2, label='Cross Validation')
plt.title('Polynomial Regression Learning Curve (Lambda = %f)' % Lambda)
plt.xlabel('Number of training examples')
plt.ylabel('Error')
plt.xlim(0, 13)
plt.ylim(0,  max(np.concatenate((error_train, error_val))) + 10)
plt.legend(loc='upper right', shadow=True, fontsize='x-large', numpoints=1)
show()

print('Polynomial Regression (Lambda = %f)\n\n' % Lambda)
print('# Training Examples\tTrain Error\tCross Validation Error')
for i in range(m):
    print('  \t%d\t\t%f\t%f' % (i, error_train[i], error_val[i]))

input('Program paused. Press Enter to continue...')
#  =========== Part 8: Validation for Selecting Lambda =============
#  You will now implement validationCurve to test various values of
#  Lambda on a validation set. You will then use this to select the
#  "best" Lambda value.

Lambda_vec, error_train, error_val = validationCurve(X_poly, y, X_poly_val, yval)
plt.figure()
plt.plot(Lambda_vec, error_train, color='b', lw=0.5, label='Train')
plt.plot(Lambda_vec, error_val, color='g', lw=0.5, label='Cross Validation')
plt.legend(loc='upper right', shadow=True, fontsize='x-large', numpoints=1)
plt.title('Polynomial Regression Validation Curve')
plt.xlabel('Lambda')
plt.ylabel('Error')
show()
# Calculating the error on a TEST set never used before
# using best_lambda found by validationCurve()
best_lambda = 3.2832832832832834
best_theta = trainLinearReg(X_poly, y, best_lambda)
X_poly_test = polyFeatures(Xtest, p)
X_poly_test = X_poly_test - mu
X_poly_test = X_poly_test / sigma
X_poly_test = np.column_stack((np.ones(X_poly_test.shape[0]), X_poly_test))  # Add Ones
error_training, _ = linearRegCostFunction(X_poly_test, ytest, best_theta, 0)
print(f"Error with best_lambda({best_lambda}) on test set = {error_training}")


print('Lambda\t\tTrain Error\tValidation Error')
for i in range(Lambda_vec.size):
    print(' %f\t%f\t%f' % (Lambda_vec[i], error_train[i], error_val[i]))

input('Program paused. Press Enter to continue...')
