import numpy as np
import matplotlib.pyplot as plt

# Load the data
x_train = np.array([1.0, 2.0])
y_train = np.array([300.0, 500.0])

#Function to calculate the cost
def compute_cost(x, y, w, b):
    m = x.shape[0]
    cost = 0.0
    for i in range(m):
        f_wb = w * x[i] + b
        cost = cost + (f_wb - y[i])**2
    total_cost = 1 / (2 * m) * cost

    return total_cost

#Function to calculate the gradient
def compute_gradient(x, y, w, b):
    m = x.shape[0]
    dj_dw = 0.0
    dj_db = 0.0
    for i in range(m):
        f_wb = w * x[i] + b
        dj_dw = dj_dw + (f_wb - y[i]) * x[i]
        dj_db = dj_db + (f_wb - y[i])
    dj_dw = dj_dw / m
    dj_db = dj_db / m

    return dj_dw, dj_db

#Function to update the weights
def update_weights(w, b, alpha, dj_dw, dj_db):
    w = w - alpha * dj_dw
    b = b - alpha * dj_db

    return w, b

#Function to train the model
def train(x, y, w, b, alpha, epochs):
    costs = []
    params = []
    for i in range(epochs):
        cost = compute_cost(x, y, w, b)
        costs.append(cost)
        dj_dw, dj_db = compute_gradient(x, y, w, b)
        w, b = update_weights(w, b, alpha, dj_dw, dj_db)
        params.append([w, b])
    return w, b, costs, params

#Initialize the weights
w = 0.0
b = 0.0
alpha = 0.01
epochs = 1000

#Train the model
w, b, costs, params = train(x_train, y_train, w, b, alpha, epochs)

#Plot the cost
plt.plot(costs)
plt.xlabel('Iterations')
plt.ylabel('Cost')
plt.title('Cost vs Iterations')
plt.show()

#Plot the line
plt.scatter(x_train, y_train)
x = np.linspace(0, 3, 100)
y = w * x + b
plt.plot(x, y, '-r')
plt.xlabel('x')
plt.ylabel('y')
plt.title('Line of best fit')
plt.show()