import numpy as np
import matplotlib.pyplot as plt

# Data represents housing prices: [square footage, price]
X = np.array([
    [1100, 199000],
    [1400, 245000],
    [1425, 319000],
    [1550, 240000],
    [1600, 312000],
    [1700, 279000],
    [1700, 310000],
    [1875, 308000],
    [2350, 405000],
    [2450, 324000]
])

# Step 0: Preprocess the data

# Normalize the data (zero mean, unit std)
mx = np.mean(X, axis=0)
sx = np.std(X, axis=0)
Xs = (X - mx) / sx

# Plot normalized data
plt.figure(1)
plt.plot(Xs[:, 0], Xs[:, 1], 'k^-')
plt.title('Normalized Data')

# Split data into input (x) and target (t)
x = Xs[:, 0]
t = Xs[:, 1]

# Step 1: Initialize training parameters and constants
m = np.random.randn()  # slope
b = np.random.randn()  # intercept
MaxIters = 60
alpha = 0.01
tol = 1e-6
Error = []

# Main loop: Steps 2-4
for i in range(MaxIters):
    y = m * x + b           # model output
    ErrVec = t - y          # error vector
    Error.append(np.sum(ErrVec**2))  # sum of squared errors

    # Step 3: Gradient calculation
    temp1 = ErrVec * (-x)
    temp2 = ErrVec * (-1)
    Em = 2 * np.sum(temp1)
    Eb = 2 * np.sum(temp2)

    # Step 4: Gradient descent update
    m = m - alpha * Em
    b = b - alpha * Eb

    # Check for convergence
    grad_norm_sq = Em**2 + Eb**2
    print(f'Gradient is {grad_norm_sq:.6f}')
    if grad_norm_sq <= tol:
        print(f'Solution found in {i + 1} steps')
        break

# Closing: Visualize results

# Plot error over iterations
plt.figure(2)
plt.plot(Error)
plt.title('Error vs Iteration')
plt.xlabel('Iteration')
plt.ylabel('Error')

# Plot regression line on original normalized data
plt.figure(1)
t_vals = np.linspace(np.min(x), np.max(x), 100)
z = m * t_vals + b
plt.plot(t_vals, z, 'r-')
plt.legend(['Data', 'Fitted Line'])
plt.show()