import numpy as np;

import matplotlib.pyplot as plt

def load(filename):
    f = open(filename, 'rb')
    return np.loadtxt(f, delimiter=",")

def plotData(x, y):
    plt.plot(x, y, 'rx', markersize=10)
    plt.ylabel('Profit in $10,000s')
    plt.xlabel('Population of City in 10,000s')
    # plt.show()

data = load('ex1data1.txt')
X = data[:,0:1]
y = data[:,1:2]

plotData(X, y)

m = len(X);
X = np.concatenate((np.ones((m, 1)), X), 1)
theta = np.zeros((2, 1))
iterations = 1500;
alpha = 0.01;

def computeCost(X, y, theta):
    m = len(y)
    return np.sum((X @ theta - y) ** 2, axis=0) / (2 * m)

# test case
# print('computeCost() test, 11.9450, 7.0175')
# print(computeCost(np.array([[1, 2], [1, 3], [1,4], [1,5]]), np.array([[7], [6], [5], [4]]), np.array([[0.1], [0.2]])))
# print(computeCost(np.array([[1, 2, 3], [1, 3, 4], [1, 4, 5], [1, 5, 6]]), np.array([[7], [6], [5], [4]]), np.array([[0.1], [0.2], [0.3]])))
# print('test over')

print(computeCost(X, y, theta))
print(computeCost(X, y, np.array([[-1],[2]])))

def gradientDescent(X, y, theta, alpha, num_iters):
    m = len(y)
    J_history = np.zeros((num_iters, 1))
    for iter in range(num_iters):
        h = (X @ theta - y)
        theta = theta - alpha * (np.sum(h * X, axis=0, keepdims=True) / m).T
        J_history[iter] = computeCost(X, y, theta)
    return (theta, J_history)

# print('gradientDescent() test')
# theta_test, J_hist = gradientDescent(np.array([[1, 5], [1, 2], [1, 4], [1, 5]]), np.array([[1], [6], [4], [2]]), np.array([[0], [0]]), 0.01, 1000)
# print(theta_test, J_hist[1], J_hist[-1])

# theta_test, J_hist = gradientDescent(np.array([[1, 5], [1, 2]]), np.array([[1], [6]]), np.array([[.5], [.5]]), 0.1, 10)
# print(theta_test, J_hist)
# print('test over')

(theta, _) = gradientDescent(X, y, theta, alpha, iterations)

print('Theta computed from gradient descent:\n', theta[0], theta[1])

plt.plot(X[:,1:2], X @ theta, '-', markersize=10)
plt.legend(['Training data', 'Linear Regression'])
plt.show()

predict1 = np.array([[1, 3.5]]) @ theta
print('For population = 35,000, we predict a profit of', predict1 * 10000)

predict2 = np.array([[1, 7]]) @ theta;
print('For population = 70,000, we predict a profit of', predict2 * 10000)

theta0_vals = np.linspace(-10, 10, 100)
theta1_vals = np.linspace(-1, 4, 100)

J_vals = np.zeros((len(theta0_vals), len(theta1_vals)))

for i in range(len(theta0_vals)):
    for j in range(len(theta1_vals)):
        t = [[theta0_vals[i]], [theta1_vals[j]]]
        J_vals[i, j] = computeCost(X, y ,t)

J_vals = J_vals.T


ax = plt.axes(projection='3d')
theta0_vals, theta1_vals = np.meshgrid(theta0_vals, theta1_vals)
ax.plot_surface(theta0_vals, theta1_vals, J_vals, cmap=plt.get_cmap('magma'))
ax.invert_xaxis()
plt.xlabel(r'$\theta_0$')
plt.ylabel(r'$\theta_1$')

plt.show()

plt.figure()
plt.contour(theta0_vals, theta1_vals, J_vals, np.logspace(-2, 3, 20))
plt.xlabel(r'$\theta_0$')
plt.ylabel(r'$\theta_1$')
plt.plot(theta[0], theta[1], 'rx', markersize=10, linewidth=2);
plt.show()
